diff --git a/packages/zero-cache/src/services/change-streamer/broadcast.test.ts b/packages/zero-cache/src/services/change-streamer/broadcast.test.ts index f3e962581b..dadd612c0c 100644 --- a/packages/zero-cache/src/services/change-streamer/broadcast.test.ts +++ b/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) { + async function consumeOne( + downstream: Subscription, + ) { 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[] = []; + const downstreams: Subscription[] = []; 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) { + async function consumeChange( + downstream: Subscription, + ) { await consumeOne(downstream); await flushCompletions(); } diff --git a/packages/zero-cache/src/services/change-streamer/change-streamer-http.pg.test.ts b/packages/zero-cache/src/services/change-streamer/change-streamer-http.pg.test.ts index ea6159caea..4fbc9b4bfa 100644 --- a/packages/zero-cache/src/services/change-streamer/change-streamer-http.pg.test.ts +++ b/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(); 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(); diff --git a/packages/zero-cache/src/services/change-streamer/change-streamer-http.ts b/packages/zero-cache/src/services/change-streamer/change-streamer-http.ts index 1fae183d55..f92f15da7c 100644 --- a/packages/zero-cache/src/services/change-streamer/change-streamer-http.ts +++ b/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(?\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> { + const ws = await this.#openChangesWebSocket(ctx); + return streamIn(this.#lc, ws, downstreamSchema); + } + + async subscribeBatched( + ctx: SubscriberContext, + ): Promise>> { + 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()}`); } } diff --git a/packages/zero-cache/src/services/change-streamer/change-streamer-service.pg.test.ts b/packages/zero-cache/src/services/change-streamer/change-streamer-service.pg.test.ts index 06c7c57a4c..5b66008d89 100644 --- a/packages/zero-cache/src/services/change-streamer/change-streamer-service.pg.test.ts +++ b/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): Queue { + function drainToQueue( + sub: Source, + ): Queue { const queue = new Queue(); 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); - 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); + 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']: [ diff --git a/packages/zero-cache/src/services/change-streamer/change-streamer-service.ts b/packages/zero-cache/src/services/change-streamer/change-streamer-service.ts index 6918013f87..29a20701ee 100644 --- a/packages/zero-cache/src/services/change-streamer/change-streamer-service.ts +++ b/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> { + subscribe(ctx: SubscriberContext): Promise> { const {protocolVersion, id, mode, replicaVersion, watermark} = ctx; if (mode === 'serving') { this.#serving.resolve(); } - const downstream = Subscription.create({ + const downstream = Subscription.create({ cleanup: () => this.#forwarder.remove(subscriber), }); const subscriber = new Subscriber( diff --git a/packages/zero-cache/src/services/change-streamer/change-streamer.ts b/packages/zero-cache/src/services/change-streamer/change-streamer.ts index 53e6fd2df5..fad55c84c3 100644 --- a/packages/zero-cache/src/services/change-streamer/change-streamer.ts +++ b/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; -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; +export interface BatchedChangeStreamer extends ChangeStreamer { + subscribeBatched( + ctx: SubscriberContext, + ): Promise>>; +} + +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, 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>; + subscribe(ctx: SubscriberContext): Promise>; /** * Notifies the change streamer of a watermark that has been backed up, diff --git a/packages/zero-cache/src/services/change-streamer/storer.pg.test.ts b/packages/zero-cache/src/services/change-streamer/storer.pg.test.ts index 6002408a1f..a67d804ce4 100644 --- a/packages/zero-cache/src/services/change-streamer/storer.pg.test.ts +++ b/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, untilWatermark?: string) { + async function drain( + sub: Subscription, + 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; diff --git a/packages/zero-cache/src/services/change-streamer/subscriber.test.ts b/packages/zero-cache/src/services/change-streamer/subscriber.test.ts index 586bbdacfe..7a2f67e4fd 100644 --- a/packages/zero-cache/src/services/change-streamer/subscriber.test.ts +++ b/packages/zero-cache/src/services/change-streamer/subscriber.test.ts @@ -1,9 +1,14 @@ import {describe, expect, test} from 'vitest'; +import type {StringifiedStreamPayload} from '../../types/streams.ts'; import {ReplicationMessages} from '../replicator/test-utils.ts'; import {createSubscriber} from './test-utils.ts'; const json = JSON.stringify; +function expandPayload(payload: StringifiedStreamPayload): readonly string[] { + return typeof payload === 'string' ? [payload] : payload; +} + describe('change-streamer/subscriber', () => { const messages = new ReplicationMessages({issues: 'id'}); @@ -284,29 +289,35 @@ describe('change-streamer/subscriber', () => { expect(sub.numPending).toBe(pending); let txNum = 0; - for await (const json of receiver) { - const msg = JSON.parse(json); - expect(sub.numProcessed).toBe(processed++); - expect(sub.numPending).toBe(pending--); - - if (msg[0] === 'begin') { - txNum++; - } - switch (txNum) { - case 1: - expect(sub.acked).toBe('00'); - break; - case 2: - expect(sub.acked).toBe('02'); - break; - case 3: - expect(sub.acked).toBe('12'); - break; - case 4: - expect(sub.acked).toBe('22'); - sub.close(); - break; + for await (const payload of receiver) { + const beforeProcessed = processed; + const beforePending = pending; + const messages = expandPayload(payload); + for (const json of messages) { + const msg = JSON.parse(json); + if (msg[0] === 'begin') { + txNum++; + } + switch (txNum) { + case 1: + expect(sub.acked).toBe('00'); + break; + case 2: + expect(sub.acked).toBe('02'); + break; + case 3: + expect(sub.acked).toBe('12'); + break; + case 4: + expect(sub.acked).toBe('22'); + sub.close(); + break; + } } + expect(sub.numProcessed).toBe(beforeProcessed); + expect(sub.numPending).toBe(beforePending); + processed += messages.length; + pending -= messages.length; } expect(sub.numProcessed).toBe(8); expect( diff --git a/packages/zero-cache/src/services/change-streamer/subscriber.ts b/packages/zero-cache/src/services/change-streamer/subscriber.ts index 715f3fa7fe..ff6391776f 100644 --- a/packages/zero-cache/src/services/change-streamer/subscriber.ts +++ b/packages/zero-cache/src/services/change-streamer/subscriber.ts @@ -3,6 +3,7 @@ import {BigIntJSON} from '../../../../shared/src/bigint-json.ts'; import type {Enum} from '../../../../shared/src/enum.ts'; import {must} from '../../../../shared/src/must.ts'; import {max} from '../../types/lexi-version.ts'; +import type {StringifiedStreamPayload} from '../../types/streams.ts'; import type {Subscription} from '../../types/subscription.ts'; import type {ChangeTag, WatermarkedChange} from './change-streamer-service.ts'; import {type Downstream, type Status} from './change-streamer.ts'; @@ -20,17 +21,22 @@ type ErrorType = Enum; export class Subscriber { readonly #protocolVersion: number; readonly id: string; - readonly #downstream: Subscription; + readonly #downstream: Subscription; readonly #latestStatus: () => Status; + // #watermark is the latest commit sent to this subscriber. #acked lags until + // the downstream stream confirms that the commit frame was consumed. #watermark: string; #acked: string; + // Non-null while catchup is reading historical changeLog rows. Live forwarded + // changes are buffered here so the subscriber sees "historical, then live" in + // one ordered stream. #backlog: WatermarkedChange[] | null; constructor( protocolVersion: number, id: string, watermark: string, - downstream: Subscription, + downstream: Subscription, latestStatus: () => Status, ) { this.#protocolVersion = protocolVersion; @@ -54,15 +60,19 @@ export class Subscriber { const [watermark] = change; if (watermark > this.#watermark) { if (this.#backlog) { + // Catchup has not finished; keep live traffic behind the historical + // replay so the subscriber never sees a gap. this.#backlog.push(change); } else { - await this.#sendChange(change); + await this.#sendChanges([change]); } } } async sendBatch(changes: readonly WatermarkedChange[]) { if (this.#backlog) { + // Entries at or before the subscriber watermark were already represented + // in its replica and must not be replayed during catchup handoff. for (const change of changes) { if (change[0] > this.#watermark) { this.#backlog.push(change); @@ -70,7 +80,7 @@ export class Subscriber { } return; } - await Promise.all(changes.map(change => this.#sendChange(change))); + await this.#sendChanges(changes); } #initialized = false; @@ -88,6 +98,8 @@ export class Subscriber { sendStatus(status: Status) { if (this.#protocolVersion >= 2 && this.#initialized) { + // Older protocol versions do not know status messages. Before initialize, + // the subscriber's requested watermark has not been validated yet. void this.#sendDownstream(['status', status]); } } @@ -95,15 +107,13 @@ export class Subscriber { /** catchup() is called on ChangeEntries loaded from the store. */ async catchup(change: WatermarkedChange) { this.#initialize(); - await this.#sendChange(change); + await this.#sendChanges([change]); } /** catchupBatch() is called on ChangeEntries loaded from the store. */ async catchupBatch(changes: readonly WatermarkedChange[]) { this.#initialize(); - for (const change of changes) { - await this.#sendChange(change); - } + await this.#sendChanges(changes); } /** @@ -120,26 +130,55 @@ export class Subscriber { // interpret the #backlog variable correctly. This is the only place // where I/O flow control is not heeded. However, it will be awaited // by the next caller to send(). - for (const change of this.#backlog) { - void this.#sendChange(change); + const backlog = this.#backlog; + for (let i = 0; i < backlog.length; i += 64) { + // Send fixed-size ranges from the existing array instead of slice()ing. + // Catchup handoff can hold thousands of live entries under load, and the + // range form avoids duplicating those references during the flush. + void this.#sendChanges(backlog, i, i + 64); } this.#backlog = null; } - async #sendChange(change: WatermarkedChange) { - const [watermark, tag, json] = change; - if (watermark <= this.watermark) { - return; + async #sendChanges( + changes: readonly WatermarkedChange[], + start = 0, + end = changes.length, + ) { + const json: string[] = []; + let commitWatermark: string | undefined; + for (let i = start; i < end && i < changes.length; i++) { + const change = changes[i]; + const [watermark, tag, payload] = change; + if (watermark <= this.watermark) { + // Catchup queries include the row at the requested watermark to prove + // continuity. The subscriber only needs changes after that point. + continue; + } + if (!this.supportsMessage(tag)) { + // Protocol-version filtering happens before JSON write so old serving + // replicas can continue catching up through newer metadata events. + continue; + } + if (tag === 'commit') { + // Only commit messages advance the public replication position; data + // messages carry an internal ordering watermark but are not resumable. + this.#watermark = watermark; + commitWatermark = watermark; + } + json.push(payload); } - if (!this.supportsMessage(tag)) { + if (json.length === 0) { + // Nothing survived the watermark/protocol filters, so there is no commit + // ACK to wait for and no downstream flow-control work to track. return; } - if (tag === 'commit') { - this.#watermark = watermark; - } - const result = await this.#sendStringifiedDownstream(json); - if (tag === 'commit' && result === 'consumed') { - this.#acked = max(this.#acked, watermark); + const payload = json.length === 1 ? json[0] : json; + const result = await this.#sendStringifiedDownstream(payload, json.length); + if (commitWatermark !== undefined && result === 'consumed') { + // Purge logic uses acked watermarks, so advance only after the downstream + // stream confirms that the commit frame was consumed. + this.#acked = max(this.#acked, commitWatermark); } } @@ -147,14 +186,19 @@ export class Subscriber { return this.#sendStringifiedDownstream(BigIntJSON.stringify(downstream)); } - async #sendStringifiedDownstream(json: string) { - this.#pending++; - const {result} = this.#downstream.push(json); + async #sendStringifiedDownstream( + payload: StringifiedStreamPayload, + messageCount = typeof payload === 'string' ? 1 : payload.length, + ) { + // pending/processed track logical downstream messages, not websocket frames; + // a stringified array still represents several change-stream messages. + this.#pending += messageCount; + const {result} = this.#downstream.push(payload); try { return await result; } finally { - this.#pending--; - this.#processed++; + this.#pending -= messageCount; + this.#processed += messageCount; } } @@ -215,7 +259,8 @@ export class Subscriber { supportsMessage(tag: ChangeTag) { switch (tag) { case 'update-table-metadata': - // update-table-row-key is only understood by subscribers >= protocol v5 + // Protocol v5 introduced relation row-key/table-metadata updates. + // Earlier replicas must skip this message instead of failing catchup. return this.#protocolVersion >= 5; } return true; diff --git a/packages/zero-cache/src/services/change-streamer/test-utils.ts b/packages/zero-cache/src/services/change-streamer/test-utils.ts index 97ea2cbaf1..2eb942b412 100644 --- a/packages/zero-cache/src/services/change-streamer/test-utils.ts +++ b/packages/zero-cache/src/services/change-streamer/test-utils.ts @@ -1,3 +1,4 @@ +import type {StringifiedStreamPayload} from '../../types/streams.ts'; import {Subscription} from '../../types/subscription.ts'; import {PROTOCOL_VERSION, type Downstream} from './change-streamer.ts'; import {Subscriber} from './subscriber.ts'; @@ -7,14 +8,15 @@ let nextID = 1; export function createSubscriber( watermark = '00', caughtUp = false, -): [Subscriber, Downstream[], Subscription] { + protocolVersion = PROTOCOL_VERSION, +): [Subscriber, Downstream[], Subscription] { const id = '' + nextID++; const received: Downstream[] = []; - const sub = Subscription.create({ - cleanup: unconsumed => received.push(...unconsumed.map(m => JSON.parse(m))), + const sub = Subscription.create({ + cleanup: unconsumed => received.push(...unconsumed.flatMap(parsePayload)), }); const subscriber = new Subscriber( - PROTOCOL_VERSION, + protocolVersion, id, watermark, sub, @@ -26,3 +28,9 @@ export function createSubscriber( return [subscriber, received, sub]; } + +function parsePayload(payload: StringifiedStreamPayload): Downstream[] { + return (typeof payload === 'string' ? [payload] : payload).map( + m => JSON.parse(m) as Downstream, + ); +} diff --git a/packages/zero-cache/src/services/replicator/incremental-sync.ts b/packages/zero-cache/src/services/replicator/incremental-sync.ts index 217c2577ef..8ca54acda0 100644 --- a/packages/zero-cache/src/services/replicator/incremental-sync.ts +++ b/packages/zero-cache/src/services/replicator/incremental-sync.ts @@ -2,11 +2,16 @@ import type {LogContext} from '@rocicorp/logger'; import {AbortError} from '../../../../shared/src/abort-error.ts'; import type {Enum} from '../../../../shared/src/enum.ts'; import {getOrCreateCounter} from '../../observability/metrics.ts'; -import type {Source} from '../../types/streams.ts'; +import { + isStreamBatch, + type Source, + type StreamInPayload, +} from '../../types/streams.ts'; import type {DownloadStatus} from '../change-source/protocol/current.ts'; import type {ChangeStreamData} from '../change-source/protocol/current/downstream.ts'; import { errorTypeToReadableName, + hasBatchedSubscribe, PROTOCOL_VERSION, type ChangeStreamer, type Downstream, @@ -24,6 +29,11 @@ import type {WriteWorkerClient} from './write-worker-client.ts'; type ErrorType = Enum; +// Batch replication messages before crossing into the write worker so +// data-heavy transactions do not pay one IPC round trip per row. The +// rm-vs-load benchmark covers the 1 RM / 16 VS case this is meant to protect; +// the cap keeps unusually large upstream transactions from sitting in the +// syncer heap while the worker is idle. const MAX_WORKER_BATCH_MESSAGES = 64; /** @@ -86,12 +96,12 @@ export class IncrementalSyncer { const {replicaVersion, watermark} = await this.#worker.getSubscriptionState(); - let downstream: Source | undefined; + let downstream: Source> | undefined; let unregister = () => {}; let err: unknown; try { - downstream = await this.#changeStreamer.subscribe({ + const ctx = { protocolVersion: PROTOCOL_VERSION, taskID: this.#taskID, id: this.#id, @@ -99,7 +109,15 @@ export class IncrementalSyncer { watermark, replicaVersion, initial: watermark === initialWatermark, - }); + }; + // Batched subscribe preserves the RM websocket frame boundary. That is + // the unit that actually arrived together on the VS, so using it as the + // worker-flush boundary gives SQLite enough contiguous work to batch + // without changing the replication protocol or commit semantics. + const useBatchedSubscribe = hasBatchedSubscribe(this.#changeStreamer); + downstream = useBatchedSubscribe + ? await this.#changeStreamer.subscribeBatched(ctx) + : await this.#changeStreamer.subscribe(ctx); this.#state.resetBackoff(); unregister = this.#state.cancelOnStop(downstream); this.#statusPublisher?.publish( @@ -112,6 +130,10 @@ export class IncrementalSyncer { const workerBatch = new WorkerMessageBatcher( this.#worker, MAX_WORKER_BATCH_MESSAGES, + // Legacy subscribe() has no outer stream batch to flush at, so commits + // remain the visibility boundary. Batched subscribe() flushes after the + // transport batch, which keeps row-heavy transactions together. + {flushOnCommit: !useBatchedSubscribe}, ); const handleWorkerResult = ( result: CommitResult | readonly CommitResult[] | null, @@ -169,8 +191,7 @@ export class IncrementalSyncer { return workerBatch.push(message)?.then(handleWorkerResult); }; - - for await (const message of downstream) { + const processChangeStreamerMessage = async (message: Downstream) => { this.#replicationEvents.add(1); switch (message[0]) { case 'status': { @@ -210,6 +231,26 @@ export class IncrementalSyncer { break; } } + }; + + for await (const payload of downstream) { + if (isStreamBatch(payload)) { + // This is the production fast path for RM -> VS traffic: apply the + // whole received websocket batch, then ACK/flush once after all of + // its messages have crossed into the write worker. + for (const message of payload.messages) { + await processChangeStreamerMessage(message); + } + await flushWorkerBatch(); + } else { + await processChangeStreamerMessage(payload); + if (useBatchedSubscribe) { + // Batched streams can still emit single messages. Flush them here + // so status/error/small frames do not sit behind a future batch + // that may never arrive. + await flushWorkerBatch(); + } + } } await flushWorkerBatch(); this.#worker.abort(); diff --git a/packages/zero-cache/src/types/streams.test.ts b/packages/zero-cache/src/types/streams.test.ts index bc9b0e45be..0601a29303 100644 --- a/packages/zero-cache/src/types/streams.test.ts +++ b/packages/zero-cache/src/types/streams.test.ts @@ -6,6 +6,7 @@ import Fastify, {type FastifyInstance} from 'fastify'; import {afterEach, beforeEach, describe, expect, test, vi} from 'vitest'; import WebSocket from 'ws'; import {unreachable} from '../../../shared/src/asserts.ts'; +import type {JSONValue} from '../../../shared/src/bigint-json.ts'; import { createSilentLogContext, TestLogSink, @@ -17,10 +18,13 @@ import * as v from '../../../shared/src/valita.ts'; import { stream, streamIn, + streamInBatches, streamOut, streamOutStringified, + type StreamInPayload, type Sink, type Source, + type StringifiedStreamPayload, } from './streams.ts'; import {Subscription, type Result} from './subscription.ts'; @@ -215,10 +219,12 @@ describe('streams with internal acks', () => { let server: FastifyInstance; let producer: Subscription; - let stringifiedProducer: Subscription; + let stringifiedProducer: Subscription; let consumed: Queue; let cleanedUp: Promise; let cleanup: (m: Message[]) => void; + let ackMessages: number[]; + let streamBatchMessages: number; let port: number; let ws: WebSocket; @@ -231,6 +237,8 @@ describe('streams with internal acks', () => { cleanup = resolve; consumed = new Queue(); + ackMessages = []; + streamBatchMessages = 1; producer = Subscription.create({ consumed: m => consumed.enqueue(m), cleanup: resolve, @@ -240,8 +248,23 @@ describe('streams with internal acks', () => { server = Fastify(); await server.register(websocket); server.get('/', {websocket: true}, ws => { - void streamOut(lc, producer, ws); - void streamOutStringified(lc, stringifiedProducer, ws); + ws.on('message', data => { + const text = data.toString(); + if (!text.startsWith('{"ack":')) { + return; + } + const {ack} = JSON.parse(text) as {ack?: unknown}; + if (typeof ack === 'number') { + ackMessages.push(ack); + } + }); + const batchOptions = + streamBatchMessages > 1 + ? {maxMessages: streamBatchMessages} + : undefined; + const streamOptions = {batch: batchOptions}; + void streamOut(lc, producer, ws, streamOptions); + void streamOutStringified(lc, stringifiedProducer, ws, streamOptions); }); // Run the server for real instead of using `injectWS()`, as that has a @@ -268,6 +291,16 @@ describe('streams with internal acks', () => { }; } + async function startBatchReceiver() { + ws = new WebSocket(`http://localhost:${port}/`); + return { + ws, + consumer: (await streamInBatches(lc, ws, messageSchema)) as Subscription< + StreamInPayload + >, + }; + } + test('one at a time', async () => { let num = 0; @@ -348,6 +381,147 @@ describe('streams with internal acks', () => { ]); }); + test('pipelined stream uses per-message acks', async () => { + const messageCount = 96; + const results: Promise[] = []; + for (let i = 0; i < messageCount; i++) { + results.push(producer.push({from: i, to: i + 1, str: `msg-${i}`}).result); + } + + const {consumer} = await startReceiver(); + await vi.waitFor(() => expect(consumer.queued).toBe(messageCount)); + + const entries = await drainPipeline(messageCount, consumer); + for (const {consumed} of entries) { + consumed(); + } + + expect(await Promise.all(results)).toEqual( + Array(messageCount).fill('consumed'), + ); + expect(ackMessages).toEqual( + Array.from({length: messageCount}, (_, i) => i + 1), + ); + consumer.cancel(); + expect(await cleanedUp).toEqual([]); + }); + + test('pipelined outbound stream batches queued messages', async () => { + streamBatchMessages = 4; + const messageCount = 10; + for (let i = 0; i < messageCount; i++) { + producer.push({from: i, to: i + 1, str: `msg-${i}`}); + } + + const {consumer} = await startReceiver(); + await vi.waitFor(() => expect(consumer.queued).toBe(messageCount)); + + const received: Message[] = []; + for await (const msg of consumer) { + received.push(msg); + if (received.length === messageCount) { + break; + } + } + + expect(received).toEqual( + Array.from({length: messageCount}, (_, i) => ({ + from: i, + to: i + 1, + str: `msg-${i}`, + })), + ); + expect(await cleanedUp).toEqual([]); + }); + + test('pipelined outbound stream uses batch frames', async () => { + streamBatchMessages = 4; + const received: JSONValue[] = []; + const messageCount = 10; + + ws = new WebSocket(`http://localhost:${port}/`); + ws.on('message', data => + received.push(JSON.parse(data.toString()) as JSONValue), + ); + + for (let i = 0; i < messageCount; i++) { + producer.push({from: i, to: i + 1, str: `msg-${i}`}); + } + + await vi.waitFor(() => expect(received).toHaveLength(3)); + expect(received).toMatchObject([ + {batch: [{id: 1}, {id: 2}, {id: 3}, {id: 4}]}, + {batch: [{id: 5}, {id: 6}, {id: 7}, {id: 8}]}, + {batch: [{id: 9}, {id: 10}]}, + ]); + ws.close(); + }); + + test('streamInBatches preserves received batch frames', async () => { + streamBatchMessages = 4; + const messageCount = 6; + for (let i = 0; i < messageCount; i++) { + producer.push({from: i, to: i + 1, str: `msg-${i}`}); + } + + const {consumer} = await startBatchReceiver(); + + const received: StreamInPayload[] = []; + for await (const payload of consumer) { + received.push(payload); + if ( + received.reduce( + (sum, payload) => + sum + ('messages' in payload ? payload.messages.length : 1), + 0, + ) === messageCount + ) { + break; + } + } + + expect(received).toEqual([ + { + tag: 'stream-batch', + messages: [ + {from: 0, to: 1, str: 'msg-0'}, + {from: 1, to: 2, str: 'msg-1'}, + {from: 2, to: 3, str: 'msg-2'}, + {from: 3, to: 4, str: 'msg-3'}, + ], + }, + { + tag: 'stream-batch', + messages: [ + {from: 4, to: 5, str: 'msg-4'}, + {from: 5, to: 6, str: 'msg-5'}, + ], + }, + ]); + await vi.waitFor(() => expect(ackMessages).toEqual([1, 2, 3, 4, 5, 6])); + expect(await cleanedUp).toEqual([]); + }); + + test('pipelined unexpected ack closes connection', async () => { + const received: JSONValue[] = []; + const closed = resolver(); + const results: Promise[] = []; + + ws = new WebSocket(`http://localhost:${port}/`); + ws.on('message', data => + received.push(JSON.parse(data.toString()) as JSONValue), + ); + ws.on('close', () => closed.resolve()); + + results.push(producer.push({from: 0, to: 1, str: 'foo'}).result); + results.push(producer.push({from: 1, to: 2, str: 'bar'}).result); + + await vi.waitFor(() => expect(received).toHaveLength(2)); + ws.send(JSON.stringify({ack: 2})); + await closed.promise; + expect(await Promise.all(results)).toEqual(['unconsumed', 'unconsumed']); + }); + test('pipelined (unconsumed)', async () => { const results: Promise[] = []; results.push(producer.push({from: 0, to: 1, str: 'foo'}).result); @@ -462,6 +636,28 @@ describe('streams with internal acks', () => { return drained; } + async function drainPipeline( + num: number, + consumer: Source, + ): Promise<{value: Message; consumed: () => void}[]> { + const {pipeline} = consumer; + expect(pipeline).not.toBeUndefined(); + if (!pipeline) { + unreachable(); + } + const iterator = pipeline[Symbol.asyncIterator](); + const drained: {value: Message; consumed: () => void}[] = []; + for (let i = 0; i < num; i++) { + const result = await iterator.next(); + expect(result.done).not.toBe(true); + if (result.done) { + unreachable(); + } + drained.push(result.value); + } + return drained; + } + test('passthrough', async () => { producer.push({from: 1, to: 2, str: 'foo', extra: 'bar'} as Message); @@ -480,6 +676,24 @@ describe('streams with internal acks', () => { ]); }); + test('stringified source flattens internal message batches', async () => { + const results = [ + stringifiedProducer.push('{"from":1,"to":2,"str":"one"}').result, + stringifiedProducer.push([ + '{"from":2,"to":3,"str":"two"}', + '{"from":3,"to":4,"str":"three"}', + ]).result, + ]; + + const {consumer} = await startReceiver(); + expect(await drain(3, consumer)).toEqual([ + {from: 1, to: 2, str: 'one'}, + {from: 2, to: 3, str: 'two'}, + {from: 3, to: 4, str: 'three'}, + ]); + expect(await Promise.all(results)).toEqual(['consumed', 'consumed']); + }); + test('bigints', async () => { producer.push({ from: 1, diff --git a/packages/zero-cache/src/types/streams.ts b/packages/zero-cache/src/types/streams.ts index a004df2581..053d104731 100644 --- a/packages/zero-cache/src/types/streams.ts +++ b/packages/zero-cache/src/types/streams.ts @@ -16,6 +16,7 @@ import { } from 'ws'; import {assert} from '../../../shared/src/asserts.ts'; import {BigIntJSON, type JSONValue} from '../../../shared/src/bigint-json.ts'; +import {HeadIndexedQueue} from '../../../shared/src/head-indexed-queue.ts'; import {Queue} from '../../../shared/src/queue.ts'; import * as v from '../../../shared/src/valita.ts'; import {Subscription, type Options} from './subscription.ts'; @@ -57,6 +58,24 @@ export type Sink = { push(message: T): void; }; +export type StreamBatch = { + readonly tag: 'stream-batch'; + readonly messages: readonly T[]; +}; + +export type StreamInPayload = T | StreamBatch; + +export function isStreamBatch( + payload: StreamInPayload, +): payload is StreamBatch { + return ( + typeof payload === 'object' && + payload !== null && + !Array.isArray(payload) && + (payload as {readonly tag?: unknown}).tag === 'stream-batch' + ); +} + /** * Back-pressure-aware transformation of a WebSocket into * upstream and downstream {@link Subscription} objects. @@ -155,7 +174,7 @@ type PipeOptions = { export function pipe({source, sink, parse, bufferMessages}: PipeOptions) { bufferMessages ??= 0; assert(bufferMessages >= 0, 'bufferMessages must be non-negative'); - const pending: Promise[] = []; + const pending = new HeadIndexedQueue>(); pipeline( source, @@ -181,12 +200,14 @@ export function pipe({source, sink, parse, bufferMessages}: PipeOptions) { pending.push(result); void result.then(() => pending.shift()); - if (pending.length <= bufferMessages) { + if (pending.size <= bufferMessages) { // immediately allow more messages callback(); } else { // wait for the oldest result in the pending queue - pending[0].then( + const oldest = pending.peek(); + assert(oldest !== undefined, 'pending result expected'); + oldest.then( () => callback(), err => callback(ensureError(err)), ); @@ -212,9 +233,22 @@ function ensureError(err: unknown) { return err instanceof Error ? err : new Error(String(err)); } -const ackSchema = v.object({ack: v.number()}); +type Ack = {ack: number}; -type Ack = v.Infer; +function parseAck(data: unknown): Ack { + if (typeof data !== 'string') { + throw new Error('Expected string message'); + } + const ack = JSON.parse(data) as {ack?: unknown}; + if ( + typeof ack.ack !== 'number' || + !Number.isSafeInteger(ack.ack) || + ack.ack < 0 + ) { + throw new Error(`Invalid ack ${String(ack.ack)}`); + } + return {ack: ack.ack}; +} type Streamed = { /** Application-level message. */ @@ -224,12 +258,22 @@ type Streamed = { id: number; }; +export type StringifiedStreamPayload = string | readonly string[]; + export function streamOut( lc: LogContext, source: Source, sink: WebSocket, + options: StreamOutOptions = {}, ): Promise { - return streamOutInternal(lc, source, sink, BigIntJSON.stringify); + return streamOutInternal( + lc, + source, + sink, + BigIntJSON.stringify, + payload => [payload], + options, + ); } /** @@ -237,29 +281,39 @@ export function streamOut( */ export function streamOutStringified( lc: LogContext, - source: Source, + source: Source, sink: WebSocket, + options: StreamOutOptions = {}, ): Promise { - return streamOutInternal(lc, source, sink, json => json); + return streamOutInternal( + lc, + source, + sink, + json => json, + payload => (typeof payload === 'string' ? [payload] : payload), + options, + ); } -async function streamOutInternal( +type StreamOutOptions = { + batch?: {maxMessages?: number | undefined} | undefined; +}; + +async function streamOutInternal( lc: LogContext, - source: Source, + source: Source, sink: WebSocket, - stringify: (payload: T) => string, + stringify: (payload: TMessage) => string, + expandPayload: (payload: TPayload) => readonly TMessage[], + options: StreamOutOptions, ): Promise { sendPingsForLiveness(lc, sink, PING_INTERVAL_MS); const closer = WebSocketCloser.forSource(lc, sink, source); - const acks = new Queue(); sink.addEventListener('message', ({data}) => { try { - if (typeof data !== 'string') { - throw new Error('Expected string message'); - } - acks.enqueue(v.parse(JSON.parse(data), ackSchema)); + acks.enqueue(parseAck(data)); } catch (e) { lc.error?.(`error parsing ack`, e); closer.close(e); @@ -270,35 +324,91 @@ async function streamOutInternal( let nextID = 0; const {pipeline} = source; if (pipeline) { + const iterator = pipeline[Symbol.asyncIterator](); + const maxBatchMessages = Math.max(1, options.batch?.maxMessages ?? 1); + lc.debug?.(`started pipelined outbound stream`); - for await (const {value: msg, consumed} of pipeline) { - const id = ++nextID; - const data = `{"id":${id},"msg":${stringify(msg)}}`; + for (;;) { + const next = await iterator.next(); + if (next.done) { + break; + } + const batch = [next.value]; + while ( + batch.length < maxBatchMessages && + source instanceof Subscription && + source.queued > 0 + ) { + // Only drain work that is already queued. Waiting for future messages + // just to fill a websocket frame would add latency to quiet streams + // and status/error delivery. + const queued = await iterator.next(); + if (queued.done) { + break; + } + batch.push(queued.value); + } + const outbound = []; + for (const {value, consumed} of batch) { + const messages = expandPayload(value); + if (messages.length === 0) { + consumed(); + continue; + } + let remaining = messages.length; + const consumeOne = () => { + remaining--; + if (remaining === 0) { + // A single source payload can expand to several websocket IDs. + // The source is consumed only after all IDs are ACKed. + consumed(); + } + }; + for (const msg of messages) { + const id = ++nextID; + outbound.push({id, msg}); + void (async () => { + const {ack} = await acks.dequeue(); + // lc.debug?.(`received ack`, ack); + if (ack !== id) { + throw new Error(`Unexpected ack for ${id}: ${ack}`); + } + consumeOne(); + })().catch(e => { + lc.error?.(`error handling ack`, e); + closer.close(e); + }); + } + } + if (outbound.length === 0) { + continue; + } + const data = stringifyOutboundBatch(outbound, stringify); // Enable for debugging. Otherwise too verbose. // lc.debug?.(`pipelining`, data); sink.send(data); - - void (async () => { - const {ack} = await acks.dequeue(); - // lc.debug?.(`received ack`, ack); - if (ack !== id) { - throw new Error(`Unexpected ack for ${id}: ${ack}`); - } - consumed(); - })(); } } else { lc.debug?.(`started synchronous outbound stream`); - for await (const msg of source) { - const id = ++nextID; - const data = `{"id":${id},"msg":${stringify(msg)}}`; + for await (const payload of source) { + const messages = expandPayload(payload); + if (messages.length === 0) { + continue; + } + const outbound = messages.map(msg => ({ + id: ++nextID, + msg, + })); + const data = stringifyOutboundBatch(outbound, stringify); // Enable for debugging. Otherwise too verbose. // lc.debug?.(`sending`, data); sink.send(data); - const {ack} = await acks.dequeue(); - if (ack !== id) { - throw new Error(`Unexpected ack for ${id}: ${ack}`); + for (const {id} of outbound) { + const {ack} = await acks.dequeue(); + if (ack !== id) { + throw new Error(`Unexpected ack for ${id}: ${ack}`); + } } } } @@ -308,11 +418,50 @@ async function streamOutInternal( } } -export async function streamIn( +function stringifyOutboundBatch( + batch: readonly {readonly id: number; readonly msg: T}[], + stringify: (payload: T) => string, +) { + if (batch.length === 1) { + const [{id, msg}] = batch; + return `{"id":${id},"msg":${stringify(msg)}}`; + } + // RM -> VS catchup emits large websocket batches. Building the frame in one + // loop avoids the extra strings and array that `map().join()` creates in the + // hottest fanout path. + let data = '{"batch":['; + for (let i = 0; i < batch.length; i++) { + if (i > 0) { + data += ','; + } + const {id, msg} = batch[i]; + data += `{"id":${id},"msg":${stringify(msg)}}`; + } + return `${data}]}`; +} + +export function streamIn( lc: LogContext, source: WebSocket, schema: v.Type, ): Promise> { + return streamInInternal(lc, source, schema, false) as Promise>; +} + +export function streamInBatches( + lc: LogContext, + source: WebSocket, + schema: v.Type, +): Promise>> { + return streamInInternal(lc, source, schema, true); +} + +async function streamInInternal( + lc: LogContext, + source: WebSocket, + schema: v.Type, + preserveBatches: boolean, +): Promise>> { expectPingsForLiveness(lc, source, PING_INTERVAL_MS); const streamedSchema = v.object({ @@ -320,28 +469,53 @@ export async function streamIn( id: v.number(), }); - const sink: Subscription> = new Subscription>( + const sink: Subscription< + StreamInPayload, + Streamed | readonly Streamed[] + > = new Subscription( { - consumed: ({id}) => source.send(JSON.stringify({ack: id} satisfies Ack)), + consumed: streamOrBatch => { + if (isStreamedBatch(streamOrBatch)) { + for (const {id} of streamOrBatch) { + source.send(JSON.stringify({ack: id} satisfies Ack)); + } + } else { + source.send(JSON.stringify({ack: streamOrBatch.id} satisfies Ack)); + } + }, cleanup: () => closer.close(), }, - ({msg}) => msg, + streamOrBatch => + isStreamedBatch(streamOrBatch) + ? { + tag: 'stream-batch', + messages: streamOrBatch.map(({msg}) => msg), + } + : streamOrBatch.msg, ); const closer = WebSocketCloser.forSink(lc, source, sink, handleMessage); function handleMessage(event: MessageEvent) { - const data = event.data.toString(); + const data = webSocketMessageDataToString(event.data); if (!sink.active) { lc.warn?.('dropping ws message received after close', data); return; } try { const value = BigIntJSON.parse(data); - const msg = v.parse(value, streamedSchema, 'passthrough'); - // Enable for debugging. Otherwise too verbose. - // lc.debug?.(`received`, data); - sink.push(msg); + const messages = parseStreamedMessages(value, streamedSchema); + if (preserveBatches && messages.length > 1) { + // The RM already grouped these messages into one websocket frame. Keep + // that boundary visible so the VS can batch one write-worker apply. + sink.push(messages); + return; + } + for (const msg of messages) { + // Enable for debugging. Otherwise too verbose. + // lc.debug?.(`received`, data); + sink.push(msg); + } } catch (e) { closer.close(e); } @@ -351,6 +525,51 @@ export async function streamIn( return sink; } +function webSocketMessageDataToString(data: MessageEvent['data']): string { + if (typeof data === 'string') { + return data; + } + if (Buffer.isBuffer(data)) { + return data.toString(); + } + if (data instanceof ArrayBuffer) { + return Buffer.from(data).toString(); + } + return Buffer.concat(data).toString(); +} + +function isStreamedBatch( + streamOrBatch: Streamed | readonly Streamed[], +): streamOrBatch is readonly Streamed[] { + return Array.isArray(streamOrBatch); +} + +function parseStreamedMessages( + value: JSONValue, + streamedSchema: v.Type>, +): Streamed[] { + if ( + typeof value === 'object' && + value !== null && + !Array.isArray(value) && + Object.hasOwn(value, 'batch') + ) { + // A batch frame is a transport optimization only. Validate each child + // message against the normal streamed schema before handing it to callers. + const batch = (value as {readonly batch?: unknown}).batch; + if (!Array.isArray(batch)) { + throw new Error('Invalid stream batch'); + } + const messages: Streamed[] = []; + messages.length = batch.length; + for (let i = 0; i < batch.length; i++) { + messages[i] = v.parse(batch[i], streamedSchema, 'passthrough'); + } + return messages; + } + return [v.parse(value, streamedSchema, 'passthrough')]; +} + class WebSocketCloser { readonly #lc: LogContext; readonly #ws: WebSocket; @@ -368,10 +587,10 @@ class WebSocketCloser { return new WebSocketCloser(lc, ws, () => stream.cancel()); } - static forSink( + static forSink( lc: LogContext, ws: WebSocket, - stream: Subscription>, + stream: Subscription, messageHandler: (e: MessageEvent) => void | undefined, ) { // If the websocket is closed, call end() to allow the downstream Sink