How we centralized our long-lived WebSockets to improve operations without compromising on developer experience

Legora runs many asynchronous tasks: large zip-file uploads, integration syncs, tabular reviews that extract 100 thousand data points, and agents that run for minutes and hours. Providing live status updates to the user through tRPC subscriptions is critical to keep them in the loop.

For each subscription, users rely on a WebSocket connection in order for our backend servers to send updates to them. When we started transition from thousands to multiple tens of thousands WebSockets, the more the following problem was felt: coordinated reconnects would create thundering herds of authorization checks (we run 6 on average for each subscription), overloading our backends and databases.

To fix this, we created a centralized, type-safe WebSocket relay.

The problem

How do we provide live updates across many features to our users?

One simple way is to use a pull model and let the client (web app, mobile, etc.) poll the status of ongoing work. This is convenient, but doesn’t scale well: every poll adds load to the system. The trade off to make with a polling setup is data freshness versus load. Frequent polling guarantees fresh data within a small margin but causes more load and superfluous work. Infrequent polling is less load, but the user sees stale data more often.

A push model flips the responsibility from the client to the server. This is more efficient (no superfluous work), but more complex. Now we need to keep track of all connected clients and what updates they are interested in receiving. A server needs to hold a long-lived stateful connection (SSE or WebSocket) with all clients for the lifetime of their sessions. For high-availability, we need to run multi replicas of this service, thus we need keep track which client is connected to which server and route incoming messages accordingly.

Initial architecture

We started with a push model. In order to solve the distributed systems routing problem, we added a central streaming service (à la Redis or NATS). In our case we picked the Redis-compatible Dragonfly system since it was already a part of our stack.

A Redis stream is an in-memory, append-only log identified by an ID. This allowed us to tackle our routing problem: When a client creates a new session, it receives a stream ID, then sends a “subscribe” request to the backend service. The backend runs authorization checks for the user requesting the subscription and, if authorized, subscribes to the corresponding redis stream and proxies messages back to the client over a WebSocket. Ironically, reading stream messages in redis is done via polling; but doing it from the backend is still much less expensive than from the front-end.

This architecture worked well, but came with a big drawback. Our authorization checks can be quite heavy, and that’s usually not a problem except in cases of huge bursts in checks - such as coordinated WebSocket reconnections across the platform. And that’s exactly what happened when we deployed new versions of our backend service.

When we deploy a new version, we bring new replicas of our backend online and shut down old ones. When doing this, the old backends will terminate WebSocket connections and clients re-connect to new backends. Since each subscription in a web socket does multiple authorization checks, a backend deployment during peak hours would cause huge load spikes on our authorization service.

Fix 1: Decoupling WebSockets from backend lifecycles

We need to be able to deploy during peak hours. And assuming we can’t make the authorization checks cheaper, we really only have one knob to tune: the rate at which we tear down and recreate WebSockets.

In the initial architecture, we handled deployments using a rollout strategy with a max surge of 25%. This means that we spin up 25% of the total number of replicas we have, and then shut down 25% of our old replicas, and we do this 4 times at which point all replicas will have been replaced. But 25% of all traffic is a sizable burst, lasting minutes. And to make things worse, when a WebSocket reconnects from a terminating server to a live server, nothing stops it from reconnecting to an old replica that’s about to get shutdown on the next rollout wave, amplifying the number of reconnects over the course of the rollout.

A simple approach to fix the authorization checks burst issue would be to catch the shutdown signal (SIGTERM), and slowly drain the WebSocket connections over a configurable period of time instead of dropping all the WebSockets at once. The problem with this is that since the WebSockets are tied to the backend replicas, it slows the rollout of the entire application and we didn’t want that. We asked ourselves how we could decouple WebSocket connections from our backend server instances.

So we started experimenting with a subscription relay service whose sole responsibility was managing WebSocket subscriptions. We had initially hoped we could simply lift-and-shift all the logic of all subscriptions into this new service, but that would mean moving authorization checks into the new service and doing so would hurt developer experience by fragmenting our business and auth logic. A product team working on a new feature would now have to think about defining their logic in two different services: domain and auth logic in the backend service as well as authorization checks in the relay service. How could we keep both domain logic and authorization checks in the backend service? The answer was to separate the responsibilities between the backend service, which would mint JWTs, and the relay service, which would simply validate the claims and hold the web socket connections.

If you squint a bit, this is similar to OAuth: We have a service making authorization checks that mints a JWT, sends it back to the user, who then uses it to connect to another service. This allows the relay service to be completely stateless and domain logic agnostic. The whole flow looks like this:

By subscribing to the relay service using a JWT that includes stream id and URL of the redis instance to use, the relay could skip auth checks and remain completely stateless.

Now there’s only one piece missing: How do we make sure we have no coordinated reconnections during deployments? We’re already helped a lot with the much lower deployment frequency of this service compared to the backend one. Aside from a few performance improvements, we’ve barely had to touch it since the initial deploy. If we’re unlucky, though, we could still run into coordinated node draining events and similar. To solve that, we had to solve the two sides of the reconnection issue: first, all WebSockets being dropped at the same time and second, reconnect amplification because of multiple rollout waves.

To solve the first problem we simply did what we mentioned earlier: we intercepted the SIGTERM signal and instead of dropping connections immediately, drain them over a period of time. The drain is simple: For each active WebSocket, we sample a number between 0 and 900 (15 minutes in seconds) from a uniform distribution, sleep for that amount of time, and then issue a reconnect on the WebSocket.

To solve the second problem we leaned into the Kubernetes machinery. We instrumented our deployments to run a single rollout wave instead of the default 4. In practice, this means that when there’s a new rollout - let’s say we have 20 replicas of the relay service - we would terminate all 20 of them at the same time and bring up 20 new ones immediately. Note that terminating still honors the 15-minute drain. So doing it this way means any time a WebSocket is disconnected during the drain period it's guaranteed to connect to a new replica that won't get terminated.

Typing

The system worked well, but we hit an interesting developer experience snag during implementation. We use typescript and tRPC for almost all our services. In the initial architecture, clients would connect directly to the backend replicas for subscriptions. The subscription is next to the domain logic, so the TS type returned by the subscription is known. In the current architecture, things are not as simple. The clients no longer go through a tRPC subscription to receive the updates, they first fetch a JWT and then connect to the relay. We need to bridge those two worlds: on the one hand, the relay is explicitly designed to be unaware of domain logic, and thus type unaware. On the other hand, the backend only returns a JWT, it no longer proxies the messages from redis. How can we retain complete type safety?

Naively, we could decide to type the subscription from the client itself. The front-end knows which subscription it’s subscribing to, so they could look at the corresponding type and manually type the return of the subscription from the relay. This is error prone and not great DX. Having been very interested in functional programming earlier in my career, the issue reminded me of the concept of a “phantom type” that I had previously seen in Haskell. In Typescript we can emulate such types using what’s called “branded types”. Let’s look at examples.

Without type safety we could type functions generating a relay subscription:

type RelaySubscription = {
  jwt:  string;
};

function signRelaySubscriptionJWT(secret: string, payload: string): RelaySubscription {
  ...
}

type RelaySubscription = {
  jwt:  string;
};

function signRelaySubscriptionJWT(secret: string, payload: string): RelaySubscription {
  ...
}

type RelaySubscription = {
  jwt:  string;
};

function signRelaySubscriptionJWT(secret: string, payload: string): RelaySubscription {
  ...
}

The intuition is that we want to “hide” some type information in the RelaySubscription type so that the front-end can unpack it and cast the return type of the relay subscription to it. This is actually quite simple:

export type RelaySubscription<TReturnType> = {
  jwt:  string;
  readonly  __streamReturnType:  TReturnType;
};

function signRelaySubscriptionJWT(secret: string, payload: string): RelaySubscription<AgentStreamReturnType> {
  ...
}

export type RelaySubscription<TReturnType> = {
  jwt:  string;
  readonly  __streamReturnType:  TReturnType;
};

function signRelaySubscriptionJWT(secret: string, payload: string): RelaySubscription<AgentStreamReturnType> {
  ...
}

export type RelaySubscription<TReturnType> = {
  jwt:  string;
  readonly  __streamReturnType:  TReturnType;
};

function signRelaySubscriptionJWT(secret: string, payload: string): RelaySubscription<AgentStreamReturnType> {
  ...
}

We‘re effectively adding a “fake” field (__streamReturnType) to carry the type information, the field itself is discarded during compilation so there‘s no runtime overhead. Then, in the front-end, we extract the phantom type this way:

type InferPayload<T> = T extends { useQuery: (...args: unknown[]) => { data?: RelaySubscription<infer P> } }
  ? P
  : never;

type InferPayload<T> = T extends { useQuery: (...args: unknown[]) => { data?: RelaySubscription<infer P> } }
  ? P
  : never;

type InferPayload<T> = T extends { useQuery: (...args: unknown[]) => { data?: RelaySubscription<infer P> } }
  ? P
  : never;

It might look scary, but it‘s essentially just a tiny helper function which, given any tRPC query generating a RelaySubscription type, infers the phantom type.

This way, we can create a generic utility function tying everything together:

export type UseSubscriptionRelayInput<TReturnType> = {
  // The tRPC procedure used to fetch the subscription JWT
  query: RelayProcedure<TReturnType>;
  // Called with each stream message, typed via the phantom type
  onSubscriptionData: (data: TReturnType) => void;
};

export function useSubscriptionRelay(
 input:  UseSubscriptionRelayInput<TReturnType>
) {
  // query is the function used to retrieve a JWT
  const { query, onSubscriptionData } = input;

  const { data } = query.useQuery();
  const token = data?.jwt;

  relay.stream.useSubscription({ token },
    onData: (data: unknown) => {
      onSubscriptionData(data as InferPayload<TReturnType>);
    }
  )

export type UseSubscriptionRelayInput<TReturnType> = {
  // The tRPC procedure used to fetch the subscription JWT
  query: RelayProcedure<TReturnType>;
  // Called with each stream message, typed via the phantom type
  onSubscriptionData: (data: TReturnType) => void;
};

export function useSubscriptionRelay(
 input:  UseSubscriptionRelayInput<TReturnType>
) {
  // query is the function used to retrieve a JWT
  const { query, onSubscriptionData } = input;

  const { data } = query.useQuery();
  const token = data?.jwt;

  relay.stream.useSubscription({ token },
    onData: (data: unknown) => {
      onSubscriptionData(data as InferPayload<TReturnType>);
    }
  )

export type UseSubscriptionRelayInput<TReturnType> = {
  // The tRPC procedure used to fetch the subscription JWT
  query: RelayProcedure<TReturnType>;
  // Called with each stream message, typed via the phantom type
  onSubscriptionData: (data: TReturnType) => void;
};

export function useSubscriptionRelay(
 input:  UseSubscriptionRelayInput<TReturnType>
) {
  // query is the function used to retrieve a JWT
  const { query, onSubscriptionData } = input;

  const { data } = query.useQuery();
  const token = data?.jwt;

  relay.stream.useSubscription({ token },
    onData: (data: unknown) => {
      onSubscriptionData(data as InferPayload<TReturnType>);
    }
  )

It‘s a bit abstract, so let’s take a concrete example with the Agent. The stream is really just producing a sequence of strings (the Agent reply chunks), so the return type of the query to get a JWT for this subscription is RelaySubscription<string> . Then, let’s look at the inferred types using our hook to subscribe to the Agent stream:

useSubscriptionRelay({
  // the return type of `getSubscriptionRelayJWT` is RelaySubscription<string>
  query: trpc.agent.getSubscriptionRelayJWT,
  // inside the hook, the typelevel utility function 
  // InferPayload<RelaySubscription<string>> evaluates to string
  // which we use to type cast the input we feed to the callback passed here
  onSubscriptionData: (data) => {
    // at this point the input `data` type is correctly defined as a string
  }
})

useSubscriptionRelay({
  // the return type of `getSubscriptionRelayJWT` is RelaySubscription<string>
  query: trpc.agent.getSubscriptionRelayJWT,
  // inside the hook, the typelevel utility function 
  // InferPayload<RelaySubscription<string>> evaluates to string
  // which we use to type cast the input we feed to the callback passed here
  onSubscriptionData: (data) => {
    // at this point the input `data` type is correctly defined as a string
  }
})

useSubscriptionRelay({
  // the return type of `getSubscriptionRelayJWT` is RelaySubscription<string>
  query: trpc.agent.getSubscriptionRelayJWT,
  // inside the hook, the typelevel utility function 
  // InferPayload<RelaySubscription<string>> evaluates to string
  // which we use to type cast the input we feed to the callback passed here
  onSubscriptionData: (data) => {
    // at this point the input `data` type is correctly defined as a string
  }
})

Impact

After rolling out this system, the peak loads on our authorization service completely vanished and we could confidently deploy many times during peak hours - and handle full node drains confidently. With this architecture we’re confident we can scale 10-50x before hitting new bottlenecks (which is likely to happen within the year here at Legora!).

The relay service is stateless and horizontally scalable so the the bottleneck we would most likely see is the Redis/dragonfly instance backing the streams. We can scale that heavily vertically (the load of streams is quite low), but at some point we would either have to shard it or migrate to something like NATS jetstream. At 100x, we would also start being more careful about resource usage and might re-write the relay service in Rust or Go to pack more WebSocket connections per pod - mainly for efficiency constraints.

If these problems speak to you, you should apply to join the Foundations team where we focus on problems at the intersection of infrastructure and backend. We’re hiring in Copenhagen, Stockholm, London, and New York City!