perf: precompute the WebSocket frames when broadcasting

darrachequesne · darrachequesne · commit 5f7b47d40f9d · 2023-01-06T08:42:51.000+01:00
Note: - only packets without binary attachments are affected - the permessage-deflate extension must be disabled (which is the default) Previous attempt: - wsPreEncoded option: 5579d40 - fix for binary packets: a33e42b - revert: 88eee59
diff --git a/lib/index.ts b/lib/index.ts
@@ -1,5 +1,6 @@
 import { EventEmitter } from "events";
 import { yeast } from "./contrib/yeast";
+import { WebSocket } from "ws";
 
 /**
  * A public ID, sent by the server at the beginning of the Socket.IO session and which can be used for private messaging
@@ -164,7 +165,7 @@ export class Adapter extends EventEmitter {
     };
 
     packet.nsp = this.nsp.name;
-    const encodedPackets = this.encoder.encode(packet);
+    const encodedPackets = this._encode(packet, packetOpts);
 
     this.apply(opts, (socket) => {
       if (typeof socket.notifyOutgoingListeners === "function") {
@@ -207,7 +208,7 @@ export class Adapter extends EventEmitter {
     // we can use the same id for each packet, since the _ids counter is common (no duplicate)
     packet.id = this.nsp._ids++;
 
-    const encodedPackets = this.encoder.encode(packet);
+    const encodedPackets = this._encode(packet, packetOpts);
 
     let clientCount = 0;
 
@@ -227,6 +228,25 @@ export class Adapter extends EventEmitter {
     clientCountCallback(clientCount);
   }
 
+  private _encode(packet: unknown, packetOpts: Record<string, unknown>) {
+    const encodedPackets = this.encoder.encode(packet);
+
+    if (encodedPackets.length === 1 && typeof encodedPackets[0] === "string") {
+      // "4" being the "message" packet type in the Engine.IO protocol
+      const data = Buffer.from("4" + encodedPackets[0]);
+      // see https://github.com/websockets/ws/issues/617#issuecomment-283002469
+      packetOpts.wsPreEncodedFrame = WebSocket.Sender.frame(data, {
+        readOnly: false,
+        mask: false,
+        rsv1: false,
+        opcode: 1,
+        fin: true,
+      });
+    }
+
+    return encodedPackets;
+  }
+
   /**
    * Gets a list of sockets by sid.
    *
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -12,6 +12,9 @@
   "main": "./dist/index.js",
   "types": "./dist/index.d.ts",
   "description": "default socket.io in-memory adapter",
+  "peerDependencies": {
+    "ws": "*"
+  },
   "devDependencies": {
     "@types/mocha": "^10.0.1",
     "@types/node": "^14.11.2",
diff --git a/test/index.ts b/test/index.ts
@@ -140,6 +140,48 @@ describe("socket.io-adapter", () => {
     expect(ids).to.eql(["s3"]);
   });
 
+  it("should precompute the WebSocket frames when broadcasting", () => {
+    function socket(id) {
+      return [
+        id,
+        {
+          id,
+          client: {
+            writeToEngine(payload, opts) {
+              expect(payload).to.eql(["123"]);
+              expect(opts.preEncoded).to.eql(true);
+              expect(opts.wsPreEncodedFrame.length).to.eql(2);
+              expect(opts.wsPreEncodedFrame[0]).to.eql(Buffer.from([129, 4]));
+              expect(opts.wsPreEncodedFrame[1]).to.eql(
+                Buffer.from([52, 49, 50, 51])
+              );
+            },
+          },
+        },
+      ];
+    }
+    const nsp = {
+      server: {
+        encoder: {
+          encode() {
+            return ["123"];
+          },
+        },
+      },
+      // @ts-ignore
+      sockets: new Map([socket("s1"), socket("s2"), socket("s3")]),
+    };
+    const adapter = new Adapter(nsp);
+    adapter.addAll("s1", new Set());
+    adapter.addAll("s2", new Set());
+    adapter.addAll("s3", new Set());
+
+    adapter.broadcast([], {
+      rooms: new Set(),
+      except: new Set(),
+    });
+  });
+
   describe("utility methods", () => {
     let adapter;
 
