feat: refactor Peer communication and schedule incoming messages on sync

Author: gdorsi

packages/cojson-transport-ws/src/BatchedOutgoingMessages.ts

@@ -12,34 +12,34 @@ export class BatchedOutgoingMessages {
   push(msg: SyncMessage) {
     const payload = addMessageToBacklog(this.backlog, msg);
 
-    if (this.timeout) {
-      clearTimeout(this.timeout);
-    }
-
     const maxChunkSizeReached =
       payload.length >= MAX_OUTGOING_MESSAGES_CHUNK_BYTES;
     const backlogExists = this.backlog.length > 0;
 
     if (maxChunkSizeReached && backlogExists) {
       this.sendMessagesInBulk();
       this.backlog = addMessageToBacklog("", msg);
-      this.timeout = setTimeout(() => {
-        this.sendMessagesInBulk();
-      }, 0);
     } else if (maxChunkSizeReached) {
       this.backlog = payload;
       this.sendMessagesInBulk();
     } else {
       this.backlog = payload;
+    }
+
+    // Throttling the sending of messages to once every 10ms
+    if (!this.timeout) {
       this.timeout = setTimeout(() => {
         this.sendMessagesInBulk();
-      }, 0);
+        this.timeout = null;
+      }, 10);
     }
   }
 
   sendMessagesInBulk() {
-    this.send(this.backlog);
-    this.backlog = "";
+    if (this.backlog.length > 0) {
+      this.send(this.backlog);
+      this.backlog = "";
+    }
   }
 
   close() {

packages/cojson-transport-ws/src/createWebSocketPeer.ts

@@ -1,11 +1,4 @@
-import {
-  type DisconnectedError,
-  type Peer,
-  type PingTimeoutError,
-  type SyncMessage,
-  cojsonInternals,
-  logger,
-} from "cojson";
+import { type Peer, type SyncMessage, cojsonInternals, logger } from "cojson";
 import { BatchedOutgoingMessages } from "./BatchedOutgoingMessages.js";
 import { deserializeMessages } from "./serialization.js";
 import type { AnyWebSocket } from "./types.js";
@@ -62,9 +55,13 @@ function waitForWebSocketOpen(websocket: AnyWebSocket) {
   });
 }
 
+const { PriorityBasedMessageQueue, CO_VALUE_PRIORITY, ConnectedPeerChannel } =
+  cojsonInternals;
+
 function createOutgoingMessagesManager(
   websocket: AnyWebSocket,
   batchingByDefault: boolean,
+  role: Peer["role"],
 ) {
   let closed = false;
   const outgoingMessages = new BatchedOutgoingMessages((messages) => {
@@ -74,34 +71,63 @@ function createOutgoingMessagesManager(
   });
 
   let batchingEnabled = batchingByDefault;
+  const queue = new PriorityBasedMessageQueue(CO_VALUE_PRIORITY.HIGH, {
+    peerRole: role,
+  });
+  let processing = false;
 
-  async function sendMessage(msg: SyncMessage) {
-    if (closed) {
-      return Promise.reject(new Error("WebSocket closed"));
+  async function processQueue() {
+    if (processing) {
+      return;
     }
 
-    if (websocket.readyState !== 1) {
-      await waitForWebSocketOpen(websocket);
-    }
+    processing = true;
 
-    while (
-      websocket.bufferedAmount > BUFFER_LIMIT &&
-      websocket.readyState === 1
-    ) {
-      await new Promise<void>((resolve) =>
-        setTimeout(resolve, BUFFER_LIMIT_POLLING_INTERVAL),
-      );
-    }
+    let msg = queue.pull();
 
-    if (websocket.readyState !== 1) {
-      return;
+    while (msg) {
+      if (closed) {
+        return;
+      }
+
+      if (websocket.readyState !== 1) {
+        await waitForWebSocketOpen(websocket);
+      }
+
+      while (
+        websocket.bufferedAmount > BUFFER_LIMIT &&
+        websocket.readyState === 1
+      ) {
+        await new Promise<void>((resolve) =>
+          setTimeout(resolve, BUFFER_LIMIT_POLLING_INTERVAL),
+        );
+      }
+
+      if (websocket.readyState !== 1) {
+        return;
+      }
+
+      if (!batchingEnabled) {
+        websocket.send(JSON.stringify(msg));
+      } else {
+        outgoingMessages.push(msg);
+      }
+
+      msg = queue.pull();
     }
 
-    if (!batchingEnabled) {
-      websocket.send(JSON.stringify(msg));
-    } else {
-      outgoingMessages.push(msg);
+    processing = false;
+  }
+
+  function sendMessage(msg: SyncMessage) {
+    if (closed) {
+      return Promise.reject(new Error("WebSocket closed"));
     }
+
+    queue.push(msg);
+    processQueue().catch((e) => {
+      logger.error("Error while processing sendMessage queue", { err: e });
+    });
   }
 
   return {
@@ -137,17 +163,11 @@ export function createWebSocketPeer({
   onSuccess,
   onClose,
 }: CreateWebSocketPeerOpts): Peer {
-  const incoming = new cojsonInternals.Channel<
-    SyncMessage | DisconnectedError | PingTimeoutError
-  >();
+  const incoming = new ConnectedPeerChannel();
   const emitClosedEvent = createClosedEventEmitter(onClose);
 
   function handleClose() {
-    incoming
-      .push("Disconnected")
-      .catch((e) =>
-        logger.error("Error while pushing disconnect msg", { err: e }),
-      );
+    incoming.push("Disconnected");
     emitClosedEvent();
   }
 
@@ -166,18 +186,19 @@ export function createWebSocketPeer({
     expectPings,
     pingTimeout,
     () => {
-      incoming
-        .push("PingTimeout")
-        .catch((e) =>
-          logger.error("Error while pushing ping timeout", { err: e }),
-        );
+      incoming.push("Disconnected");
+      logger.error("Ping timeout from peer", {
+        peerId: id,
+        peerRole: role,
+      });
       emitClosedEvent();
     },
   );
 
   const outgoingMessages = createOutgoingMessagesManager(
     websocket,
     batchingByDefault,
+    role,
   );
   let isFirstMessage = true;
 
@@ -211,45 +232,50 @@ export function createWebSocketPeer({
 
     for (const msg of messages) {
       if (msg && "action" in msg) {
-        incoming
-          .push(msg)
-          .catch((e) =>
-            logger.error("Error while pushing incoming msg", { err: e }),
-          );
+        incoming.push(msg);
       }
     }
   }
 
   websocket.addEventListener("message", handleIncomingMsg);
 
+  const outgoing = new ConnectedPeerChannel();
+
+  outgoing.onMessage((msg) => {
+    if (msg === "Disconnected") {
+      handleClose();
+      return;
+    }
+
+    outgoingMessages.sendMessage(msg);
+  });
+
+  outgoing.onClose(() => {
+    outgoingMessages.close();
+
+    websocket.removeEventListener("message", handleIncomingMsg);
+    websocket.removeEventListener("close", handleClose);
+    pingTimeoutListener.clear();
+    emitClosedEvent();
+
+    if (websocket.readyState === 0) {
+      websocket.addEventListener(
+        "open",
+        function handleClose() {
+          websocket.close();
+        },
+        { once: true },
+      );
+    } else if (websocket.readyState === 1) {
+      websocket.close();
+    }
+  });
+
   return {
     id,
     incoming,
-    outgoing: {
-      push: outgoingMessages.sendMessage,
-      close() {
-        outgoingMessages.close();
-
-        websocket.removeEventListener("message", handleIncomingMsg);
-        websocket.removeEventListener("close", handleClose);
-        pingTimeoutListener.clear();
-        emitClosedEvent();
-
-        if (websocket.readyState === 0) {
-          websocket.addEventListener(
-            "open",
-            function handleClose() {
-              websocket.close();
-            },
-            { once: true },
-          );
-        } else if (websocket.readyState === 1) {
-          websocket.close();
-        }
-      },
-    },
+    outgoing,
     role,
-    crashOnClose: false,
     deletePeerStateOnClose,
   };
 }

packages/cojson-transport-ws/src/tests/BatchedOutgoingMessages.test.ts

@@ -113,7 +113,7 @@ describe("BatchedOutgoingMessages", () => {
     expect(sendMock).toHaveBeenCalledWith(JSON.stringify(message));
   });
 
-  test("should clear timeout when pushing new messages", () => {
+  test("should throttle the messages send", () => {
     const { sendMock, batchedMessages } = setup();
     const message1: SyncMessage = {
       action: "known",
@@ -129,13 +129,8 @@ describe("BatchedOutgoingMessages", () => {
     };
 
     batchedMessages.push(message1);
-
-    const clearTimeoutSpy = vi.spyOn(global, "clearTimeout");
-
     batchedMessages.push(message2);
 
-    expect(clearTimeoutSpy).toHaveBeenCalled();
-
     vi.runAllTimers();
 
     expect(sendMock).toHaveBeenCalledTimes(1);

packages/cojson-transport-ws/src/tests/createWebSocketPeer.test.ts

@@ -48,18 +48,14 @@ describe("createWebSocketPeer", () => {
     expect(peer).toHaveProperty("incoming");
     expect(peer).toHaveProperty("outgoing");
     expect(peer).toHaveProperty("role", "client");
-    expect(peer).toHaveProperty("crashOnClose", false);
   });
 
   test("should handle disconnection", async () => {
     expect.assertions(1);
 
     const { listeners, peer } = setup();
 
-    const incoming = peer.incoming as Channel<
-      SyncMessage | "Disconnected" | "PingTimeout"
-    >;
-    const pushSpy = vi.spyOn(incoming, "push");
+    const pushSpy = vi.spyOn(peer.incoming, "push");
 
     const closeHandler = listeners.get("close");
 
@@ -72,18 +68,15 @@ describe("createWebSocketPeer", () => {
     vi.useFakeTimers();
     const { listeners, peer } = setup();
 
-    const incoming = peer.incoming as Channel<
-      SyncMessage | "Disconnected" | "PingTimeout"
-    >;
-    const pushSpy = vi.spyOn(incoming, "push");
+    const pushSpy = vi.spyOn(peer.incoming, "push");
 
     const messageHandler = listeners.get("message");
 
     messageHandler?.(new MessageEvent("message", { data: "{}" }));
 
     await vi.advanceTimersByTimeAsync(10_000);
 
-    expect(pushSpy).toHaveBeenCalledWith("PingTimeout");
+    expect(pushSpy).toHaveBeenCalledWith("Disconnected");
 
     vi.useRealTimers();
   });
@@ -97,15 +90,14 @@ describe("createWebSocketPeer", () => {
       header: false,
       sessions: {},
     };
-    const promise = peer.outgoing.push(testMessage);
+
+    peer.outgoing.push(testMessage);
 
     await waitFor(() => {
       expect(mockWebSocket.send).toHaveBeenCalledWith(
         JSON.stringify(testMessage),
       );
     });
-
-    await expect(promise).resolves.toBeUndefined();
   });
 
   test("should stop sending messages when the websocket is closed", async () => {
@@ -153,23 +145,6 @@ describe("createWebSocketPeer", () => {
     expect(mockWebSocket.close).toHaveBeenCalled();
   });
 
-  test("should return a rejection if a message is sent after the peer is closed", async () => {
-    const { peer } = setup();
-
-    peer.outgoing.close();
-
-    const message: SyncMessage = {
-      action: "known",
-      id: "co_ztest",
-      header: false,
-      sessions: {},
-    };
-
-    await expect(peer.outgoing.push(message)).rejects.toThrow(
-      "WebSocket closed",
-    );
-  });
-
   test("should call onSuccess handler after receiving first message", () => {
     const onSuccess = vi.fn();
     const { listeners } = setup({ onSuccess });

packages/cojson/package.json

@@ -39,7 +39,6 @@
     "@scure/base": "1.2.1",
     "jazz-crypto-rs": "0.0.7",
     "neverthrow": "^7.0.1",
-    "queueueue": "^4.1.2",
     "unicode-segmenter": "^0.12.0"
   },
   "scripts": {