perf(zero-cache): batch serving replica apply work

packages/zero-cache/src/services/replicator/change-processor.bench.ts

@@ -0,0 +1,136 @@
+import {bench, describe, use} from '../../../../shared/src/bench.ts';
+import {createSilentLogContext} from '../../../../shared/src/logging-test-utils.ts';
+import {Database} from '../../../../zqlite/src/db.ts';
+import {StatementRunner} from '../../db/statements.ts';
+import type {DataOrSchemaChange} from '../change-source/protocol/current/data.ts';
+import type {ChangeStreamData} from '../change-source/protocol/current/downstream.ts';
+import {ChangeProcessor} from './change-processor.ts';
+import {initReplicationState} from './schema/replication-state.ts';
+import {ReplicationMessages} from './test-utils.ts';
+
+const ROWS = 1_000;
+const lc = createSilentLogContext();
+const messages = new ReplicationMessages({foo: ['id', 'shard']});
+
+describe('change-processor DML apply', () => {
+  const stableInsertChanges = transaction(
+    '10',
+    Array.from({length: ROWS}, (_, i) =>
+      messages.insert('foo', {
+        id: i,
+        shard: i % 10,
+        alpha: i * 2,
+        beta: i * 3,
+      }),
+    ),
+  );
+  const stableInsertUncached = setupBenchmark(false);
+  const stableInsertCached = setupBenchmark(true);
+
+  bench(
+    '1k stable-shape inserts (uncached plans)',
+    () => {
+      use(apply(stableInsertUncached, stableInsertChanges));
+    },
+    {min_cpu_time: 5e8, min_samples: 20},
+  );
+
+  bench(
+    '1k stable-shape inserts (cached plans)',
+    () => {
+      use(apply(stableInsertCached, stableInsertChanges));
+    },
+    {min_cpu_time: 5e8, min_samples: 20},
+  );
+
+  const mixedDmlChanges = transaction('20', [
+    ...Array.from({length: ROWS}, (_, i) =>
+      messages.insert('foo', {
+        id: i,
+        shard: i % 10,
+        alpha: i * 2,
+        beta: i * 3,
+      }),
+    ),
+    ...Array.from({length: ROWS}, (_, i) =>
+      i % 2 === 0
+        ? messages.update('foo', {
+            id: i,
+            shard: i % 10,
+            alpha: i * 5,
+            beta: i * 7,
+          })
+        : messages.update('foo', {
+            beta: i * 7,
+            shard: i % 10,
+            id: i,
+            alpha: i * 5,
+          }),
+    ),
+    ...Array.from({length: ROWS}, (_, i) =>
+      i % 2 === 0
+        ? messages.delete('foo', {id: i, shard: i % 10})
+        : messages.delete('foo', {shard: i % 10, id: i}),
+    ),
+  ]);
+  const mixedDmlUncached = setupBenchmark(false);
+  const mixedDmlCached = setupBenchmark(true);
+
+  bench(
+    '1k inserts + updates + deletes (uncached plans)',
+    () => {
+      use(apply(mixedDmlUncached, mixedDmlChanges));
+    },
+    {min_cpu_time: 5e8, min_samples: 20},
+  );
+
+  bench(
+    '1k inserts + updates + deletes (cached plans)',
+    () => {
+      use(apply(mixedDmlCached, mixedDmlChanges));
+    },
+    {min_cpu_time: 5e8, min_samples: 20},
+  );
+});
+
+function setupBenchmark(cacheDmlSqlPlans: boolean) {
+  const replica = new Database(lc, ':memory:');
+  initReplicationState(replica, ['zero_data'], '02');
+  replica.exec(`
+    CREATE TABLE foo(
+      id INTEGER,
+      shard INTEGER,
+      alpha INTEGER,
+      beta INTEGER,
+      _0_version TEXT,
+      PRIMARY KEY(id, shard)
+    );
+  `);
+  return new ChangeProcessor(
+    new StatementRunner(replica),
+    'backup',
+    (_, err) => {
+      throw err;
+    },
+    {cacheDmlSqlPlans},
+  );
+}
+
+function transaction(
+  watermark: string,
+  dataMessages: DataOrSchemaChange[],
+): ChangeStreamData[] {
+  return [
+    ['begin', messages.begin(), {commitWatermark: watermark}],
+    ...dataMessages.map(msg => ['data', msg] as ChangeStreamData),
+    ['commit', messages.commit(), {watermark}],
+  ];
+}
+
+function apply(processor: ChangeProcessor, changes: ChangeStreamData[]) {
+  let result = null;
+  for (const change of changes) {
+    result = processor.processMessage(lc, change);
+  }
+  return result;
+}

packages/zero-cache/src/services/replicator/change-processor.test.ts

@@ -3639,6 +3639,98 @@ describe('replicator/change-processor', () => {
       }
     });
   }
+
+  test('DML plan cache binds repeated shapes by column name', () => {
+    initDB(
+      servingReplica,
+      `
+      CREATE TABLE foo(
+        id INTEGER,
+        shard INTEGER,
+        alpha INTEGER,
+        beta INTEGER,
+        _0_version TEXT,
+        PRIMARY KEY(id, shard)
+      );
+      `,
+    );
+
+    const foo = new ReplicationMessages({foo: ['id', 'shard']});
+    for (const change of [
+      ['begin', foo.begin(), {commitWatermark: '07'}],
+      ['data', foo.insert('foo', {id: 1, shard: 10, alpha: 100, beta: 200})],
+      ['data', foo.insert('foo', {beta: 400, shard: 20, id: 2, alpha: 300})],
+      ['data', foo.update('foo', {id: 1, shard: 10, alpha: 101, beta: 201})],
+      ['data', foo.update('foo', {beta: 401, shard: 20, id: 2, alpha: 301})],
+      ['data', foo.delete('foo', {shard: 20, id: 2})],
+      ['commit', foo.commit(), {watermark: '07'}],
+    ] satisfies ChangeStreamData[]) {
+      servingProcessor.processMessage(lc, change);
+    }
+
+    expectTables(
+      servingReplica,
+      {
+        foo: [
+          {
+            id: 1n,
+            shard: 10n,
+            alpha: 101n,
+            beta: 201n,
+            ['_0_version']: '07',
+          },
+        ],
+      },
+      'bigint',
+    );
+  });
+
+  test('processMessages applies a stream batch and returns each commit', () => {
+    initDB(
+      servingReplica,
+      `
+      CREATE TABLE foo(
+        id INTEGER PRIMARY KEY,
+        _0_version TEXT
+      );
+      `,
+    );
+
+    const foo = new ReplicationMessages({foo: 'id'});
+    const result = servingProcessor.processMessages(lc, [
+      ['begin', foo.begin(), {commitWatermark: '07'}],
+      ['data', foo.insert('foo', {id: 1})],
+      ['commit', foo.commit(), {watermark: '07'}],
+      ['begin', foo.begin(), {commitWatermark: '0a'}],
+      ['data', foo.insert('foo', {id: 2})],
+      ['commit', foo.commit(), {watermark: '0a'}],
+    ]);
+
+    expect(result).toEqual([
+      {
+        watermark: '07',
+        completedBackfill: undefined,
+        schemaUpdated: false,
+        changeLogUpdated: true,
+      },
+      {
+        watermark: '0a',
+        completedBackfill: undefined,
+        schemaUpdated: false,
+        changeLogUpdated: true,
+      },
+    ]);
+    expectTables(
+      servingReplica,
+      {
+        foo: [
+          {id: 1n, ['_0_version']: '07'},
+          {id: 2n, ['_0_version']: '0a'},
+        ],
+      },
+      'bigint',
+    );
+  });
 });
 
 describe('replicator/change-processor-errors', () => {
@@ -3790,6 +3882,11 @@ describe('replicator/change-processor-errors', () => {
       'data',
       messages.insert('auto_rollback', {id: 123}),
     ]);
+    processor.processMessage(lc, [
+      'commit',
+      messages.commit(),
+      {watermark: '0e'},
+    ]);
 
     expect(failures).toHaveLength(1);
     expect(failures[0]).toBeInstanceOf(Error);