Refactor streaming promise management.

rkistner · rkistner · commit f09dea31ca6a · 2025-12-01T11:27:13.000+02:00
diff --git a/modules/module-postgres/src/replication/PostgresSnapshotter.ts b/modules/module-postgres/src/replication/PostgresSnapshotter.ts
@@ -344,23 +344,28 @@ export class PostgresSnapshotter {
   }
 
   async replicationLoop() {
-    if (this.queue.size == 0) {
-      // Special case where we start with no tables to snapshot
-      await this.markSnapshotDone();
-    }
-    while (!this.abort_signal.aborted) {
-      const table = this.queue.values().next().value;
-      if (table == null) {
-        this.initialSnapshotDone.resolve();
-        await timers.setTimeout(500, { signal: this.abort_signal });
-        continue;
-      }
-
-      await this.replicateTable(table);
-      this.queue.delete(table);
+    try {
       if (this.queue.size == 0) {
+        // Special case where we start with no tables to snapshot
         await this.markSnapshotDone();
       }
+      while (!this.abort_signal.aborted) {
+        const table = this.queue.values().next().value;
+        if (table == null) {
+          this.initialSnapshotDone.resolve();
+          await timers.setTimeout(500, { signal: this.abort_signal });
+          continue;
+        }
+
+        await this.replicateTable(table);
+        this.queue.delete(table);
+        if (this.queue.size == 0) {
+          await this.markSnapshotDone();
+        }
+      }
+    } catch (e) {
+      this.initialSnapshotDone.reject(e);
+      throw e;
     }
   }
 
@@ -410,7 +415,7 @@ export class PostgresSnapshotter {
    * If (partial) replication was done before on this slot, this clears the state
    * and starts again from scratch.
    */
-  async startReplication(db: pgwire.PgConnection) {
+  async queueSnapshotTables(db: pgwire.PgConnection) {
     const sourceTables = this.sync_rules.getSourceTables();
 
     await this.storage.startBatch(
diff --git a/modules/module-postgres/src/replication/WalStream.ts b/modules/module-postgres/src/replication/WalStream.ts
@@ -101,9 +101,10 @@ export class WalStream {
 
   private connections: PgManager;
 
-  private abort_signal: AbortSignal;
+  private abortController = new AbortController();
+  private abortSignal: AbortSignal = this.abortController.signal;
 
-  private streamPromise: Promise<void> | null = null;
+  private initPromise: Promise<void> | null = null;
   private snapshotter: PostgresSnapshotter;
 
   private relationCache = new RelationCache((relation: number | SourceTable) => {
@@ -135,9 +136,16 @@ export class WalStream {
     this.slot_name = options.storage.slot_name;
     this.connections = options.connections;
 
-    this.abort_signal = options.abort_signal;
-    this.snapshotter = new PostgresSnapshotter(options);
-    this.abort_signal.addEventListener(
+    // We wrap in our own abort controller so we can trigger abort internally.
+    options.abort_signal.addEventListener('abort', () => {
+      this.abortController.abort();
+    });
+    if (options.abort_signal.aborted) {
+      this.abortController.abort();
+    }
+
+    this.snapshotter = new PostgresSnapshotter({ ...options, abort_signal: this.abortSignal });
+    this.abortSignal.addEventListener(
       'abort',
       () => {
         if (this.startedStreaming) {
@@ -159,7 +167,7 @@ export class WalStream {
   }
 
   get stopped() {
-    return this.abort_signal.aborted;
+    return this.abortSignal.aborted;
   }
 
   async handleRelation(options: {
@@ -328,52 +336,62 @@ export class WalStream {
     return null;
   }
 
+  /**
+   * Start replication loop, and continue until aborted or error.
+   */
   async replicate() {
     try {
-      await this.initReplication();
-
-      await this.streamChanges();
+      this.initPromise = this.initReplication();
+      await this.initPromise;
+      const streamPromise = this.streamChanges();
+      const loopPromise = this.snapshotter.replicationLoop();
+      await Promise.race([loopPromise, streamPromise]);
     } catch (e) {
       await this.storage.reportError(e);
       throw e;
+    } finally {
+      this.abortController.abort();
+    }
+  }
+
+  public async waitForInitialSnapshot() {
+    if (this.initPromise == null) {
+      throw new ReplicationAssertionError('replicate() must be called before waitForInitialSnapshot()');
     }
+    await this.initPromise;
+
+    await this.snapshotter.waitForInitialSnapshot();
   }
 
-  async initReplication() {
+  /**
+   * Initialize replication.
+   * Start replication loop, and continue until aborted, error or initial snapshot completed.
+   */
+  private async initReplication() {
     const result = await this.snapshotter.checkSlot();
     const db = await this.connections.snapshotConnection();
     try {
       await this.snapshotter.setupSlot(db, result);
-      // Trigger here, but we await elsewhere
-      // TODO: fail on the first error
-      this.streamChanges().catch((_) => {});
       if (result.needsInitialSync) {
-        await this.snapshotter.startReplication(db);
+        await this.snapshotter.queueSnapshotTables(db);
       }
-
-      // FIXME: handle the background promise
-      this.snapshotter.replicationLoop();
-      await this.snapshotter.waitForInitialSnapshot();
     } finally {
       await db.end();
     }
   }
 
-  async streamChanges() {
-    this.streamPromise ??= (async () => {
-      const streamReplicationConnection = await this.connections.replicationConnection();
-      try {
-        await this.streamChangesInternal(streamReplicationConnection);
-      } catch (e) {
-        if (isReplicationSlotInvalidError(e)) {
-          throw new MissingReplicationSlotError(e.message, e);
-        }
-        throw e;
-      } finally {
-        await streamReplicationConnection.end();
+  private async streamChanges() {
+    const streamReplicationConnection = await this.connections.replicationConnection();
+    try {
+      await this.streamChangesInternal(streamReplicationConnection);
+    } catch (e) {
+      if (isReplicationSlotInvalidError(e)) {
+        throw new MissingReplicationSlotError(e.message, e);
       }
-    })();
-    await this.streamPromise;
+      throw e;
+    } finally {
+      await streamReplicationConnection.end();
+    }
   }
 
   private async streamChangesInternal(replicationConnection: pgwire.PgConnection) {
@@ -451,7 +469,7 @@ export class WalStream {
         for await (const chunk of replicationStream.pgoutputDecode()) {
           this.touch();
 
-          if (this.abort_signal.aborted) {
+          if (this.abortSignal.aborted) {
             break;
           }
 
diff --git a/modules/module-postgres/test/src/checkpoints.test.ts b/modules/module-postgres/test/src/checkpoints.test.ts
@@ -35,11 +35,9 @@ const checkpointTests = (factory: TestStorageFactory) => {
 
     await pool.query(`CREATE TABLE test_data(id text primary key, description text, other text)`);
 
-    await context.replicateSnapshot();
-
-    context.startStreaming();
     // Wait for a consistent checkpoint before we start.
-    await context.getCheckpoint();
+    await context.initializeReplication();
+
     const storage = context.storage!;
 
     const controller = new AbortController();
diff --git a/modules/module-postgres/test/src/chunked_snapshots.test.ts b/modules/module-postgres/test/src/chunked_snapshots.test.ts
@@ -142,7 +142,8 @@ function defineBatchTests(factory: TestStorageFactory) {
     await p;
 
     // 5. Logical replication picks up the UPDATE above, but it is missing the TOAST column.
-    context.startStreaming();
+    // Note: logical replication now runs concurrently with the snapshot.
+    // TODO: re-check the test logic here.
 
     // 6. If all went well, the "resnapshot" process would take care of this.
     const data = await context.getBucketData('global[]', undefined, {});
diff --git a/modules/module-postgres/test/src/large_batch.test.ts b/modules/module-postgres/test/src/large_batch.test.ts
@@ -39,8 +39,6 @@ function defineBatchTests(factory: storage.TestStorageFactory) {
 
     const start = Date.now();
 
-    context.startStreaming();
-
     const checkpoint = await context.getCheckpoint({ timeout: 100_000 });
     const duration = Date.now() - start;
     const used = Math.round(process.memoryUsage().heapUsed / 1024 / 1024);
@@ -87,7 +85,6 @@ function defineBatchTests(factory: storage.TestStorageFactory) {
       const start = Date.now();
 
       await context.replicateSnapshot();
-      context.startStreaming();
 
       const checkpoint = await context.getCheckpoint({ timeout: 100_000 });
       const duration = Date.now() - start;
@@ -139,8 +136,6 @@ function defineBatchTests(factory: storage.TestStorageFactory) {
 
     const start = Date.now();
 
-    context.startStreaming();
-
     const checkpoint = await context.getCheckpoint({ timeout: 50_000 });
     const duration = Date.now() - start;
     const used = Math.round(process.memoryUsage().heapUsed / 1024 / 1024);
@@ -226,8 +221,6 @@ function defineBatchTests(factory: storage.TestStorageFactory) {
     });
     await context.replicateSnapshot();
 
-    context.startStreaming();
-
     const checkpoint = await context.getCheckpoint({ timeout: 50_000 });
     const checksum = await context.storage!.getChecksums(checkpoint, ['global[]']);
     expect(checksum.get('global[]')!.count).toEqual((numDocs + 2) * 4);
diff --git a/modules/module-postgres/test/src/resuming_snapshots.test.ts b/modules/module-postgres/test/src/resuming_snapshots.test.ts
@@ -104,7 +104,6 @@ async function testResumingReplication(factory: TestStorageFactory, stopAfter: n
   await context2.loadNextSyncRules();
   await context2.replicateSnapshot();
 
-  context2.startStreaming();
   const data = await context2.getBucketData('global[]', undefined, {});
 
   const deletedRowOps = data.filter((row) => row.object_type == 'test_data2' && row.object_id === String(id1));
diff --git a/modules/module-postgres/test/src/schema_changes.test.ts b/modules/module-postgres/test/src/schema_changes.test.ts
@@ -37,7 +37,6 @@ function defineTests(factory: storage.TestStorageFactory) {
     await pool.query(`INSERT INTO test_data(id, description) VALUES('t1', 'test1')`);
 
     await context.replicateSnapshot();
-    context.startStreaming();
 
     await pool.query(`INSERT INTO test_data(id, description) VALUES('t2', 'test2')`);
 
@@ -68,7 +67,6 @@ function defineTests(factory: storage.TestStorageFactory) {
     const { pool } = context;
 
     await context.replicateSnapshot();
-    context.startStreaming();
 
     await pool.query(`CREATE TABLE test_data(id text primary key, description text)`);
     await pool.query(`INSERT INTO test_data(id, description) VALUES('t1', 'test1')`);
@@ -103,7 +101,6 @@ function defineTests(factory: storage.TestStorageFactory) {
     await pool.query(`INSERT INTO test_data_old(id, description) VALUES('t1', 'test1')`);
 
     await context.replicateSnapshot();
-    context.startStreaming();
 
     await pool.query(
       { statement: `ALTER TABLE test_data_old RENAME TO test_data` },
@@ -143,7 +140,6 @@ function defineTests(factory: storage.TestStorageFactory) {
     await pool.query(`INSERT INTO test_data1(id, description) VALUES('t1', 'test1')`);
 
     await context.replicateSnapshot();
-    context.startStreaming();
 
     await pool.query(
       { statement: `ALTER TABLE test_data1 RENAME TO test_data2` },
@@ -186,7 +182,6 @@ function defineTests(factory: storage.TestStorageFactory) {
     await pool.query(`INSERT INTO test_data(id, description) VALUES('t1', 'test1')`);
 
     await context.replicateSnapshot();
-    context.startStreaming();
 
     await pool.query(
       { statement: `ALTER TABLE test_data RENAME TO test_data_na` },
@@ -216,7 +211,6 @@ function defineTests(factory: storage.TestStorageFactory) {
     await pool.query(`INSERT INTO test_data(id, description) VALUES('t1', 'test1')`);
 
     await context.replicateSnapshot();
-    context.startStreaming();
 
     await pool.query(
       { statement: `ALTER TABLE test_data REPLICA IDENTITY FULL` },
@@ -259,7 +253,6 @@ function defineTests(factory: storage.TestStorageFactory) {
     await pool.query(`INSERT INTO test_data(id, description) VALUES('t1', 'test1')`);
 
     await context.replicateSnapshot();
-    context.startStreaming();
 
     await pool.query(
       { statement: `ALTER TABLE test_data ADD COLUMN other TEXT` },
@@ -301,7 +294,6 @@ function defineTests(factory: storage.TestStorageFactory) {
     await pool.query(`INSERT INTO test_data(id, description) VALUES('t1', 'test1')`);
 
     await context.replicateSnapshot();
-    context.startStreaming();
 
     await pool.query(
       { statement: `ALTER TABLE test_data ALTER COLUMN id TYPE varchar` },
@@ -345,7 +337,6 @@ function defineTests(factory: storage.TestStorageFactory) {
     await pool.query(`INSERT INTO test_data(id, description) VALUES('t1', 'test1')`);
 
     await context.replicateSnapshot();
-    context.startStreaming();
 
     await pool.query(`INSERT INTO test_data(id, description) VALUES('t2', 'test2')`);
 
@@ -396,7 +387,6 @@ function defineTests(factory: storage.TestStorageFactory) {
     await pool.query(`INSERT INTO test_data(id, description) VALUES('t1', 'test1')`);
 
     await context.replicateSnapshot();
-    context.startStreaming();
 
     await pool.query(`INSERT INTO test_data(id, description) VALUES('t2', 'test2')`);
 
@@ -440,7 +430,6 @@ function defineTests(factory: storage.TestStorageFactory) {
     await pool.query(`INSERT INTO test_other(id, description) VALUES('t1', 'test1')`);
 
     await context.replicateSnapshot();
-    context.startStreaming();
 
     await pool.query(`INSERT INTO test_other(id, description) VALUES('t2', 'test2')`);
 
@@ -468,7 +457,6 @@ function defineTests(factory: storage.TestStorageFactory) {
     await pool.query(`INSERT INTO test_data(id, description) VALUES('t1', 'test1')`);
 
     await context.replicateSnapshot();
-    context.startStreaming();
 
     await pool.query(`INSERT INTO test_data(id, description) VALUES('t2', 'test2')`);
 
@@ -506,7 +494,6 @@ function defineTests(factory: storage.TestStorageFactory) {
     await pool.query(`INSERT INTO test_data(id, description) VALUES('t1', 'test1')`);
 
     await context.replicateSnapshot();
-    context.startStreaming();
 
     await pool.query(`INSERT INTO test_data(id, description) VALUES('t2', 'test2')`);
 
@@ -555,7 +542,6 @@ function defineTests(factory: storage.TestStorageFactory) {
     await pool.query(`INSERT INTO test_data_old(id, num) VALUES('t2', 0)`);
 
     await context.replicateSnapshot();
-    context.startStreaming();
 
     await pool.query(
       { statement: `ALTER TABLE test_data_old RENAME TO test_data` },
@@ -627,7 +613,6 @@ config:
     await pool.query(`INSERT INTO test_data(id) VALUES ('t1')`);
 
     await context.replicateSnapshot();
-    context.startStreaming();
 
     await pool.query(
       { statement: `CREATE TYPE composite AS (foo bool, bar int4);` },
diff --git a/modules/module-postgres/test/src/wal_stream.test.ts b/modules/module-postgres/test/src/wal_stream.test.ts
diff --git a/modules/module-postgres/test/src/wal_stream_utils.ts b/modules/module-postgres/test/src/wal_stream_utils.ts
diff --git a/packages/service-core/src/sync/util.ts b/packages/service-core/src/sync/util.ts
