Refactor replication connections.

Author: rkistner

modules/module-postgres/src/replication/WalStream.ts

@@ -374,7 +374,7 @@ WHERE  oid = $1::regclass`,
    * If (partial) replication was done before on this slot, this clears the state
    * and starts again from scratch.
    */
-  async startInitialReplication(replicationConnection: pgwire.PgConnection, status: InitResult) {
+  async startInitialReplication(status: InitResult) {
     // If anything here errors, the entire replication process is aborted,
     // and all connections are closed, including this one.
     const db = await this.connections.snapshotConnection();
@@ -396,7 +396,12 @@ WHERE  oid = $1::regclass`,
 
       // We use the replication connection here, not a pool.
       // The replication slot must be created before we start snapshotting tables.
-      await replicationConnection.query(`CREATE_REPLICATION_SLOT ${slotName} LOGICAL pgoutput`);
+      const initReplicationConnection = await this.connections.replicationConnection();
+      try {
+        await initReplicationConnection.query(`CREATE_REPLICATION_SLOT ${slotName} LOGICAL pgoutput`);
+      } finally {
+        await initReplicationConnection.end();
+      }
 
       this.logger.info(`Created replication slot ${slotName}`);
     }
@@ -815,37 +820,34 @@ WHERE  oid = $1::regclass`,
 
   async replicate() {
     try {
-      // If anything errors here, the entire replication process is halted, and
-      // all connections automatically closed, including this one.
-      const initReplicationConnection = await this.connections.replicationConnection();
-      await this.initReplication(initReplicationConnection);
-      await initReplicationConnection.end();
+      await this.initReplication();
 
-      // At this point, the above connection has often timed out, so we start a new one
-      const streamReplicationConnection = await this.connections.replicationConnection();
-      await this.streamChanges(streamReplicationConnection);
-      await streamReplicationConnection.end();
+      // At this point, the above connection has often timed out, so we start a new one in streamChanges().
+      await this.streamChanges();
     } catch (e) {
       await this.storage.reportError(e);
       throw e;
     }
   }
 
-  async initReplication(replicationConnection: pgwire.PgConnection) {
+  async initReplication() {
     const result = await this.initSlot();
     if (result.needsInitialSync) {
-      await this.startInitialReplication(replicationConnection, result);
+      await this.startInitialReplication(result);
     }
   }
 
-  async streamChanges(replicationConnection: pgwire.PgConnection) {
+  async streamChanges() {
+    const streamReplicationConnection = await this.connections.replicationConnection();
     try {
-      await this.streamChangesInternal(replicationConnection);
+      await this.streamChangesInternal(streamReplicationConnection);
     } catch (e) {
       if (isReplicationSlotInvalidError(e)) {
         throw new MissingReplicationSlotError(e.message, e);
       }
       throw e;
+    } finally {
+      await streamReplicationConnection.end();
     }
   }
 

modules/module-postgres/test/src/slow_tests.test.ts

@@ -70,7 +70,6 @@ function defineSlowTests(factory: storage.TestStorageFactory) {
 
   async function testRepeatedReplication(testOptions: { compact: boolean; maxBatchSize: number; numBatches: number }) {
     const connections = new PgManager(TEST_CONNECTION_OPTIONS, {});
-    const replicationConnection = await connections.replicationConnection();
     const pool = connections.pool;
     await clearTestDb(pool);
     await using f = await factory();
@@ -97,9 +96,9 @@ bucket_definitions:
     );
     await pool.query(`ALTER TABLE test_data REPLICA IDENTITY FULL`);
 
-    await walStream.initReplication(replicationConnection);
+    await walStream.initReplication();
     let abort = false;
-    streamPromise = walStream.streamChanges(replicationConnection).finally(() => {
+    streamPromise = walStream.streamChanges().finally(() => {
       abort = true;
     });
     const start = Date.now();
@@ -331,7 +330,6 @@ bucket_definitions:
       i += 1;
 
       const connections = new PgManager(TEST_CONNECTION_OPTIONS, {});
-      const replicationConnection = await connections.replicationConnection();
 
       abortController = new AbortController();
       const options: WalStreamOptions = {
@@ -347,9 +345,9 @@ bucket_definitions:
       // 3. Start initial replication, then streaming, but don't wait for any of this
       let initialReplicationDone = false;
       streamPromise = (async () => {
-        await walStream.initReplication(replicationConnection);
+        await walStream.initReplication();
         initialReplicationDone = true;
-        await walStream.streamChanges(replicationConnection);
+        await walStream.streamChanges();
       })()
         .catch((e) => {
           initialReplicationDone = true;

modules/module-postgres/test/src/wal_stream_utils.ts

@@ -19,7 +19,6 @@ export class WalStreamTestContext implements AsyncDisposable {
   private abortController = new AbortController();
   private streamPromise?: Promise<void>;
   public storage?: SyncRulesBucketStorage;
-  private replicationConnection?: pgwire.PgConnection;
   private snapshotPromise?: Promise<void>;
 
   /**
@@ -150,18 +149,14 @@ export class WalStreamTestContext implements AsyncDisposable {
 
   async replicateSnapshot() {
     const promise = (async () => {
-      this.replicationConnection = await this.connectionManager.replicationConnection();
-      await this.walStream.initReplication(this.replicationConnection);
+      await this.walStream.initReplication();
     })();
     this.snapshotPromise = promise.catch((e) => e);
     await promise;
   }
 
   startStreaming() {
-    if (this.replicationConnection == null) {
-      throw new Error('Call replicateSnapshot() before startStreaming()');
-    }
-    this.streamPromise = this.walStream.streamChanges(this.replicationConnection!);
+    this.streamPromise = this.walStream.streamChanges();
   }
 
   async getCheckpoint(options?: { timeout?: number }) {