perf(zero-cache): batch v6 websocket change frames

packages/zero-cache/src/services/change-streamer/broadcast.test.ts

@@ -2,6 +2,7 @@ import {afterEach, describe, expect, test, vi} from 'vitest';
 import {BigIntJSON} from '../../../../shared/src/bigint-json.ts';
 import {createSilentLogContext} from '../../../../shared/src/logging-test-utils.ts';
 import {sleep} from '../../../../shared/src/sleep.ts';
+import type {StringifiedStreamPayload} from '../../types/streams.ts';
 import type {Subscription} from '../../types/subscription.ts';
 import {ReplicationMessages} from '../replicator/test-utils.ts';
 import {Broadcast} from './broadcast.ts';
@@ -24,7 +25,9 @@ describe('change-streamer/broadcast', () => {
     vi.restoreAllMocks();
   });
 
-  async function consumeOne(downstream: Subscription<string>) {
+  async function consumeOne(
+    downstream: Subscription<StringifiedStreamPayload>,
+  ) {
     const pipeline = downstream.pipeline;
     if (pipeline === undefined) {
       throw new Error('Expected pipelined subscription');
@@ -43,7 +46,7 @@ describe('change-streamer/broadcast', () => {
 
   async function makeSubscribers(count: number) {
     const subscribers: Subscriber[] = [];
-    const downstreams: Subscription<string>[] = [];
+    const downstreams: Subscription<StringifiedStreamPayload>[] = [];
     for (let i = 0; i < count; i++) {
       const [subscriber, , downstream] = createSubscriber('00', true);
       subscribers.push(subscriber);
@@ -54,7 +57,9 @@ describe('change-streamer/broadcast', () => {
     return {subscribers, downstreams};
   }
 
-  async function consumeChange(downstream: Subscription<string>) {
+  async function consumeChange(
+    downstream: Subscription<StringifiedStreamPayload>,
+  ) {
     await consumeOne(downstream);
     await flushCompletions();
   }

packages/zero-cache/src/services/change-streamer/change-streamer-http.pg.test.ts

@@ -65,8 +65,7 @@ describe('change-streamer/http', () => {
     const {promise, resolve: cleanup} = resolver<Downstream[]>();
     connectionClosed = promise;
     downstream = Subscription.create({
-      cleanup: msgs =>
-        cleanup(msgs.map(m => BigIntJSON.parse(m) as Downstream)),
+      cleanup: msgs => cleanup(msgs.flatMap(parseDownstreamPayload)),
     });
     snapshotStream = Subscription.create();
     subscribeFn = vi.fn();
@@ -145,6 +144,10 @@ describe('change-streamer/http', () => {
     return drained;
   }
 
+  function parseDownstreamPayload(msg: string): Downstream[] {
+    return [BigIntJSON.parse(msg) as Downstream];
+  }
+
   test('health checks and keepalives', async () => {
     const [parent] = inProcChannel();
     const service = resolver();

packages/zero-cache/src/services/change-streamer/change-streamer-http.ts

@@ -10,9 +10,11 @@ import {type Worker} from '../../types/processes.ts';
 import {type ShardID} from '../../types/shards.ts';
 import {
   streamIn,
+  streamInBatches,
   streamOut,
   streamOutStringified,
   type Source,
+  type StreamInPayload,
 } from '../../types/streams.ts';
 import {URLParams} from '../../types/url-params.ts';
 import {installWebSocketReceiver} from '../../types/websocket-handoff.ts';
@@ -35,6 +37,7 @@ const MIN_SUPPORTED_PROTOCOL_VERSION = 1;
 const SNAPSHOT_PATH_PATTERN = '/replication/:version/snapshot';
 const CHANGES_PATH_PATTERN = '/replication/:version/changes';
 const PATH_REGEX = /\/replication\/v(?<version>\d+)\/(changes|snapshot)$/;
+const STREAM_BATCH_MESSAGES = 256;
 
 const SNAPSHOT_PATH = `/replication/v${PROTOCOL_VERSION}/snapshot`;
 const CHANGES_PATH = `/replication/v${PROTOCOL_VERSION}/changes`;
@@ -101,13 +104,6 @@ export class ChangeStreamerHttpServer extends HttpService {
         return this.#reserveSnapshot(ws, msg);
       case 'changes':
         return this.#subscribe(ws, msg);
-      default:
-        closeWithError(
-          this._lc,
-          ws,
-          `invalid action "${action}" received in handoff`,
-        );
-        return;
     }
   };
 
@@ -143,7 +139,20 @@ export class ChangeStreamerHttpServer extends HttpService {
         // end the reservation to safely resume scheduling cleanup.
         this.#backupMonitor.endReservation(ctx.taskID);
       }
-      void streamOutStringified(this._lc, downstream, ws);
+      const url = new URL(
+        req.url ?? '',
+        req.headers.origin ?? 'http://localhost',
+      );
+      // Serving replicas request streamBatch=1 so the RM can preserve a
+      // websocket batch boundary instead of forcing the VS to rediscover
+      // batching after parse/ACK. Older clients omit the flag and keep the
+      // original one-message stream shape.
+      const streamBatchRequested = url.searchParams.get('streamBatch') === '1';
+      void streamOutStringified(this._lc, downstream, ws, {
+        batch: streamBatchRequested
+          ? {maxMessages: STREAM_BATCH_MESSAGES}
+          : undefined,
+      });
     } catch (err) {
       closeWithError(this._lc, ws, err, PROTOCOL_ERROR);
     }
@@ -233,12 +242,25 @@ export class ChangeStreamerHttpClient implements ChangeStreamer {
   }
 
   async subscribe(ctx: SubscriberContext): Promise<Source<Downstream>> {
+    const ws = await this.#openChangesWebSocket(ctx);
+    return streamIn(this.#lc, ws, downstreamSchema);
+  }
+
+  async subscribeBatched(
+    ctx: SubscriberContext,
+  ): Promise<Source<StreamInPayload<Downstream>>> {
+    const ws = await this.#openChangesWebSocket(ctx);
+    // Keep websocket batches visible to the incremental syncer so it can hand a
+    // whole received frame to the write worker before ACKing the frame messages.
+    return streamInBatches(this.#lc, ws, downstreamSchema);
+  }
+
+  async #openChangesWebSocket(ctx: SubscriberContext) {
     const uri = await this.#resolveChangeStreamer(CHANGES_PATH);
 
     const params = getParams(ctx);
-    const ws = new WebSocket(uri + `?${params.toString()}`);
-
-    return streamIn(this.#lc, ws, downstreamSchema);
+    params.set('streamBatch', '1');
+    return new WebSocket(uri + `?${params.toString()}`);
   }
 }
 

packages/zero-cache/src/services/change-streamer/change-streamer-service.pg.test.ts

@@ -13,7 +13,7 @@ import {Database} from '../../../../zqlite/src/db.ts';
 import {StatementRunner} from '../../db/statements.ts';
 import {expectTables, test, type PgTest} from '../../test/db.ts';
 import type {PostgresDB} from '../../types/pg.ts';
-import type {Source} from '../../types/streams.ts';
+import type {Source, StringifiedStreamPayload} from '../../types/streams.ts';
 import {Subscription, type Result} from '../../types/subscription.ts';
 import type {ChangeSource} from '../change-source/change-source.ts';
 import {type ChangeStreamMessage} from '../change-source/protocol/current/downstream.ts';
@@ -110,11 +110,16 @@ describe('change-streamer/service', () => {
     };
   });
 
-  function drainToQueue(sub: Source<string>): Queue<Downstream> {
+  function drainToQueue(
+    sub: Source<StringifiedStreamPayload>,
+  ): Queue<Downstream> {
     const queue = new Queue<Downstream>();
     void (async () => {
-      for await (const msg of sub) {
-        queue.enqueue(BigIntJSON.parse(msg) as Downstream);
+      for await (const payload of sub) {
+        const entries = typeof payload === 'string' ? [payload] : payload;
+        for (const msg of entries) {
+          queue.enqueue(BigIntJSON.parse(msg) as Downstream);
+        }
       }
     })();
     return queue;
@@ -891,18 +896,20 @@ describe('change-streamer/service', () => {
     expect(setTimeoutFn).toHaveBeenCalledTimes(3);
 
     drainToQueue(sub1);
-    for await (const json of sub2) {
-      const msg: Downstream = BigIntJSON.parse(json) as Downstream;
+    const sub2Queue = drainToQueue(sub2);
+    for (;;) {
+      const msg = await sub2Queue.dequeue();
       if (msg[0] === 'commit' && msg[2].watermark === '08') {
-        // Now that sub2 has consumed past '06',
-        // a purge should successfully clear records before '06'
-        await (setTimeoutFn.mock.calls[2][0]() as unknown as Promise<void>);
-        expect(
-          await sql`SELECT watermark FROM "zoro_3/cdc"."changeLog"`.values(),
-        ).toEqual([['06'], ['07'], ['08']]);
         break;
       }
     }
+    // Now that sub2 has consumed past '06',
+    // a purge should successfully clear records before '06'
+    await (setTimeoutFn.mock.calls[2][0]() as unknown as Promise<void>);
+    expect(
+      await sql`SELECT watermark FROM "zoro_3/cdc"."changeLog"`.values(),
+    ).toEqual([['06'], ['07'], ['08']]);
+
     // replicationState is unaffected
     await expectTables(sql, {
       ['zoro_3/cdc.replicationState']: [

packages/zero-cache/src/services/change-streamer/change-streamer-service.ts

@@ -11,7 +11,7 @@ import {
 } from '../../types/lexi-version.ts';
 import type {PostgresDB} from '../../types/pg.ts';
 import type {ShardID} from '../../types/shards.ts';
-import type {Source} from '../../types/streams.ts';
+import type {Source, StringifiedStreamPayload} from '../../types/streams.ts';
 import {Subscription} from '../../types/subscription.ts';
 import type {
   ChangeSource,
@@ -523,12 +523,12 @@ class ChangeStreamerImpl implements ChangeStreamerService {
     }
   }
 
-  subscribe(ctx: SubscriberContext): Promise<Source<string>> {
+  subscribe(ctx: SubscriberContext): Promise<Source<StringifiedStreamPayload>> {
     const {protocolVersion, id, mode, replicaVersion, watermark} = ctx;
     if (mode === 'serving') {
       this.#serving.resolve();
     }
-    const downstream = Subscription.create<string>({
+    const downstream = Subscription.create<StringifiedStreamPayload>({
       cleanup: () => this.#forwarder.remove(subscriber),
     });
     const subscriber = new Subscriber(

packages/zero-cache/src/services/change-streamer/change-streamer.ts

@@ -1,6 +1,10 @@
 import type {Enum} from '../../../../shared/src/enum.ts';
 import * as v from '../../../../shared/src/valita.ts';
-import type {Source} from '../../types/streams.ts';
+import type {
+  Source,
+  StreamInPayload,
+  StringifiedStreamPayload,
+} from '../../types/streams.ts';
 import {changeStreamDataSchema} from '../change-source/protocol/current/downstream.ts';
 import type {ReplicatorMode} from '../replicator/replicator.ts';
 import {changeSourceTimingsSchema} from '../replicator/reporter/report-schema.ts';
@@ -75,7 +79,6 @@ export interface ChangeStreamer {
 //   - Adds support for `backfill` messages
 // v6: v1.0.1  (backwards compatible, no version change)
 //   - Adds lag reporting to status messages
-
 export const PROTOCOL_VERSION = 6;
 
 export type SubscriberContext = {
@@ -168,19 +171,6 @@ export const downstreamSchema = v.union(
 
 export type Error = v.Infer<typeof errorSchema>;
 
-export function errorTypeToReadableName(val: ErrorType) {
-  switch (val) {
-    case ErrorType.WrongReplicaVersion:
-      return 'WrongReplicaVersion';
-    case ErrorType.WatermarkTooOld:
-      return 'WatermarkTooOld';
-    case ErrorType.Unknown:
-      return 'Unknown';
-    default:
-      return 'Unknown';
-  }
-}
-
 /**
  * A stream of transactions, each starting with a {@link Begin} message,
  * containing one or more {@link Data} messages, and ending with a
@@ -194,13 +184,40 @@ export function errorTypeToReadableName(val: ErrorType) {
  */
 export type Downstream = v.Infer<typeof downstreamSchema>;
 
+export interface BatchedChangeStreamer extends ChangeStreamer {
+  subscribeBatched(
+    ctx: SubscriberContext,
+  ): Promise<Source<StreamInPayload<Downstream>>>;
+}
+
+export function hasBatchedSubscribe(
+  changeStreamer: ChangeStreamer,
+): changeStreamer is BatchedChangeStreamer {
+  return 'subscribeBatched' in changeStreamer;
+}
+
+export function errorTypeToReadableName(val: ErrorType) {
+  switch (val) {
+    case ErrorType.WrongReplicaVersion:
+      return 'WrongReplicaVersion';
+    case ErrorType.WatermarkTooOld:
+      return 'WatermarkTooOld';
+    case ErrorType.Unknown:
+      return 'Unknown';
+    default:
+      return 'Unknown';
+  }
+}
+
 export interface ChangeStreamerService
   extends Omit<ChangeStreamer, 'subscribe'>, Service {
   /**
    * The server-side interface overrides `subscribe()` to return a stream
-   * of already-stringified {@link Downstream} payloads.
+   * of already-stringified {@link Downstream} payloads. A payload can contain
+   * several ordered messages so the stream layer can flatten them into the
+   * normal wire format without adding one producer queue entry per message.
    */
-  subscribe(ctx: SubscriberContext): Promise<Source<string>>;
+  subscribe(ctx: SubscriberContext): Promise<Source<StringifiedStreamPayload>>;
 
   /**
    * Notifies the change streamer of a watermark that has been backed up,

packages/zero-cache/src/services/change-streamer/storer.pg.test.ts

@@ -7,6 +7,7 @@ import {Queue} from '../../../../shared/src/queue.ts';
 import {sleep} from '../../../../shared/src/sleep.ts';
 import {type PgTest, test} from '../../test/db.ts';
 import type {PostgresDB} from '../../types/pg.ts';
+import type {StringifiedStreamPayload} from '../../types/streams.ts';
 import type {Subscription} from '../../types/subscription.ts';
 import {
   type ChangeStreamData,
@@ -93,13 +94,19 @@ describe('change-streamer/storer', () => {
 
   const messages = new ReplicationMessages({issues: 'id'});
 
-  async function drain(sub: Subscription<string>, untilWatermark?: string) {
+  async function drain(
+    sub: Subscription<StringifiedStreamPayload>,
+    untilWatermark?: string,
+  ) {
     const msgs: Downstream[] = [];
-    for await (const json of sub) {
-      const msg: Downstream = JSON.parse(json);
-      msgs.push(msg);
-      if (msg[0] === 'commit' && msg[2].watermark === untilWatermark) {
-        return msgs;
+    for await (const payload of sub) {
+      const entries = typeof payload === 'string' ? [payload] : payload;
+      for (const json of entries) {
+        const msg: Downstream = JSON.parse(json);
+        msgs.push(msg);
+        if (msg[0] === 'commit' && msg[2].watermark === untilWatermark) {
+          return msgs;
+        }
       }
     }
     return msgs;