From 7338bf8af4bbdca7ca9d3725875d3919885e45e4 Mon Sep 17 00:00:00 2001 From: Alp Karavil <46430499+Karavil@users.noreply.github.com> Date: Tue, 26 May 2026 13:31:53 -0400 Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=20perf:=20batch=20replication=20write?= =?UTF-8?q?-worker=20messages?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Why: avoid one worker-thread round trip for every RM stream row while preserving commit and rollback boundaries * add processMessages() to the write-worker client protocol * add WorkerMessageBatcher for bounded transaction-shaped batches * expose worker-batch mode in the RM vs load benchmark --- .../services/replicator/incremental-sync.ts | 107 +++++++++++------- .../replicator/worker-message-batcher.test.ts | 86 ++++++++++++++ .../replicator/worker-message-batcher.ts | 77 +++++++++++++ .../replicator/write-worker-client.ts | 29 ++++- .../services/replicator/write-worker.test.ts | 72 +++++++++++- .../src/services/replicator/write-worker.ts | 4 + 6 files changed, 331 insertions(+), 44 deletions(-) create mode 100644 packages/zero-cache/src/services/replicator/worker-message-batcher.test.ts create mode 100644 packages/zero-cache/src/services/replicator/worker-message-batcher.ts diff --git a/packages/zero-cache/src/services/replicator/incremental-sync.ts b/packages/zero-cache/src/services/replicator/incremental-sync.ts index 810b5c6396..217c2577ef 100644 --- a/packages/zero-cache/src/services/replicator/incremental-sync.ts +++ b/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; +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 | 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 | 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; diff --git a/packages/zero-cache/src/services/replicator/worker-message-batcher.test.ts b/packages/zero-cache/src/services/replicator/worker-message-batcher.test.ts new file mode 100644 index 0000000000..623dc89651 --- /dev/null +++ b/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}; +} diff --git a/packages/zero-cache/src/services/replicator/worker-message-batcher.ts b/packages/zero-cache/src/services/replicator/worker-message-batcher.ts new file mode 100644 index 0000000000..1a938e38af --- /dev/null +++ b/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 | undefined { + this.#messages.push(message); + return this.#shouldFlush(message) ? this.flush() : undefined; + } + + flush(): Promise | 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; + } +} diff --git a/packages/zero-cache/src/services/replicator/write-worker-client.ts b/packages/zero-cache/src/services/replicator/write-worker-client.ts index ff4b5044c8..7301eb76cd 100644 --- a/packages/zero-cache/src/services/replicator/write-worker-client.ts +++ b/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; processMessage(downstream: ChangeStreamData): Promise; + processMessages( + downstreams: readonly ChangeStreamData[], + ): Promise; abort(): void; stop(): Promise; 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 { + 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)); +} diff --git a/packages/zero-cache/src/services/replicator/write-worker.test.ts b/packages/zero-cache/src/services/replicator/write-worker.test.ts index 933c66edd8..9ef57137f1 100644 --- a/packages/zero-cache/src/services/replicator/write-worker.test.ts +++ b/packages/zero-cache/src/services/replicator/write-worker.test.ts @@ -97,6 +97,74 @@ describe('write-worker', () => { expect(state.watermark).toBe('06'); }); + test('processMessages handles a full transaction batch', async () => { + const issues = new ReplicationMessages({issues: ['issueID', 'bool']}); + + const messages: ChangeStreamData[] = [ + ['begin', issues.begin(), {commitWatermark: '06'}], + ['data', issues.insert('issues', {issueID: 123, bool: true})], + ['data', issues.insert('issues', {issueID: 456, bool: false})], + ['commit', issues.commit(), {watermark: '06'}], + ]; + + await expect(worker.processMessages(messages)).resolves.toEqual({ + watermark: '06', + completedBackfill: undefined, + schemaUpdated: false, + changeLogUpdated: true, + }); + + const rows = mainDb + .prepare('SELECT issueID, bool, _0_version FROM issues ORDER BY issueID') + .all(); + expect(rows).toEqual([ + {issueID: 123, bool: 1, _0_version: '06'}, + {issueID: 456, bool: 0, _0_version: '06'}, + ]); + + const state = await worker.getSubscriptionState(); + expect(state.watermark).toBe('06'); + }); + + test('processMessages returns every committed transaction in a batch', async () => { + const issues = new ReplicationMessages({issues: ['issueID', 'bool']}); + + const messages: ChangeStreamData[] = [ + ['begin', issues.begin(), {commitWatermark: '06'}], + ['data', issues.insert('issues', {issueID: 123, bool: true})], + ['commit', issues.commit(), {watermark: '06'}], + ['begin', issues.begin(), {commitWatermark: '07'}], + ['data', issues.insert('issues', {issueID: 456, bool: false})], + ['commit', issues.commit(), {watermark: '07'}], + ]; + + await expect(worker.processMessages(messages)).resolves.toEqual([ + { + watermark: '06', + completedBackfill: undefined, + schemaUpdated: false, + changeLogUpdated: true, + }, + { + watermark: '07', + completedBackfill: undefined, + schemaUpdated: false, + changeLogUpdated: true, + }, + ]); + + const rows = mainDb + .prepare('SELECT issueID, bool, _0_version FROM issues ORDER BY issueID') + .all(); + expect(rows).toEqual([ + {issueID: 123, bool: 1, _0_version: '06'}, + {issueID: 456, bool: 0, _0_version: '07'}, + ]); + + const state = await worker.getSubscriptionState(); + expect(state.watermark).toBe('07'); + }); + test('abort rolls back pending transaction', async () => { const issues = new ReplicationMessages({issues: ['issueID', 'bool']}); @@ -134,7 +202,7 @@ describe('write-worker', () => { }); test('stop shuts down cleanly', async () => { - await worker.stop(); + await expect(worker.stop()).resolves.toBeUndefined(); // Create a new worker for afterEach cleanup worker = new ThreadWriteWorkerClient(); await worker.init( @@ -172,7 +240,7 @@ describe('write-worker', () => { new: {id: [1, 'int4']}, }, ]), - ).rejects.toThrow(); + ).rejects.toThrow(/Received message outside of transaction/); expect(errorReceived).toBeDefined(); }); diff --git a/packages/zero-cache/src/services/replicator/write-worker.ts b/packages/zero-cache/src/services/replicator/write-worker.ts index 194f71792e..0fa76931b7 100644 --- a/packages/zero-cache/src/services/replicator/write-worker.ts +++ b/packages/zero-cache/src/services/replicator/write-worker.ts @@ -65,6 +65,10 @@ function createAPI(): API { return must(processor).processMessage(must(lc), downstream); }, + processMessages(downstreams: readonly ChangeStreamData[]) { + return must(processor).processMessages(must(lc), downstreams); + }, + abort() { must(processor).abort(must(lc)); createProcessor();