How to manage graphql-ws client at the gateway for different users

[quant-daddy]

In my architecture, each client (browser) connects to the gateway. The gateway then connects with the appropriate upstream graphql subscription service, creating a separate upstream websocket connection for each downstream connection using the code below

const clients: {
  [key: string]: { [userId: string]: Client };
} = {};

const getClient = (url: string, userId: string | null) => {
  let client: Client | null = null;
  if (!userId || !clients[url] || !clients[url][userId]) {
    client = createClient({
      url,
      webSocketImpl: ws,
      connectionParams: {
        userId: userId,
      },
      on: {
        closed: () => {
          console.log("closed url ", url);
        },
        connected: () => {
          console.log("connected url ", url);
        },
      },
    });
    // if a client is create for a user, the cache it
    if (userId) {
      clients[url] = {
        ...clients[url],
        [userId]: client,
      };
    }
  }
  // return the cached version for a user
  if (userId) {
    return clients[url][userId];
  }
  // for non users, return the generated client
  return client as Client;
};

What is recommended way to manage connections at scale: Can we use the same websocket connection for multiple downstream client connection? I think not. I have added more details in this issue. Would love to hear your thoughts @enisdenjo.

[enisdenjo]

I converted this to a Discussion - it is better fitted here. Lets jump right in...

I'd start by saying, be very careful with creating too many WS connections, since each WS connection creates its own TCP connection which can be quite wasteful. However, the maximum connections is not as for SSE (6-8 depending on the browser), you can have (I think) around 200 for Chrome. But, I guess the snippet above is on the server - so the limits dont quite apply (or are configurable at least).

WebSockets are quite resilient so using 1 WS for multiple downstream clients should be just fine. Nevertheless, wouldnt you agree that creating separate upstream connections for each browser is an overkill? You might end up consuming much more resources than necessary when lots of clients connect. 100 clients -> 100 tunnel connections.

For scaling, I'd recommend simple load balancers and services on-demand. A single server can handle quite an amount of WS connections without too many issues - especially if you're using something very efficient, like µWebSockets.js.

[enisdenjo]

I understand the "how?", but I dont quite understand the "why?".

Can you elaborate a bit more about why are you creating persistent connections for each browser? What are you trying to achieve?

Also, how does the connection process look like and how do you actually stream the messages from the broker client back to the browser? Is this diagram correct:

browser <-> (gql-ws server <-> gql-ws client) <-> somewhere else

If the broker server implements a gql-ws server, you can identify browsers easily though distinct IDs passed through the connectionParams. With this setup you can also have one WS connection to the upstream for all browsers connected to the broker server.

Finally, do you have any code examples that you can share?

[quant-daddy]

By persistent connection, I meant a websocket connection that is used for the browser window (or a mobile app or any client). Poor choice of words :)
Yes the diagram is absolutely correct. So the gateway listens for websocket connections like so:

const wsServer = new ws.Server({
  // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
  server: this.server!,
  path: "/graphql",
});
useServer(
  {
    schema: schema,
    onConnect: (ctx) => {
      debug("on connect called");
      if (config.SHOULD_ALLOW_UNAUTHENTICATED_SUBSCRIPTION) {
        return true;
      }
      const userId = getUserIdFromReq(ctx.extra.request);
      if (!userId) {
        debug("on connect rejected unauthenticated used");
        return false;
      }
      return true;
    },
    onSubscribe: async (ctx, msg) => {
      let doc = msg.payload.query;
      const context: ApolloContext = {
        userId: getUserIdFromReq(ctx.extra.request),
      };
      return {
        schema: schema,
        document: parse(doc),
        variableValues: msg.payload.variables,
        contextValue: context,
      };
    },
  },
  wsServer
);

For every subscription request the gateway receives, graphql schema stitching execute that using a remote subscriber to the upstream like so

export function makeRemoteSubscriber(url: string) {
  return async ({
    document,
    variables,
    context,
  }: {
    document: string | ASTNode;
    variables?: {};
    context?: ApolloContext;
  }) => {
    const client = getClient(url); // this just creates a client using the `createClient` method from graphql-ws
    const pending: unknown[] = [];
    let deferred: null | {
      resolve: (value: any) => void;
      reject: (value: any) => void;
    } = null;
    let error: Error | null = null;
    let done = false;

    const query = typeof document === "string" ? document : print(document);
    const dispose = client.subscribe(
      {
        query,
        variables,
      },
      {
        next: (data) => {
          pending.push(data);
          deferred && deferred.resolve(false);
        },
        error: (err: Error) => {
          error = err;
          deferred && deferred.reject(error);
        },
        complete: () => {
          done = true;
          deferred && deferred.resolve(true);
        },
      }
    );

    return {
      [Symbol.asyncIterator]() {
        return this;
      },
      async next() {
        if (done) return { done: true, value: undefined };
        if (error) throw error;
        if (pending.length) return { value: pending.shift() };
        return (await new Promise(
          (resolve, reject) => (deferred = { resolve, reject })
        ))
          ? { done: true, value: undefined }
          : { value: pending.shift() };
      },
      async return() {
        dispose();
        return { done: true, value: undefined };
      },
    };
  };
}

To test this setup, I created two nodejs clients that make a graphql subscritption and they both get the right data for their subscription (because each subscription has an ID associated with it and data is pushed to the subscription with the appropriate ID). So this is great and works as you suggested. We are using a single WS connection to the upstream from the broker server using the getClient method.
My next concern is how to send authentication data (in ApolloContext) for each subscription in the client.subscribe function call in the broker server? Since the client is shared across all the browser, we have to somehow tell the upstream about the user identity for each subscription. I tried passing userId in addition to query and variables in the subscribe call but it gives a typescript error, but this can be ignored. Am I going in the right direction?

[enisdenjo]

Cool, glad it works! From the top of my head, I see 2 potential solutions to your auth problem:

1. Just ignore the TS error and pass along the token/userID to the main server, do the auth there and respond.
2. Another way, and the way I'd recommend, is doing auth on the broker server. You reduce roundtrips for auth, you can decline access directly in the onConnect callback, and you're always safer with bringing the auth logic closer to the user. If you do indeed need the main server for auth, you may as well just perform a REST call from any of graphql-ws server callbacks.

[quant-daddy]

Sounds great! Thanks. The first option works: the upstream needs access to the userId for each subscription so I pass along the connectionParams (which contains the userId) for each subscription.

{
  query,
  variables,
  // @ts-ignore
  connectionParams: context?.connectionParams,
},

My another concern is about making it production ready.

1. How many connections should be created from the broker to the upstream? I guess it would depend on the limit of throughput of each TCP connection right? So may be one can have a pool of clients created on the broker server or dynamically create new client as traffic goes up. Since we are only using one broker to upstream connection, would using uWebsocket have any significant performance benefit?
2. How to keep the connection between the broker and the upstream alive? I am assuming that even if the socket is closed, every time subscribe is called on the broker, it will recreate the connection. So this should not be a problem.

Thank you so much again!

[enisdenjo]

You're in luck, [email protected] just landed, adding the extensions field in the subscribe message payload - you now can type-safely pass along the connectionParams through that field on each subscribe. 😄

Answering your questions:

1. Do prefer having as little as possible. I'd also recommend using a provisioned load balancer for handling the traffic. Using µWebSockets will significantly boost the performance, as advertised, it has an "unfair advantage" over other implementations. Do check the benchmarks and decide yourself. Quoting a statement from the referenced benchmarks page:

   µWebSockets performs with 98% the theoretical maximum for any user space Linux process - if anything would ever be faster, it would only be so by less than 2%.

2. Well, since you're using the graphql-ws client on the broker to connect to the upstream - the connection will be managed on-demand automatically for you as long as you have the lazy option set to true (which is the default). Meaning, it will connect on first subscribe and disconnect on last unsubscribe.

And, you're very welcome. 😄

[pdambrauskas]

Hey!

I have a related question:
When we have a gateway that forwards subscriptions upstream through a single websocket connection as you discussed, how should we handle cases, when gateway clients terminate without sending CompleteMessage message?

As I understand, upstream subscription will not get terminated in that case? How should we address that? Extend the protocol to send custom ClientPing { payload: { clientId: ... } } messages & kill upstream subscriptions related to that client ID, or do you have any other ideas?

[enisdenjo]

Yes, using the ping/pong messages should be the way to go.