cojson-transport-ws 0.8.12 → 0.8.16

package/src/index.ts

@@ -1,151 +1,211 @@
 import {
-    DisconnectedError,
-    Peer,
-    PingTimeoutError,
-    SyncMessage,
-    cojsonInternals,
+  DisconnectedError,
+  Peer,
+  PingTimeoutError,
+  SyncMessage,
+  cojsonInternals,
 } from "cojson";
-import { AnyWebSocket } from "./types.js";
 import { BatchedOutgoingMessages } from "./BatchedOutgoingMessages.js";
 import { deserializeMessages } from "./serialization.js";
+import { AnyWebSocket } from "./types.js";
 
 export const BUFFER_LIMIT = 100_000;
 export const BUFFER_LIMIT_POLLING_INTERVAL = 10;
 
 export type CreateWebSocketPeerOpts = {
-    id: string;
-    websocket: AnyWebSocket;
-    role: Peer["role"];
-    expectPings?: boolean;
-    batchingByDefault?: boolean;
+  id: string;
+  websocket: AnyWebSocket;
+  role: Peer["role"];
+  expectPings?: boolean;
+  batchingByDefault?: boolean;
+  onClose?: () => void;
 };
 
-export function createWebSocketPeer({
-    id,
-    websocket,
-    role,
-    expectPings = true,
-    batchingByDefault = true,
-}: CreateWebSocketPeerOpts): Peer {
-    const incoming = new cojsonInternals.Channel<
-        SyncMessage | DisconnectedError | PingTimeoutError
-    >();
-
-    websocket.addEventListener("close", function handleClose() {
-        incoming
-            .push("Disconnected")
-            .catch((e) =>
-                console.error("Error while pushing disconnect msg", e),
-            );
-    });
+function createPingTimeoutListener(enabled: boolean, callback: () => void) {
+  if (!enabled) {
+    return {
+      reset() {},
+      clear() {},
+    };
+  }
+
+  let pingTimeout: ReturnType<typeof setTimeout> | null = null;
+
+  return {
+    reset() {
+      pingTimeout && clearTimeout(pingTimeout);
+      pingTimeout = setTimeout(() => {
+        callback();
+      }, 10_000);
+    },
+    clear() {
+      pingTimeout && clearTimeout(pingTimeout);
+    },
+  };
+}
 
-    let pingTimeout: ReturnType<typeof setTimeout> | null = null;
+function waitForWebSocketOpen(websocket: AnyWebSocket) {
+  return new Promise<void>((resolve) => {
+    if (websocket.readyState === 1) {
+      resolve();
+    } else {
+      websocket.addEventListener("open", resolve, { once: true });
+    }
+  });
+}
 
-    let supportsBatching = batchingByDefault;
+function createOutgoingMessagesManager(
+  websocket: AnyWebSocket,
+  batchingByDefault: boolean,
+) {
+  const outgoingMessages = new BatchedOutgoingMessages((messages) => {
+    if (websocket.readyState === 1) {
+      websocket.send(messages);
+    }
+  });
 
-    websocket.addEventListener("message", function handleIncomingMsg(event) {
-        const result = deserializeMessages(event.data as string);
+  let batchingEnabled = batchingByDefault;
 
-        if (!result.ok) {
-            console.error("Error while deserializing messages", event.data, result.error);
-            return;
-        }
+  async function sendMessage(msg: SyncMessage) {
+    if (websocket.readyState !== 1) {
+      await waitForWebSocketOpen(websocket);
+    }
 
-        const { messages } = result;
+    while (
+      websocket.bufferedAmount > BUFFER_LIMIT &&
+      websocket.readyState === 1
+    ) {
+      await new Promise<void>((resolve) =>
+        setTimeout(resolve, BUFFER_LIMIT_POLLING_INTERVAL),
+      );
+    }
 
-        if (!supportsBatching && messages.length > 1) {
-            // If more than one message is received, the other peer supports batching
-            supportsBatching = true;
-        }
+    if (websocket.readyState !== 1) {
+      return;
+    }
 
-        if (expectPings) {
-            pingTimeout && clearTimeout(pingTimeout);
-            pingTimeout = setTimeout(() => {
-                incoming
-                    .push("PingTimeout")
-                    .catch((e) =>
-                        console.error("Error while pushing ping timeout", e),
-                    );
-            }, 10_000);
-        }
+    if (!batchingEnabled) {
+      websocket.send(JSON.stringify(msg));
+    } else {
+      outgoingMessages.push(msg);
+    }
+  }
+
+  return {
+    sendMessage,
+    setBatchingEnabled(enabled: boolean) {
+      batchingEnabled = enabled;
+    },
+    close() {
+      outgoingMessages.close();
+    },
+  };
+}
 
-        for (const msg of messages) {
-            if (msg && "action" in msg) {
-                incoming
-                    .push(msg)
-                    .catch((e) =>
-                        console.error("Error while pushing incoming msg", e),
-                    );
-            }
-        }
-    });
+function createClosedEventEmitter(callback = () => {}) {
+  let disconnected = false;
 
-    const websocketOpen = new Promise<void>((resolve) => {
-        if (websocket.readyState === 1) {
-            resolve();
-        } else {
-            websocket.addEventListener("open", resolve, { once: true });
-        }
-    });
+  return () => {
+    if (disconnected) return;
+    disconnected = true;
+    callback();
+  };
+}
 
-    const outgoingMessages = new BatchedOutgoingMessages((messages) => {
-        if (websocket.readyState === 1) {
-            websocket.send(
-                messages,
-            );
-        }
-    });
+export function createWebSocketPeer({
+  id,
+  websocket,
+  role,
+  expectPings = true,
+  batchingByDefault = true,
+  onClose,
+}: CreateWebSocketPeerOpts): Peer {
+  const incoming = new cojsonInternals.Channel<
+    SyncMessage | DisconnectedError | PingTimeoutError
+  >();
+  const emitClosedEvent = createClosedEventEmitter(onClose);
+
+  function handleClose() {
+    incoming
+      .push("Disconnected")
+      .catch((e) => console.error("Error while pushing disconnect msg", e));
+    emitClosedEvent();
+  }
+
+  websocket.addEventListener("close", handleClose);
+
+  const pingTimeout = createPingTimeoutListener(expectPings, () => {
+    incoming
+      .push("PingTimeout")
+      .catch((e) => console.error("Error while pushing ping timeout", e));
+    emitClosedEvent();
+  });
+
+  const outgoingMessages = createOutgoingMessagesManager(
+    websocket,
+    batchingByDefault,
+  );
+
+  function handleIncomingMsg(event: { data: unknown }) {
+    const result = deserializeMessages(event.data);
+
+    if (!result.ok) {
+      console.error(
+        "Error while deserializing messages",
+        event.data,
+        result.error,
+      );
+      return;
+    }
 
-    async function pushMessage(msg: SyncMessage) {
-        if (websocket.readyState !== 1) {
-            await websocketOpen;
-        }
+    const { messages } = result;
 
-        while (
-            websocket.bufferedAmount > BUFFER_LIMIT &&
-            websocket.readyState === 1
-        ) {
-            await new Promise<void>((resolve) =>
-                setTimeout(resolve, BUFFER_LIMIT_POLLING_INTERVAL),
-            );
-        }
+    if (messages.length > 1) {
+      // If more than one message is received, the other peer supports batching
+      outgoingMessages.setBatchingEnabled(true);
+    }
 
-        if (websocket.readyState !== 1) {
-            return;
-        }
+    pingTimeout.reset();
 
-        if (!supportsBatching) {
-            websocket.send(JSON.stringify(msg));
-        } else {
-            outgoingMessages.push(msg);
-        }
+    for (const msg of messages) {
+      if (msg && "action" in msg) {
+        incoming
+          .push(msg)
+          .catch((e) => console.error("Error while pushing incoming msg", e));
+      }
     }
+  }
 
-    return {
-        id,
-        incoming,
-        outgoing: {
-            push: pushMessage,
-            close() {
-                console.log("Trying to close", id, websocket.readyState);
-                if (supportsBatching) {
-                    outgoingMessages.close();
-                }
-
-                if (websocket.readyState === 0) {
-                    websocket.addEventListener(
-                        "open",
-                        function handleClose() {
-                            websocket.close();
-                        },
-                        { once: true },
-                    );
-                } else if (websocket.readyState == 1) {
-                    websocket.close();
-                }
+  websocket.addEventListener("message", handleIncomingMsg);
+
+  return {
+    id,
+    incoming,
+    outgoing: {
+      push: outgoingMessages.sendMessage,
+      close() {
+        console.log("Trying to close", id, websocket.readyState);
+        outgoingMessages.close();
+
+        websocket.removeEventListener("message", handleIncomingMsg);
+        websocket.removeEventListener("close", handleClose);
+        pingTimeout.clear();
+        emitClosedEvent();
+
+        if (websocket.readyState === 0) {
+          websocket.addEventListener(
+            "open",
+            function handleClose() {
+              websocket.close();
             },
-        },
-        role,
-        crashOnClose: false,
-    };
+            { once: true },
+          );
+        } else if (websocket.readyState == 1) {
+          websocket.close();
+        }
+      },
+    },
+    role,
+    crashOnClose: false,
+  };
 }

package/src/serialization.ts

@@ -2,23 +2,32 @@ import { SyncMessage } from "cojson";
 import { PingMsg } from "./types.js";
 
 export function addMessageToBacklog(backlog: string, message: SyncMessage) {
-    if (!backlog) {
-        return JSON.stringify(message);
-    }
-    return `${backlog}\n${JSON.stringify(message)}`;
+  if (!backlog) {
+    return JSON.stringify(message);
+  }
+  return `${backlog}\n${JSON.stringify(message)}`;
 }
 
-export function deserializeMessages(messages: string)  {
-    try {
-        return {
-            ok: true,
-            messages: messages.split("\n").map((msg) => JSON.parse(msg)) as SyncMessage[] | PingMsg[],
-        } as const;
-    } catch (e) {
-        console.error("Error while deserializing messages", e);
-        return {
-            ok: false,
-            error: e,
-        } as const;
-    }
+export function deserializeMessages(messages: unknown) {
+  if (typeof messages !== "string") {
+    return {
+      ok: false,
+      error: new Error("Expected a string"),
+    } as const;
+  }
+
+  try {
+    return {
+      ok: true,
+      messages: messages.split("\n").map((msg) => JSON.parse(msg)) as
+        | SyncMessage[]
+        | PingMsg[],
+    } as const;
+  } catch (e) {
+    console.error("Error while deserializing messages", e);
+    return {
+      ok: false,
+      error: e,
+    } as const;
+  }
 }

package/src/tests/BatchedOutgoingMessages.test.ts

@@ -1,114 +1,146 @@
-import { describe, test, expect, vi, beforeEach, afterEach } from "vitest";
-import { BatchedOutgoingMessages, MAX_OUTGOING_MESSAGES_CHUNK_BYTES } from "../BatchedOutgoingMessages.js";
 import { SyncMessage } from "cojson";
 import { CoValueKnownState } from "cojson/src/sync.js";
+import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
+import {
+  BatchedOutgoingMessages,
+  MAX_OUTGOING_MESSAGES_CHUNK_BYTES,
+} from "../BatchedOutgoingMessages.js";
 
 beforeEach(() => {
-    vi.useFakeTimers();
-})
+  vi.useFakeTimers();
+});
 
 afterEach(() => {
-    vi.useRealTimers();
-})
+  vi.useRealTimers();
+});
 
 describe("BatchedOutgoingMessages", () => {
-    function setup() {
-        const sendMock = vi.fn();
-        const batchedMessages = new BatchedOutgoingMessages(sendMock);
-        return { sendMock, batchedMessages };
-    }
-
-    test("should batch messages and send them after a timeout", () => {
-        const { sendMock, batchedMessages } = setup();
-        const message1: SyncMessage = { action: "known", id: "co_z1", header: false, sessions: {} };
-        const message2: SyncMessage = { action: "known", id: "co_z2", header: false, sessions: {} };
-
-        batchedMessages.push(message1);
-        batchedMessages.push(message2);
-
-        expect(sendMock).not.toHaveBeenCalled();
-
-        vi.runAllTimers();
-
-        expect(sendMock).toHaveBeenCalledTimes(1);
-        expect(sendMock).toHaveBeenCalledWith(
-            `${JSON.stringify(message1)}\n${JSON.stringify(message2)}`
-        );
-    });
-
-    test("should send messages immediately when reaching MAX_OUTGOING_MESSAGES_CHUNK_BYTES", () => {
-        const { sendMock, batchedMessages } = setup();
-        const largeMessage: SyncMessage = {
-            action: "known",
-            id: "co_z_large",
-            header: false,
-            sessions: {
-                // Add a large payload to exceed MAX_OUTGOING_MESSAGES_CHUNK_BYTES
-                payload: "x".repeat(MAX_OUTGOING_MESSAGES_CHUNK_BYTES)
-            } as CoValueKnownState['sessions'],
-           
-        };
-
-        batchedMessages.push(largeMessage);
-
-        expect(sendMock).toHaveBeenCalledTimes(1);
-        expect(sendMock).toHaveBeenCalledWith(JSON.stringify(largeMessage));
-    });
-
-    test("should send accumulated messages before a large message", () => {
-        const { sendMock, batchedMessages } = setup();
-        const smallMessage: SyncMessage = { action: "known", id: "co_z_small", header: false, sessions: {} };
-        const largeMessage: SyncMessage = {
-            action: "known",
-            id: "co_z_large",
-            header: false,
-            sessions: {
-                // Add a large payload to exceed MAX_OUTGOING_MESSAGES_CHUNK_BYTES
-                payload: "x".repeat(MAX_OUTGOING_MESSAGES_CHUNK_BYTES)
-            } as CoValueKnownState['sessions'],
-        };
-
-        batchedMessages.push(smallMessage);
-        batchedMessages.push(largeMessage);
-
-        vi.runAllTimers();
-
-        expect(sendMock).toHaveBeenCalledTimes(2);
-        expect(sendMock).toHaveBeenNthCalledWith(1, JSON.stringify(smallMessage));
-        expect(sendMock).toHaveBeenNthCalledWith(2, JSON.stringify(largeMessage));
-    });
-
-    test("should send remaining messages on close", () => {
-        const { sendMock, batchedMessages } = setup();
-        const message: SyncMessage = { action: "known", id: "co_z_test", header: false, sessions: {} };
-
-        batchedMessages.push(message);
-        expect(sendMock).not.toHaveBeenCalled();
-
-        batchedMessages.close();
-
-        expect(sendMock).toHaveBeenCalledTimes(1);
-        expect(sendMock).toHaveBeenCalledWith(JSON.stringify(message));
-    });
-
-    test("should clear timeout when pushing new messages", () => {
-        const { sendMock, batchedMessages } = setup();
-        const message1: SyncMessage = { action: "known", id: "co_z1", header: false, sessions: {} };
-        const message2: SyncMessage = { action: "known", id: "co_z2", header: false, sessions: {} };
-
-        batchedMessages.push(message1);
-        
-        const clearTimeoutSpy = vi.spyOn(global, 'clearTimeout');
-        
-        batchedMessages.push(message2);
-
-        expect(clearTimeoutSpy).toHaveBeenCalled();
-
-        vi.runAllTimers();
-
-        expect(sendMock).toHaveBeenCalledTimes(1);
-        expect(sendMock).toHaveBeenCalledWith(
-            `${JSON.stringify(message1)}\n${JSON.stringify(message2)}`
-        );
-    });
-});
+  function setup() {
+    const sendMock = vi.fn();
+    const batchedMessages = new BatchedOutgoingMessages(sendMock);
+    return { sendMock, batchedMessages };
+  }
+
+  test("should batch messages and send them after a timeout", () => {
+    const { sendMock, batchedMessages } = setup();
+    const message1: SyncMessage = {
+      action: "known",
+      id: "co_z1",
+      header: false,
+      sessions: {},
+    };
+    const message2: SyncMessage = {
+      action: "known",
+      id: "co_z2",
+      header: false,
+      sessions: {},
+    };
+
+    batchedMessages.push(message1);
+    batchedMessages.push(message2);
+
+    expect(sendMock).not.toHaveBeenCalled();
+
+    vi.runAllTimers();
+
+    expect(sendMock).toHaveBeenCalledTimes(1);
+    expect(sendMock).toHaveBeenCalledWith(
+      `${JSON.stringify(message1)}\n${JSON.stringify(message2)}`,
+    );
+  });
+
+  test("should send messages immediately when reaching MAX_OUTGOING_MESSAGES_CHUNK_BYTES", () => {
+    const { sendMock, batchedMessages } = setup();
+    const largeMessage: SyncMessage = {
+      action: "known",
+      id: "co_z_large",
+      header: false,
+      sessions: {
+        // Add a large payload to exceed MAX_OUTGOING_MESSAGES_CHUNK_BYTES
+        payload: "x".repeat(MAX_OUTGOING_MESSAGES_CHUNK_BYTES),
+      } as CoValueKnownState["sessions"],
+    };
+
+    batchedMessages.push(largeMessage);
+
+    expect(sendMock).toHaveBeenCalledTimes(1);
+    expect(sendMock).toHaveBeenCalledWith(JSON.stringify(largeMessage));
+  });
+
+  test("should send accumulated messages before a large message", () => {
+    const { sendMock, batchedMessages } = setup();
+    const smallMessage: SyncMessage = {
+      action: "known",
+      id: "co_z_small",
+      header: false,
+      sessions: {},
+    };
+    const largeMessage: SyncMessage = {
+      action: "known",
+      id: "co_z_large",
+      header: false,
+      sessions: {
+        // Add a large payload to exceed MAX_OUTGOING_MESSAGES_CHUNK_BYTES
+        payload: "x".repeat(MAX_OUTGOING_MESSAGES_CHUNK_BYTES),
+      } as CoValueKnownState["sessions"],
+    };
+
+    batchedMessages.push(smallMessage);
+    batchedMessages.push(largeMessage);
+
+    vi.runAllTimers();
+
+    expect(sendMock).toHaveBeenCalledTimes(2);
+    expect(sendMock).toHaveBeenNthCalledWith(1, JSON.stringify(smallMessage));
+    expect(sendMock).toHaveBeenNthCalledWith(2, JSON.stringify(largeMessage));
+  });
+
+  test("should send remaining messages on close", () => {
+    const { sendMock, batchedMessages } = setup();
+    const message: SyncMessage = {
+      action: "known",
+      id: "co_z_test",
+      header: false,
+      sessions: {},
+    };
+
+    batchedMessages.push(message);
+    expect(sendMock).not.toHaveBeenCalled();
+
+    batchedMessages.close();
+
+    expect(sendMock).toHaveBeenCalledTimes(1);
+    expect(sendMock).toHaveBeenCalledWith(JSON.stringify(message));
+  });
+
+  test("should clear timeout when pushing new messages", () => {
+    const { sendMock, batchedMessages } = setup();
+    const message1: SyncMessage = {
+      action: "known",
+      id: "co_z1",
+      header: false,
+      sessions: {},
+    };
+    const message2: SyncMessage = {
+      action: "known",
+      id: "co_z2",
+      header: false,
+      sessions: {},
+    };
+
+    batchedMessages.push(message1);
+
+    const clearTimeoutSpy = vi.spyOn(global, "clearTimeout");
+
+    batchedMessages.push(message2);
+
+    expect(clearTimeoutSpy).toHaveBeenCalled();
+
+    vi.runAllTimers();
+
+    expect(sendMock).toHaveBeenCalledTimes(1);
+    expect(sendMock).toHaveBeenCalledWith(
+      `${JSON.stringify(message1)}\n${JSON.stringify(message2)}`,
+    );
+  });
+});