feat: broadcast and expect multiple acks

This feature was added in `[email protected]`:

```js
io.timeout(1000).emit("some-event", (err, responses) => {
  // ...
});
```

Thanks to this change, it will now work with multiple Socket.IO
servers.

Related: socketio/socket.io@8b20457

Author: darrachequesne

lib/index.ts

@@ -21,6 +21,8 @@ enum EventType {
   FETCH_SOCKETS_RESPONSE,
   SERVER_SIDE_EMIT,
   SERVER_SIDE_EMIT_RESPONSE,
+  BROADCAST_CLIENT_COUNT,
+  BROADCAST_ACK,
 }
 
 interface Request {
@@ -32,6 +34,12 @@ interface Request {
   responses: any[];
 }
 
+interface AckRequest {
+  type: EventType.BROADCAST;
+  clientCountCallback: (clientCount: number) => void;
+  ack: (...args: any[]) => void;
+}
+
 /**
  * UID of an emitter using the `@socket.io/postgres-emitter` package
  */
@@ -151,6 +159,7 @@ export class PostgresAdapter extends Adapter {
   private heartbeatTimer: NodeJS.Timeout | undefined;
   private cleanupTimer: NodeJS.Timeout | undefined;
   private requests: Map<string, Request> = new Map();
+  private ackRequests: Map<string, AckRequest> = new Map();
 
   /**
    * Adapter constructor.
@@ -271,12 +280,54 @@ export class PostgresAdapter extends Adapter {
       }
       case EventType.BROADCAST: {
         debug("broadcast with opts %j", document.data.opts);
-        super.broadcast(
-          document.data.packet,
-          PostgresAdapter.deserializeOptions(document.data.opts)
-        );
+
+        const withAck = document.data.requestId !== undefined;
+        if (withAck) {
+          super.broadcastWithAck(
+            document.data.packet,
+            PostgresAdapter.deserializeOptions(document.data.opts),
+            (clientCount) => {
+              debug("waiting for %d client acknowledgements", clientCount);
+              this.publish({
+                type: EventType.BROADCAST_CLIENT_COUNT,
+                data: {
+                  requestId: document.data.requestId,
+                  clientCount,
+                },
+              });
+            },
+            (arg) => {
+              debug("received acknowledgement with value %j", arg);
+              this.publish({
+                type: EventType.BROADCAST_ACK,
+                data: {
+                  requestId: document.data.requestId,
+                  packet: arg,
+                },
+              });
+            }
+          );
+        } else {
+          super.broadcast(
+            document.data.packet,
+            PostgresAdapter.deserializeOptions(document.data.opts)
+          );
+        }
+        break;
+      }
+
+      case EventType.BROADCAST_CLIENT_COUNT: {
+        const request = this.ackRequests.get(document.data.requestId);
+        request?.clientCountCallback(document.data.clientCount);
+        break;
+      }
+
+      case EventType.BROADCAST_ACK: {
+        const request = this.ackRequests.get(document.data.requestId);
+        request?.ack(document.data.packet);
         break;
       }
+
       case EventType.SOCKETS_JOIN: {
         debug("calling addSockets with opts %j", document.data.opts);
         super.addSockets(
@@ -285,6 +336,7 @@ export class PostgresAdapter extends Adapter {
         );
         break;
       }
+
       case EventType.SOCKETS_LEAVE: {
         debug("calling delSockets with opts %j", document.data.opts);
         super.delSockets(
@@ -419,6 +471,7 @@ export class PostgresAdapter extends Adapter {
       if (
         [
           EventType.BROADCAST,
+          EventType.BROADCAST_ACK,
           EventType.SERVER_SIDE_EMIT,
           EventType.SERVER_SIDE_EMIT_RESPONSE,
         ].includes(document.type) &&
@@ -506,6 +559,48 @@ export class PostgresAdapter extends Adapter {
     });
   }
 
+  public broadcastWithAck(
+    packet: any,
+    opts: BroadcastOptions,
+    clientCountCallback: (clientCount: number) => void,
+    ack: (...args: any[]) => void
+  ) {
+    const onlyLocal = opts?.flags?.local;
+    if (!onlyLocal) {
+      const requestId = randomId();
+
+      this.publish({
+        type: EventType.BROADCAST,
+        data: {
+          packet,
+          requestId,
+          opts: PostgresAdapter.serializeOptions(opts),
+        },
+      });
+
+      this.ackRequests.set(requestId, {
+        type: EventType.BROADCAST,
+        clientCountCallback,
+        ack,
+      });
+
+      // we have no way to know at this level whether the server has received an acknowledgement from each client, so we
+      // will simply clean up the ackRequests map after the given delay
+      setTimeout(() => {
+        this.ackRequests.delete(requestId);
+      }, opts.flags!.timeout);
+    }
+
+    // packets with binary contents are modified by the broadcast method, hence the nextTick()
+    process.nextTick(() => {
+      super.broadcastWithAck(packet, opts, clientCountCallback, ack);
+    });
+  }
+
+  public serverCount(): Promise<number> {
+    return Promise.resolve(1 + this.nodesMap.size);
+  }
+
   addSockets(opts: BroadcastOptions, rooms: Room[]) {
     super.addSockets(opts, rooms);
 

package-lock.json

@@ -23,7 +23,7 @@
     "@types/pg": "^8.6.1",
     "debug": "~4.3.1",
     "pg": "^8.0.0",
-    "socket.io-adapter": "~2.3.0"
+    "socket.io-adapter": "~2.4.0"
   },
   "devDependencies": {
     "@types/expect.js": "^0.3.29",
@@ -33,7 +33,7 @@
     "mocha": "^8.4.0",
     "nyc": "^15.1.0",
     "prettier": "^2.1.2",
-    "socket.io": "^4.1.1",
+    "socket.io": "^4.5.0",
     "socket.io-client": "^4.1.1",
     "ts-node": "^9.1.1",
     "typescript": "^4.0.5"

package.json

@@ -2,7 +2,7 @@ import { createServer } from "http";
 import { Server, Socket as ServerSocket } from "socket.io";
 import { io as ioc, Socket as ClientSocket } from "socket.io-client";
 import expect = require("expect.js");
-import { createAdapter, PostgresAdapter } from "..";
+import { createAdapter, PostgresAdapter } from "../lib";
 import type { AddressInfo } from "net";
 import { times, sleep } from "./util";
 import { Pool } from "pg";
@@ -175,6 +175,91 @@ describe("@socket.io/postgres-adapter", () => {
 
       servers[0].local.emit("test");
     });
+
+    it("broadcasts with multiple acknowledgements", (done) => {
+      clientSockets[0].on("test", (cb) => {
+        cb(1);
+      });
+
+      clientSockets[1].on("test", (cb) => {
+        cb(2);
+      });
+
+      clientSockets[2].on("test", (cb) => {
+        cb(3);
+      });
+
+      servers[0].timeout(500).emit("test", (err: Error, responses: any[]) => {
+        expect(err).to.be(null);
+        expect(responses).to.contain(1);
+        expect(responses).to.contain(2);
+        expect(responses).to.contain(3);
+
+        setTimeout(() => {
+          // @ts-ignore
+          expect(servers[0].of("/").adapter.ackRequests.size).to.eql(0);
+
+          done();
+        }, 500);
+      });
+    });
+
+    it("broadcasts with multiple acknowledgements (binary content)", (done) => {
+      clientSockets[0].on("test", (cb) => {
+        cb(Buffer.from([1]));
+      });
+
+      clientSockets[1].on("test", (cb) => {
+        cb(Buffer.from([2]));
+      });
+
+      clientSockets[2].on("test", (cb) => {
+        cb(Buffer.from([3]));
+      });
+
+      servers[0].timeout(500).emit("test", (err: Error, responses: any[]) => {
+        expect(err).to.be(null);
+        responses.forEach((response) => {
+          expect(Buffer.isBuffer(response)).to.be(true);
+        });
+
+        done();
+      });
+    });
+
+    it("broadcasts with multiple acknowledgements (no client)", (done) => {
+      servers[0]
+        .to("abc")
+        .timeout(500)
+        .emit("test", (err: Error, responses: any[]) => {
+          expect(err).to.be(null);
+          expect(responses).to.eql([]);
+
+          done();
+        });
+    });
+
+    it("broadcasts with multiple acknowledgements (timeout)", (done) => {
+      clientSockets[0].on("test", (cb) => {
+        cb(1);
+      });
+
+      clientSockets[1].on("test", (cb) => {
+        cb(2);
+      });
+
+      clientSockets[2].on("test", (cb) => {
+        // do nothing
+      });
+
+      servers[0].timeout(500).emit("test", (err: Error, responses: any[]) => {
+        expect(err).to.be.an(Error);
+        expect(responses).to.contain(1);
+        expect(responses).to.contain(2);
+
+        done();
+      });
+    });
   });
 
   describe("socketsJoin", () => {
@@ -281,7 +366,7 @@ describe("@socket.io/postgres-adapter", () => {
     });
 
     it("returns a single socket instance", async () => {
-      serverSockets[1].data = "test";
+      serverSockets[1].data = "test" as any;
 
       const [remoteSocket] = await servers[0]
         .in(serverSockets[1].id)