perf(zero-cache): batch replication write-worker messages

packages/zero-cache/src/services/replicator/incremental-sync.ts

@@ -19,10 +19,13 @@ import type {ReplicationStatusPublisher} from './replication-status.ts';
 import type {ReplicaState, ReplicatorMode} from './replicator.ts';
 import {ReplicationReportRecorder} from './reporter/recorder.ts';
 import type {ReplicationReport} from './reporter/report-schema.ts';
+import {WorkerMessageBatcher} from './worker-message-batcher.ts';
 import type {WriteWorkerClient} from './write-worker-client.ts';
 
 type ErrorType = Enum<typeof ErrorType>;
 
+const MAX_WORKER_BATCH_MESSAGES = 64;
+
 /**
  * The {@link IncrementalSyncer} manages a logical replication stream from upstream,
  * handling application lifecycle events (start, stop) and retrying the
@@ -85,7 +88,7 @@ export class IncrementalSyncer {
 
       let downstream: Source<Downstream> | undefined;
       let unregister = () => {};
-      let err: unknown | undefined;
+      let err: unknown;
 
       try {
         downstream = await this.#changeStreamer.subscribe({
@@ -106,6 +109,66 @@ export class IncrementalSyncer {
         );
 
         let backfillStatus: DownloadStatus | undefined;
+        const workerBatch = new WorkerMessageBatcher(
+          this.#worker,
+          MAX_WORKER_BATCH_MESSAGES,
+        );
+        const handleWorkerResult = (
+          result: CommitResult | readonly CommitResult[] | null,
+        ) => {
+          if (result === null) {
+            return;
+          }
+          const results = Array.isArray(result) ? result : [result];
+          for (const commitResult of results) {
+            this.#handleResult(lc, commitResult);
+            if (commitResult?.completedBackfill) {
+              backfillStatus = undefined;
+            }
+          }
+        };
+        const flushWorkerBatch = async () => {
+          const result = workerBatch.flush();
+          if (result) {
+            handleWorkerResult(await result);
+          }
+        };
+        const processChangeStreamData = (
+          message: ChangeStreamData,
+        ): Promise<void> | undefined => {
+          const msg = message[1];
+          if (msg.tag === 'backfill' && msg.status) {
+            const {status} = msg;
+            if (!backfillStatus) {
+              // Start publishing the status every 3 seconds.
+              backfillStatus = status;
+              this.#statusPublisher?.publish(
+                lc,
+                'Replicating',
+                `Backfilling ${msg.relation.name} table`,
+                3000,
+                () =>
+                  backfillStatus
+                    ? {
+                        downloadStatus: [
+                          {
+                            ...backfillStatus,
+                            table: msg.relation.name,
+                            columns: [
+                              ...msg.relation.rowKey.columns,
+                              ...msg.columns,
+                            ],
+                          },
+                        ],
+                      }
+                    : {},
+              );
+            }
+            backfillStatus = status; // Update the current status
+          }
+
+          return workerBatch.push(message)?.then(handleWorkerResult);
+        };
 
         for await (const message of downstream) {
           this.#replicationEvents.add(1);
@@ -140,49 +203,15 @@ export class IncrementalSyncer {
               break;
             }
             default: {
-              const msg = message[1];
-              if (msg.tag === 'backfill' && msg.status) {
-                const {status} = msg;
-                if (!backfillStatus) {
-                  // Start publishing the status every 3 seconds.
-                  backfillStatus = status;
-                  this.#statusPublisher?.publish(
-                    lc,
-                    'Replicating',
-                    `Backfilling ${msg.relation.name} table`,
-                    3000,
-                    () =>
-                      backfillStatus
-                        ? {
-                            downloadStatus: [
-                              {
-                                ...backfillStatus,
-                                table: msg.relation.name,
-                                columns: [
-                                  ...msg.relation.rowKey.columns,
-                                  ...msg.columns,
-                                ],
-                              },
-                            ],
-                          }
-                        : {},
-                  );
-                }
-                backfillStatus = status; // Update the current status
-              }
-
-              const result = await this.#worker.processMessage(
-                message as ChangeStreamData,
-              );
-
-              this.#handleResult(lc, result);
-              if (result?.completedBackfill) {
-                backfillStatus = undefined;
+              const result = processChangeStreamData(message);
+              if (result) {
+                await result;
               }
               break;
             }
           }
         }
+        await flushWorkerBatch();
         this.#worker.abort();
       } catch (e) {
         err = e;

packages/zero-cache/src/services/replicator/worker-message-batcher.test.ts

@@ -0,0 +1,86 @@
+import {expect, test, vi} from 'vitest';
+import type {ChangeStreamData} from '../change-source/protocol/current/downstream.ts';
+import {WorkerMessageBatcher} from './worker-message-batcher.ts';
+import type {WriteWorkerClient} from './write-worker-client.ts';
+
+const begin = [
+  'begin',
+  {tag: 'begin'},
+  {commitWatermark: '02'},
+] satisfies ChangeStreamData;
+const row = [
+  'data',
+  {
+    tag: 'insert',
+    relation: {
+      schema: 'public',
+      name: 'issues',
+      rowKey: {columns: ['id']},
+    },
+    new: {id: '1'},
+  },
+] satisfies ChangeStreamData;
+const commit = [
+  'commit',
+  {tag: 'commit'},
+  {watermark: '02'},
+] satisfies ChangeStreamData;
+
+test('batches until transaction boundary', async () => {
+  const {worker, processMessages} = mockWorker();
+  const batcher = new WorkerMessageBatcher(worker, 64);
+
+  expect(batcher.push(begin)).toBeUndefined();
+  expect(batcher.push(row)).toBeUndefined();
+  expect(processMessages).not.toHaveBeenCalled();
+
+  await batcher.push(commit);
+
+  expect(processMessages).toHaveBeenCalledTimes(1);
+  expect(processMessages).toHaveBeenCalledWith([begin, row, commit]);
+  expect(batcher.size).toBe(0);
+});
+
+test('flushes large batches before commit', async () => {
+  const {worker, processMessages} = mockWorker();
+  const batcher = new WorkerMessageBatcher(worker, 2);
+
+  expect(batcher.push(begin)).toBeUndefined();
+  await batcher.push(row);
+
+  expect(processMessages).toHaveBeenCalledTimes(1);
+  expect(processMessages).toHaveBeenCalledWith([begin, row]);
+  expect(batcher.size).toBe(0);
+});
+
+test('can defer commit flush until caller finishes a stream batch', async () => {
+  const {worker, processMessages} = mockWorker();
+  const batcher = new WorkerMessageBatcher(worker, 64, {
+    flushOnCommit: false,
+  });
+
+  expect(batcher.push(begin)).toBeUndefined();
+  expect(batcher.push(row)).toBeUndefined();
+  expect(batcher.push(commit)).toBeUndefined();
+  expect(processMessages).not.toHaveBeenCalled();
+  expect(batcher.size).toBe(3);
+
+  await batcher.flush();
+
+  expect(processMessages).toHaveBeenCalledTimes(1);
+  expect(processMessages).toHaveBeenCalledWith([begin, row, commit]);
+  expect(batcher.size).toBe(0);
+});
+
+function mockWorker() {
+  const processMessages = vi.fn().mockResolvedValue(null);
+  const worker: WriteWorkerClient = {
+    getSubscriptionState: vi.fn(),
+    processMessage: vi.fn(),
+    processMessages,
+    abort: vi.fn(),
+    stop: vi.fn(),
+    onError: vi.fn(),
+  };
+  return {worker, processMessages};
+}

packages/zero-cache/src/services/replicator/worker-message-batcher.ts

@@ -0,0 +1,77 @@
+import type {ChangeStreamData} from '../change-source/protocol/current/downstream.ts';
+import type {CommitResult} from './change-processor.ts';
+import type {WriteWorkerClient} from './write-worker-client.ts';
+
+type ProcessMessagesResult = CommitResult | readonly CommitResult[] | null;
+
+type WorkerMessageBatcherOptions = {
+  // Legacy subscribe() delivers one message at a time, so commits still need to
+  // flush immediately to keep replication status and backfill completion moving.
+  // Batched subscribe() preserves RM websocket frames and flushes after the frame.
+  readonly flushOnCommit?: boolean | undefined;
+};
+
+// RM -> VS traffic arrives in ordered stream batches. Grouping write-worker
+// handoff around those batches keeps row-heavy transactions from paying a
+// promise/IPC boundary per row; callers that need per-commit ACKs can leave
+// flushOnCommit enabled.
+export class WorkerMessageBatcher {
+  readonly #worker: WriteWorkerClient;
+  readonly #maxMessages: number;
+  readonly #flushOnCommit: boolean;
+  #messages: ChangeStreamData[] = [];
+
+  constructor(
+    worker: WriteWorkerClient,
+    maxMessages: number,
+    options: WorkerMessageBatcherOptions = {},
+  ) {
+    this.#worker = worker;
+    this.#maxMessages = maxMessages;
+    this.#flushOnCommit = options.flushOnCommit ?? true;
+  }
+
+  get size() {
+    return this.#messages.length;
+  }
+
+  push(message: ChangeStreamData): Promise<ProcessMessagesResult> | undefined {
+    this.#messages.push(message);
+    return this.#shouldFlush(message) ? this.flush() : undefined;
+  }
+
+  flush(): Promise<ProcessMessagesResult> | undefined {
+    if (this.#messages.length === 0) {
+      return undefined;
+    }
+    const messages = this.#messages;
+    this.#messages = [];
+    return this.#worker.processMessages(messages);
+  }
+
+  clear() {
+    this.#messages = [];
+  }
+
+  #shouldFlush(message: ChangeStreamData) {
+    const tag = message[0];
+    if (this.#flushOnCommit && tag === 'commit') {
+      // A commit is the first point where the write worker can return a durable
+      // watermark, schema-update, or backfill-complete result. Callers on the
+      // one-message-at-a-time path must see that result before later stream work.
+      return true;
+    }
+    if (tag === 'rollback') {
+      // Rollback terminates the current upstream transaction without a commit
+      // result. Flushing here keeps rolled-back rows from sharing a worker call
+      // with the next transaction, which preserves the stream's transaction shape.
+      return true;
+    }
+    if (this.#messages.length >= this.#maxMessages) {
+      // The batcher is meant to remove per-row worker calls, not to become an
+      // unbounded in-memory queue if a very large transaction arrives.
+      return true;
+    }
+    return false;
+  }
+}

packages/zero-cache/src/services/replicator/write-worker-client.ts

@@ -11,6 +11,7 @@ import type {SubscriptionState} from './schema/replication-state.ts';
 export type PragmaConfig = {
   busyTimeout: number;
   analysisLimit: number;
+  synchronous?: 'OFF' | 'NORMAL' | 'FULL' | undefined;
   walAutocheckpoint?: number | undefined;
 };
 
@@ -22,6 +23,9 @@ type ErrorHandler = (err: Error) => void;
 export interface WriteWorkerClient {
   getSubscriptionState(): Promise<SubscriptionState>;
   processMessage(downstream: ChangeStreamData): Promise<CommitResult | null>;
+  processMessages(
+    downstreams: readonly ChangeStreamData[],
+  ): Promise<CommitResult | readonly CommitResult[] | null>;
   abort(): void;
   stop(): Promise<void>;
   onError(handler: ErrorHandler): void;
@@ -32,6 +36,7 @@ export type ArgsMap = {
   init: [string, ChangeProcessorMode, PragmaConfig, LogConfig];
   getSubscriptionState: [];
   processMessage: [ChangeStreamData];
+  processMessages: [readonly ChangeStreamData[]];
   abort: [];
   stop: [];
 };
@@ -44,6 +49,7 @@ export type ResultMap = {
   init: void;
   getSubscriptionState: SubscriptionState;
   processMessage: CommitResult | null;
+  processMessages: CommitResult | readonly CommitResult[] | null;
   abort: void;
   stop: void;
 };
@@ -57,6 +63,9 @@ export type WriteError = {writeError: Error};
 export function applyPragmas(db: Database, pragmas: PragmaConfig) {
   db.pragma(`busy_timeout = ${pragmas.busyTimeout}`);
   db.pragma(`analysis_limit = ${pragmas.analysisLimit}`);
+  if (pragmas.synchronous !== undefined) {
+    db.pragma(`synchronous = ${pragmas.synchronous}`);
+  }
   if (pragmas.walAutocheckpoint !== undefined) {
     db.pragma(`wal_autocheckpoint = ${pragmas.walAutocheckpoint}`);
   }
@@ -89,9 +98,7 @@ export class ThreadWriteWorkerClient implements WriteWorkerClient {
       if (!r) return; // stale abort response
       this.#pending = null;
       if (msg.error !== undefined) {
-        r.reject(
-          msg.error instanceof Error ? msg.error : new Error(String(msg.error)),
-        );
+        r.reject(ensureError(msg.error));
       } else {
         r.resolve(msg.result);
       }
@@ -145,6 +152,12 @@ export class ThreadWriteWorkerClient implements WriteWorkerClient {
     return this.#call('processMessage', [downstream]);
   }
 
+  processMessages(
+    downstreams: readonly ChangeStreamData[],
+  ): Promise<CommitResult | readonly CommitResult[] | null> {
+    return this.#call('processMessages', [downstreams]);
+  }
+
   abort(): void {
     if (!this.#terminated) {
       this.#worker.postMessage({method: 'abort', args: []} satisfies Request);
@@ -162,3 +175,13 @@ export class ThreadWriteWorkerClient implements WriteWorkerClient {
     this.#errorHandler = handler;
   }
 }
+
+function ensureError(err: unknown): Error {
+  if (err instanceof Error) {
+    return err;
+  }
+  if (typeof err === 'string') {
+    return new Error(err);
+  }
+  return new Error(JSON.stringify(err));
+}