Split transition messages over 5MB into multiple TransitionChunk messages (#41607)

To prevent large Transition messages on slow network connections from making the server appear unresponsive while they download, split Transition messages into chunks.

This currently requires npm version 1.27.5 (unreleased).

GitOrigin-RevId: 376609422e5496f62c07f3913a0f476733c84b43

Author: thomasballinger

src/browser/sync/client.ts

@@ -23,8 +23,8 @@ import {
   MutationRequest,
   QueryId,
   QueryJournal,
-  RequestId,
   ServerMessage,
+  RequestId,
   TS,
   UserIdentityAttributes,
 } from "./protocol.js";
@@ -485,8 +485,6 @@ export class BaseConvexClient {
               void this.webSocketManager.terminate();
               throw error;
             }
-            case "Ping":
-              break; // do nothing
             default: {
               serverMessage satisfies never;
             }

src/browser/sync/client_node.test.ts

@@ -9,7 +9,7 @@ import {
   MutationRequest,
   parseServerMessage,
   RequestId,
-  ServerMessage,
+  WireServerMessage,
 } from "./protocol.js";
 import {
   encodeServerMessage,
@@ -62,7 +62,7 @@ test("BaseConvexClient closes cleanly", () => {
 });
 
 test("Tests can encode longs in server messages", () => {
-  const orig: ServerMessage = {
+  const orig: WireServerMessage = {
     type: "Transition",
     startVersion: { querySet: 0, identity: 0, ts: Long.fromNumber(0) },
     endVersion: { querySet: 1, identity: 0, ts: Long.fromNumber(1) },
@@ -103,7 +103,7 @@ test("Actions can be called immediately", async () => {
   });
 });
 
-function actionSuccess(requestId: RequestId): ServerMessage {
+function actionSuccess(requestId: RequestId): WireServerMessage {
   return {
     type: "ActionResponse",
     requestId: requestId,
@@ -779,3 +779,74 @@ test("Query unsubscription triggers empty transition for listeners", async () =>
     await client.close();
   }, true);
 });
+
+test("TransitionChunk messages are assembled into a Transition", async () => {
+  await withInMemoryWebSocket(async ({ address, receive, send }) => {
+    const client = new BaseConvexClient(address, () => null, {
+      webSocketConstructor: nodeWebSocket,
+      unsavedChangesWarning: false,
+    });
+    expect((await receive()).type).toEqual("Connect");
+    expect((await receive()).type).toEqual("ModifyQuerySet");
+
+    const fullTransition: WireServerMessage = {
+      type: "Transition",
+      startVersion: {
+        querySet: 0,
+        ts: Long.fromNumber(0),
+        identity: 0,
+      },
+      endVersion: {
+        querySet: 1,
+        ts: Long.fromNumber(1000),
+        identity: 0,
+      },
+      modifications: [
+        {
+          type: "QueryUpdated",
+          queryId: 0,
+          value: { result: "chunk test data" },
+          logLines: [],
+          journal: null,
+        },
+      ],
+    };
+
+    const fullJson = encodeServerMessage(fullTransition);
+    const midpoint = Math.floor(fullJson.length / 2);
+    const chunk1 = fullJson.substring(0, midpoint);
+    const chunk2 = fullJson.substring(midpoint);
+
+    send({
+      type: "TransitionChunk",
+      chunk: chunk1,
+      partNumber: 0,
+      totalParts: 2,
+      messageLength: fullJson.length,
+    });
+
+    expect(client.getMaxObservedTimestamp()).toBeUndefined();
+
+    send({
+      type: "TransitionChunk",
+      chunk: chunk2,
+      partNumber: 1,
+      totalParts: 2,
+      messageLength: fullJson.length,
+    });
+
+    // synchronously, it's still undefined
+    expect(client.getMaxObservedTimestamp()).toEqual(undefined);
+
+    for (let i = 0; i < 10; i++) {
+      if (client.getMaxObservedTimestamp()) {
+        break;
+      }
+      await new Promise((resolve) => setTimeout(resolve, 100));
+    }
+
+    expect(client.getMaxObservedTimestamp()).toEqual(Long.fromNumber(1000));
+
+    await client.close();
+  });
+});

src/browser/sync/client_node_test_helpers.ts

@@ -7,15 +7,15 @@ import WebSocket, { WebSocketServer } from "ws";
 // Let's pretend this ws WebSocket is a browser WebSocket (it's very close)
 export const nodeWebSocket = WebSocket as unknown as typeof window.WebSocket;
 
-import { ClientMessage, ServerMessage } from "./protocol.js";
+import { ClientMessage, WireServerMessage } from "./protocol.js";
 import { QueryToken } from "./udf_path_utils.js";
 import { BaseConvexClient } from "./client.js";
 
 export type InMemoryWebSocketTest = (args: {
   address: string;
   socket: () => WebSocket;
   receive: () => Promise<ClientMessage>;
-  send: (message: ServerMessage) => void;
+  send: (message: WireServerMessage) => void;
   close: () => void;
 }) => Promise<void>;
 
@@ -66,7 +66,7 @@ export async function withInMemoryWebSocket(
     const structured = JSON.parse(text);
     return structured;
   }
-  function send(message: ServerMessage) {
+  function send(message: WireServerMessage) {
     // eslint-disable-next-line no-console
     if (debug) console.debug(`      <--${message.type}-- server`);
     socket!.send(encodeServerMessage(message));
@@ -97,7 +97,7 @@ export async function withInMemoryWebSocket(
   }
 }
 
-export function encodeServerMessage(message: ServerMessage): string {
+export function encodeServerMessage(message: WireServerMessage): string {
   function replacer(_key: string, value: any) {
     if (Long.isLong(value)) {
       return encodeLong(value);

src/browser/sync/protocol.ts

@@ -19,11 +19,12 @@ export function longToU64(raw: U64): EncodedU64 {
 
 export function parseServerMessage(
   encoded: EncodedServerMessage,
-): ServerMessage {
+): WireServerMessage {
   switch (encoded.type) {
     case "FatalError":
     case "AuthError":
     case "ActionResponse":
+    case "TransitionChunk":
     case "Ping": {
       return { ...encoded };
     }
@@ -262,6 +263,14 @@ export type Transition = {
   serverTs?: number;
 };
 
+export type TransitionChunk = {
+  type: "TransitionChunk";
+  chunk: string;
+  partNumber: number;
+  totalParts: number;
+  messageLength: number;
+};
+
 type MutationSuccess = {
   type: "MutationResponse";
   requestId: RequestId;
@@ -312,11 +321,20 @@ type Ping = {
   type: "Ping";
 };
 
+// Server Messages without the messages only visible to WebSocketManager
 export type ServerMessage =
   | Transition
   | MutationResponse
   | ActionResponse
   | FatalError
+  | AuthError;
+
+export type WireServerMessage =
+  | Transition
+  | TransitionChunk
+  | MutationResponse
+  | ActionResponse
+  | FatalError
   | AuthError
   | Ping;
 
@@ -329,6 +347,7 @@ type EncodedMutationResponse = MutationFailed | EncodedMutationSuccess;
 
 type EncodedServerMessage =
   | EncodedTransition
+  | TransitionChunk
   | EncodedMutationResponse
   | ActionResponse
   | FatalError

src/browser/sync/web_socket_manager.ts

@@ -5,6 +5,7 @@ import {
   parseServerMessage,
   ServerMessage,
   Transition,
+  TransitionChunk,
 } from "./protocol.js";
 
 const CLOSE_NORMAL = 1000;
@@ -162,6 +163,13 @@ export class WebSocketManager {
     | (string & {}) // a full serverErrorReason (not just the prefix) or a new one
     | null;
 
+  // State for assembling the split-up Transition currently being received.
+  private transitionChunkBuffer: {
+    chunks: string[];
+    totalParts: number;
+    messageLength: number;
+  } | null = null;
+
   /** Upon HTTPS/WSS failure, the first jittered backoff duration, in ms. */
   private readonly defaultInitialBackoff: number;
 
@@ -239,6 +247,60 @@ export class WebSocketManager {
     this.markConnectionStateDirty();
   }
 
+  private assembleTransition(chunk: TransitionChunk): Transition | null {
+    if (
+      chunk.partNumber < 0 ||
+      chunk.partNumber >= chunk.totalParts ||
+      chunk.totalParts === 0 ||
+      (this.transitionChunkBuffer &&
+        (this.transitionChunkBuffer.totalParts !== chunk.totalParts ||
+          this.transitionChunkBuffer.messageLength !== chunk.messageLength))
+    ) {
+      // Throwing an error doesn't crash the client, so clear the buffer.
+      this.transitionChunkBuffer = null;
+      throw new Error("Invalid TransitionChunk");
+    }
+
+    if (this.transitionChunkBuffer === null) {
+      this.transitionChunkBuffer = {
+        chunks: [],
+        totalParts: chunk.totalParts,
+        messageLength: chunk.messageLength,
+      };
+    }
+
+    if (chunk.partNumber !== this.transitionChunkBuffer.chunks.length) {
+      // Throwing an error doesn't crash the client, so clear the buffer.
+      const expectedLength = this.transitionChunkBuffer.chunks.length;
+      this.transitionChunkBuffer = null;
+      throw new Error(
+        `TransitionChunk received out of order: expected part ${expectedLength}, got ${chunk.partNumber}`,
+      );
+    }
+
+    this.transitionChunkBuffer.chunks.push(chunk.chunk);
+
+    if (this.transitionChunkBuffer.chunks.length === chunk.totalParts) {
+      const fullJson = this.transitionChunkBuffer.chunks.join("");
+      this.transitionChunkBuffer = null;
+
+      if (fullJson.length !== chunk.messageLength) {
+        throw new Error(
+          `Assembled Transition length mismatch: expected ${chunk.messageLength}, got ${fullJson.length}`,
+        );
+      }
+      const transition = parseServerMessage(JSON.parse(fullJson));
+      if (transition.type !== "Transition") {
+        throw new Error(
+          `Expected Transition, got ${transition.type} after assembling chunks`,
+        );
+      }
+      return transition;
+    }
+
+    return null;
+  }
+
   private connect() {
     if (this.socket.state === "terminated") {
       return;
@@ -304,6 +366,7 @@ export class WebSocketManager {
     };
     // NB: The WebSocket API calls `onclose` even if connection fails, so we can route all error paths through `onclose`.
     ws.onerror = (error) => {
+      this.transitionChunkBuffer = null;
       const message = (error as ErrorEvent).message;
       if (message) {
         this.logger.log(`WebSocket error message: ${message}`);
@@ -312,8 +375,32 @@ export class WebSocketManager {
     ws.onmessage = (message) => {
       this.resetServerInactivityTimeout();
       const messageLength = message.data.length;
-      const serverMessage = parseServerMessage(JSON.parse(message.data));
+      let serverMessage = parseServerMessage(JSON.parse(message.data));
       this._logVerbose(`received ws message with type ${serverMessage.type}`);
+
+      // Ping's only purpose is to reset the server inactivity timer.
+      if (serverMessage.type === "Ping") {
+        return;
+      }
+
+      // TransitionChunks never reach the main client logic.
+      if (serverMessage.type === "TransitionChunk") {
+        const transition = this.assembleTransition(serverMessage);
+        if (!transition) {
+          return;
+        }
+        serverMessage = transition;
+        this._logVerbose(
+          `assembled full ws message of type ${serverMessage.type}`,
+        );
+      }
+
+      if (this.transitionChunkBuffer !== null) {
+        throw new Error(
+          `Received unexpected ${serverMessage.type} while buffering TransitionChunks`,
+        );
+      }
+
       if (serverMessage.type === "Transition") {
         this.reportLargeTransition({
           messageLength,
@@ -329,6 +416,7 @@ export class WebSocketManager {
     };
     ws.onclose = (event) => {
       this._logVerbose("begin ws.onclose");
+      this.transitionChunkBuffer = null;
       if (this.lastCloseReason === null) {
         // event.reason is often an empty string
         this.lastCloseReason = event.reason || `closed with code ${event.code}`;

src/react/auth_websocket.test.tsx

@@ -14,7 +14,7 @@ import { Long } from "../vendor/long.js";
 import {
   AuthError,
   ClientMessage,
-  ServerMessage,
+  WireServerMessage,
 } from "../browser/sync/protocol.js";
 
 const testReactClient = (address: string, options?: ConvexReactClientOptions) =>
@@ -1152,7 +1152,7 @@ async function assertReconnectWithAuth(
 }
 
 async function simulateAuthError(args: {
-  send: (message: ServerMessage) => void;
+  send: (message: WireServerMessage) => void;
   close: () => void;
   authError: AuthError;
 }) {

src/react/react_node.test.ts

@@ -5,7 +5,7 @@ import { ConvexReactClient } from "./client.js";
 import {
   ClientMessage,
   QuerySetModification,
-  ServerMessage,
+  WireServerMessage,
 } from "../browser/sync/protocol.js";
 import {
   nodeWebSocket,
@@ -60,7 +60,7 @@ const expectQuerySetModification = (
   return message;
 };
 
-function transition(): ServerMessage {
+function transition(): WireServerMessage {
   return {
     type: "Transition",
     startVersion: { querySet: 0, identity: 0, ts: Long.fromNumber(0) },