feat: use PUB/SUB for fetchSockets() and serverSideEmit() requests

Author: darrachequesne

README.md

@@ -131,6 +131,8 @@ io.listen(3000);
 |---------------------|-----------------------------------------------------------------------------------------------------------------------|----------------|
 | `streamName`        | The name of the Redis stream.                                                                                         | `socket.io`    |
 | `streamCount`       | The number of streams to use to scale horizontally.                                                                   | `1`            |
+| `channelPrefix`     | The prefix of the Redis PUB/SUB channels used to communicate between the nodes.                                       | `socket.io`    |
+| `useShardedPubSub`  | Whether to use sharded PUB/SUB (added in Redis 7.0) to communicate between the nodes.                                 | `false`        |
 | `maxLen`            | The maximum size of the stream. Almost exact trimming (~) is used.                                                    | `10_000`       |
 | `readCount`         | The number of elements to fetch per XREAD call.                                                                       | `100`          |
 | `blockTimeInMs`     | The number of ms before the XREAD call times out.                                                                     | `5_000`        |

lib/adapter.ts

@@ -1,4 +1,4 @@
-import { ClusterAdapterWithHeartbeat, MessageType } from "socket.io-adapter";
+import { ClusterAdapter, MessageType } from "socket.io-adapter";
 import type {
   ClusterAdapterOptions,
   ClusterMessage,
@@ -13,11 +13,16 @@ import {
   hasBinary,
   GETDEL,
   SET,
+  SUBSCRIBE,
   XADD,
   XRANGE,
   XREAD,
   hashCode,
   duplicateClient,
+  SPUBLISH,
+  PUBLISH,
+  PUBSUB,
+  SSUBSCRIBE,
 } from "./util";
 
 const debug = debugModule("socket.io-redis-streams-adapter");
@@ -42,6 +47,17 @@ export interface RedisStreamsAdapterOptions {
    * @default 1
    */
   streamCount?: number;
+  /**
+   * The prefix of the Redis PUB/SUB channels used to communicate between the nodes.
+   * @default "socket.io"
+   */
+  channelPrefix?: string;
+  /**
+   * Whether to use sharded PUB/SUB (added in Redis 7.0) to communicate between the nodes.
+   * @default false
+   * @see https://redis.io/docs/latest/develop/pubsub/#sharded-pubsub
+   */
+  useShardedPubSub?: boolean;
   /**
    * The maximum size of the stream. Almost exact trimming (~) is used.
    * @default 10_000
@@ -168,6 +184,8 @@ export function createAdapter(
     {
       streamName: "socket.io",
       streamCount: 1,
+      channelPrefix: "socket.io",
+      useShardedPubSub: false,
       maxLen: 10_000,
       readCount: 100,
       blockTimeInMs: 5_000,
@@ -197,8 +215,19 @@ export function createAdapter(
     }
   });
 
+  const subClientPromise = duplicateClient(redisClient);
+
+  controller.signal.addEventListener("abort", () => {
+    subClientPromise.then((subClient) => subClient.disconnect());
+  });
+
   return function (nsp) {
-    const adapter = new RedisStreamsAdapter(nsp, redisClient, options);
+    const adapter = new RedisStreamsAdapter(
+      nsp,
+      redisClient,
+      subClientPromise,
+      options
+    );
     namespaceToAdapters.set(nsp.name, adapter);
 
     const defaultClose = adapter.close;
@@ -229,28 +258,71 @@ function computeStreamName(
   }
 }
 
-class RedisStreamsAdapter extends ClusterAdapterWithHeartbeat {
+function isEphemeral(message: ClusterMessage) {
+  const isBroadcastWithAck =
+    message.type === MessageType.BROADCAST &&
+    message.data.requestId !== undefined;
+
+  return (
+    isBroadcastWithAck ||
+    [MessageType.SERVER_SIDE_EMIT, MessageType.FETCH_SOCKETS].includes(
+      message.type
+    )
+  );
+}
+
+class RedisStreamsAdapter extends ClusterAdapter {
   readonly #redisClient: any;
   readonly #opts: Required<RedisStreamsAdapterOptions>;
   readonly #streamName: string;
+  readonly #publicChannel: string;
+  readonly #privateChannel: string;
 
   constructor(
-    nsp,
-    redisClient,
+    nsp: any,
+    redisClient: any,
+    subClientPromise: Promise<any>,
     opts: Required<RedisStreamsAdapterOptions> & ClusterAdapterOptions
   ) {
-    super(nsp, opts);
+    super(nsp);
     this.#redisClient = redisClient;
     this.#opts = opts;
     // each namespace is routed to a specific stream to ensure the ordering of messages
     this.#streamName = computeStreamName(nsp.name, opts);
 
-    this.init();
+    this.#publicChannel = `${opts.channelPrefix}#${nsp.name}#`;
+    this.#privateChannel = `${opts.channelPrefix}#${nsp.name}#${this.uid}#`;
+
+    subClientPromise.then((subClient) => {
+      (this.#opts.useShardedPubSub ? SSUBSCRIBE : SUBSCRIBE)(
+        subClient,
+        [this.#publicChannel, this.#privateChannel],
+        (payload: Buffer) => {
+          try {
+            const message = decode(payload) as ClusterMessage;
+            this.onMessage(message);
+          } catch (e) {
+            return debug("invalid format: %s", e.message);
+          }
+        }
+      );
+    });
   }
 
   override doPublish(message: ClusterMessage) {
     debug("publishing %o", message);
 
+    if (isEphemeral(message)) {
+      // ephemeral messages are sent with Redis PUB/SUB
+      const payload = Buffer.from(encode(message));
+      (this.#opts.useShardedPubSub ? SPUBLISH : PUBLISH)(
+        this.#redisClient,
+        this.#publicChannel,
+        payload
+      );
+      return Promise.resolve("");
+    }
+
     return XADD(
       this.#redisClient,
       this.#streamName,
@@ -263,8 +335,15 @@ class RedisStreamsAdapter extends ClusterAdapterWithHeartbeat {
     requesterUid: ServerId,
     response: ClusterResponse
   ): Promise<void> {
-    // @ts-ignore
-    return this.doPublish(response);
+    const responseChannel = `${this.#opts.channelPrefix}#${
+      this.nsp.name
+    }#${requesterUid}#`;
+    const payload = Buffer.from(encode(response));
+    return (this.#opts.useShardedPubSub ? SPUBLISH : PUBLISH)(
+      this.#redisClient,
+      responseChannel,
+      payload
+    ).then();
   }
 
   private encode(message: ClusterMessage): RawClusterMessage {
@@ -335,6 +414,14 @@ class RedisStreamsAdapter extends ClusterAdapterWithHeartbeat {
     return message;
   }
 
+  override serverCount(): Promise<number> {
+    return PUBSUB(
+      this.#redisClient,
+      this.#opts.useShardedPubSub ? "SHARDNUMSUB" : "NUMSUB",
+      this.#publicChannel
+    );
+  }
+
   override persistSession(session) {
     debug("persisting session %o", session);
     const sessionKey = this.#opts.sessionKeyPrefix + session.pid;

lib/util.ts

@@ -203,3 +203,102 @@ export function hashCode(str: string) {
   }
   return hash;
 }
+
+export function PUBLISH(redisClient: any, channel: string, payload: Buffer) {
+  return redisClient.publish(channel, payload);
+}
+
+export function SPUBLISH(redisClient: any, channel: string, payload: Buffer) {
+  if (isRedisV4Client(redisClient)) {
+    return redisClient.sPublish(channel, payload);
+  } else {
+    return redisClient.spublish(channel, payload);
+  }
+}
+
+const RETURN_BUFFERS = true;
+
+export function SUBSCRIBE(
+  subClient: any,
+  channels: string[],
+  listener: (payload: Uint8Array) => void
+) {
+  if (isRedisV4Client(subClient)) {
+    subClient.subscribe(channels, listener, RETURN_BUFFERS);
+  } else {
+    subClient.subscribe(channels);
+    subClient.on("messageBuffer", (channel: string, payload: Uint8Array) => {
+      if (channels.includes(channel)) {
+        listener(payload);
+      }
+    });
+  }
+}
+
+export function SSUBSCRIBE(
+  subClient: any,
+  channels: string[],
+  listener: (payload: Uint8Array) => void
+) {
+  if (isRedisV4Client(subClient)) {
+    subClient.sSubscribe(channels, listener, RETURN_BUFFERS);
+  } else {
+    subClient.ssubscribe(channels);
+    subClient.on("smessageBuffer", (channel: string, payload: Uint8Array) => {
+      if (channels.includes(channel)) {
+        listener(payload);
+      }
+    });
+  }
+}
+
+function parseNumSubResponse(res: string[]) {
+  return parseInt(res[1], 10);
+}
+
+function sumValues(values) {
+  return values.reduce((acc, val) => acc + val, 0);
+}
+
+export function PUBSUB(
+  redisClient: any,
+  arg: "NUMSUB" | "SHARDNUMSUB",
+  channel: string
+) {
+  if (redisClient.constructor.name === "Cluster" || redisClient.isCluster) {
+    // ioredis cluster
+    return Promise.all(
+      redisClient.nodes().map((node) => {
+        return node
+          .send_command("PUBSUB", [arg, channel])
+          .then(parseNumSubResponse);
+      })
+    ).then(sumValues);
+  } else if (isRedisV4Client(redisClient)) {
+    const isCluster = Array.isArray(redisClient.masters);
+    if (isCluster) {
+      // redis@4 cluster
+      const nodes = redisClient.masters;
+      return Promise.all(
+        nodes.map((node) => {
+          return node.client
+            .sendCommand(["PUBSUB", arg, channel])
+            .then(parseNumSubResponse);
+        })
+      ).then(sumValues);
+    } else {
+      // redis@4 standalone
+      return redisClient
+        .sendCommand(["PUBSUB", arg, channel])
+        .then(parseNumSubResponse);
+    }
+  } else {
+    // ioredis / redis@3 standalone
+    return new Promise((resolve, reject) => {
+      redisClient.send_command("PUBSUB", [arg, channel], (err, numSub) => {
+        if (err) return reject(err);
+        resolve(parseNumSubResponse(numSub));
+      });
+    });
+  }
+}

test/index.ts

@@ -349,7 +349,9 @@ export function testSuite(
         this.timeout(6000);
 
         servers[0].serverSideEmit("hello", (err: Error, response: any) => {
-          expect(err.message).to.be("timeout reached: missing 1 responses");
+          expect(err.message).to.be(
+            "timeout reached: only 1 responses received out of 2"
+          );
           expect(response).to.be.an(Array);
           expect(response).to.contain(2);
           done();
@@ -362,7 +364,7 @@ export function testSuite(
         });
       });
 
-      it("succeeds even if an instance leaves the cluster", (done) => {
+      it.skip("succeeds even if an instance leaves the cluster", (done) => {
         servers[0].on("hello", shouldNotHappen(done));
         servers[1].on("hello", (cb) => cb(2));
         servers[2].on("hello", () => {

test/test-runner.ts

@@ -71,6 +71,23 @@ describe("@socket.io/redis-streams-adapter", () => {
     );
   });
 
+  describe("redis with Redis cluster and sharded PUB/SUB", () => {
+    testSuites(
+      async () => {
+        const redisClient = createCluster({
+          rootNodes: CLUSTER_ROOT_NODES,
+        });
+
+        await redisClient.connect();
+
+        return redisClient;
+      },
+      {
+        useShardedPubSub: true,
+      }
+    );
+  });
+
   describe("ioredis with single Redis node", () => {
     testSuites(() => {
       return new Redis();

test/util.ts

@@ -74,20 +74,11 @@ export function setup(
 
     function isReady() {
       return (
-        servers.length === NODES_COUNT &&
-        clientSockets.length === NODES_COUNT &&
-        servers.every((server) => {
-          const serverCount = server.of("/").adapter.nodesMap.size;
-          return serverCount === NODES_COUNT - 1;
-        })
+        servers.length === NODES_COUNT && clientSockets.length === NODES_COUNT
       );
     }
 
     while (!isReady()) {
-      if (servers.length > 0) {
-        // notify other servers in the cluster
-        servers[0]?.of("/").adapter.init();
-      }
       await sleep(100);
     }