Refactor commit logic to better handle concurrent replication.

rkistner · rkistner · commit ce500a10f98c · 2025-11-20T11:23:42.000+02:00
diff --git a/modules/module-mongodb-storage/src/storage/implementation/MongoBucketBatch.ts b/modules/module-mongodb-storage/src/storage/implementation/MongoBucketBatch.ts
@@ -661,32 +661,7 @@ export class MongoBucketBatch
 
     await this.flush(options);
 
-    if (this.last_checkpoint_lsn != null && lsn < this.last_checkpoint_lsn) {
-      // When re-applying transactions, don't create a new checkpoint until
-      // we are past the last transaction.
-      this.logger.info(`Re-applied transaction ${lsn} - skipping checkpoint`);
-      // Cannot create a checkpoint yet - return false
-      return false;
-    }
-
-    if (!createEmptyCheckpoints && this.persisted_op == null) {
-      // Nothing to commit - also return true
-      await this.autoActivate(lsn);
-      return true;
-    }
-
     const now = new Date();
-    const update: Partial<SyncRuleDocument> = {
-      last_checkpoint_lsn: lsn,
-      last_checkpoint_ts: now,
-      last_keepalive_ts: now,
-      last_fatal_error: null,
-      keepalive_op: null
-    };
-
-    if (this.persisted_op != null) {
-      update.last_checkpoint = this.persisted_op;
-    }
 
     // Mark relevant write checkpoints as "processed".
     // This makes it easier to identify write checkpoints that are "valid" in order.
@@ -705,58 +680,101 @@ export class MongoBucketBatch
       }
     );
 
-    const updateResult = await this.db.sync_rules.updateOne(
+    const updateResult = await this.db.sync_rules.findOneAndUpdate(
       {
-        _id: this.group_id,
-        snapshot_done: true,
-        $or: [{ no_checkpoint_before: null }, { no_checkpoint_before: { $lte: lsn } }]
+        _id: this.group_id
       },
       [
         {
           $set: {
-            last_checkpoint_lsn: { $literal: lsn },
-            last_checkpoint_ts: { $literal: now },
+            _can_checkpoint: {
+              $and: [
+                { $eq: ['$snapshot_done', true] },
+                {
+                  $or: [{ $eq: ['$last_checkpoint_lsn', null] }, { $lte: ['$last_checkpoint_lsn', { $literal: lsn }] }]
+                },
+                {
+                  $or: [
+                    { $eq: ['$no_checkpoint_before', null] },
+                    { $lte: ['$no_checkpoint_before', { $literal: lsn }] }
+                  ]
+                }
+              ]
+            }
+          }
+        },
+        {
+          $set: {
+            last_checkpoint_lsn: {
+              $cond: ['$_can_checkpoint', { $literal: lsn }, '$last_checkpoint_lsn']
+            },
+            last_checkpoint_ts: {
+              $cond: ['$_can_checkpoint', { $literal: now }, '$last_checkpoint_ts']
+            },
             last_keepalive_ts: { $literal: now },
-            last_fatal_error: { $literal: null },
-            keepalive_op: { $literal: null },
+            last_fatal_error: {
+              $cond: ['$_can_checkpoint', { $literal: null }, '$last_fatal_error']
+            },
+            keepalive_op: {
+              $cond: [
+                '$_can_checkpoint',
+                // Checkpoint: set to null
+                { $literal: null },
+                // Not checkpoint: update to max of existing keepalive_op and persisted_op
+                {
+                  $toString: {
+                    $max: [{ $toLong: '$keepalive_op' }, { $literal: this.persisted_op }]
+                  }
+                }
+              ]
+            },
             last_checkpoint: {
-              $max: ['$last_checkpoint', { $literal: this.persisted_op }, { $toLong: '$keepalive_op' }]
+              $cond: [
+                '$_can_checkpoint',
+                {
+                  $max: ['$last_checkpoint', { $literal: this.persisted_op }, { $toLong: '$keepalive_op' }]
+                },
+                '$last_checkpoint'
+              ]
             }
           }
+        },
+        {
+          $unset: '_can_checkpoint'
         }
       ],
-      { session: this.session }
+      {
+        session: this.session,
+        returnDocument: 'after',
+        projection: {
+          snapshot_done: 1,
+          last_checkpoint_lsn: 1,
+          no_checkpoint_before: 1,
+          keepalive_op: 1
+        }
+      }
     );
-    if (updateResult.matchedCount == 0) {
+    if (updateResult == null) {
+      throw new ReplicationAssertionError('Failed to load sync_rules document during checkpoint update');
+    }
+    const checkpointCreated = updateResult.snapshot_done === true && updateResult.last_checkpoint_lsn === lsn;
+
+    if (!checkpointCreated) {
       // Failed on snapshot_done or no_checkpoint_before.
-      if (Date.now() - this.lastWaitingLogThottled > 5_000) {
+      if (Date.now() - this.lastWaitingLogThottled > 5_000 || true) {
         this.logger.info(
-          `Waiting until before creating checkpoint, currently at ${lsn}. Persisted op: ${this.persisted_op}`
-        );
-        this.lastWaitingLogThottled = Date.now();
-      }
-
-      if (this.persisted_op != null) {
-        await this.db.sync_rules.updateOne(
-          {
-            _id: this.group_id
-          },
-          [
-            // KLUDGE: the string format is a pain here, not sure why we ever had it as a string
+          `Waiting before creating checkpoint, currently at ${lsn}. Persisted op: ${this.persisted_op}. Current state: ${JSON.stringify(
             {
-              $set: {
-                keepalive_op: {
-                  $toString: {
-                    $max: [{ $toLong: '$keepalive_op' }, { $literal: this.persisted_op }]
-                  }
-                }
-              }
+              snapshot_done: updateResult.snapshot_done,
+              last_checkpoint_lsn: updateResult.last_checkpoint_lsn,
+              no_checkpoint_before: updateResult.no_checkpoint_before
             }
-          ],
-          { session: this.session }
+          )}`
         );
+        this.lastWaitingLogThottled = Date.now();
       }
     } else {
+      this.logger.info(`Created checkpoint at ${lsn}. Persisted op: ${this.persisted_op}`);
       await this.autoActivate(lsn);
       await this.db.notifyCheckpoint();
       this.persisted_op = null;
diff --git a/modules/module-postgres/src/replication/PostgresSnapshotter.ts b/modules/module-postgres/src/replication/PostgresSnapshotter.ts
@@ -190,7 +190,6 @@ export class PostgresSnapshotter {
           objectId: relid,
           replicaIdColumns: cresult.replicationColumns
         } as SourceEntityDescriptor,
-        snapshot: false,
         referencedTypeIds: columnTypes
       });
 
@@ -329,8 +328,12 @@ export class PostgresSnapshotter {
         },
         async (batch) => {
           await this.snapshotTableInTx(batch, db, table);
+          // This commit ensures we set keepalive_op.
+          // It may be better if that is automatically set when flushing.
+          await batch.commit(ZERO_LSN);
         }
       );
+      this.logger.info(`Flushed snapshot at ${flushResults?.flushed_op}`);
     } finally {
       await db.end();
     }
@@ -442,7 +445,8 @@ export class PostgresSnapshotter {
     }
   }
 
-  public async queueSnapshot(table: storage.SourceTable) {
+  public async queueSnapshot(batch: storage.BucketStorageBatch, table: storage.SourceTable) {
+    await batch.markTableSnapshotRequired(table);
     this.queue.add(table);
   }
 
@@ -477,6 +481,7 @@ export class PostgresSnapshotter {
       tableLsnNotBefore = rs.rows[0][0];
       // Side note: A ROLLBACK would probably also be fine here, since we only read in this transaction.
       await db.query('COMMIT');
+      this.logger.info(`Snapshot complete for table ${table.qualifiedName}, resume at ${tableLsnNotBefore}`);
       const [resultTable] = await batch.markTableSnapshotDone([table], tableLsnNotBefore);
       this.relationCache.update(resultTable);
       return resultTable;
@@ -614,10 +619,9 @@ export class PostgresSnapshotter {
   async handleRelation(options: {
     batch: storage.BucketStorageBatch;
     descriptor: SourceEntityDescriptor;
-    snapshot: boolean;
     referencedTypeIds: number[];
   }) {
-    const { batch, descriptor, snapshot, referencedTypeIds } = options;
+    const { batch, descriptor, referencedTypeIds } = options;
 
     if (!descriptor.objectId && typeof descriptor.objectId != 'number') {
       throw new ReplicationAssertionError(`objectId expected, got ${typeof descriptor.objectId}`);
@@ -637,31 +641,6 @@ export class PostgresSnapshotter {
     // Ensure we have a description for custom types referenced in the table.
     await this.connections.types.fetchTypes(referencedTypeIds);
 
-    // Snapshot if:
-    // 1. Snapshot is requested (false for initial snapshot, since that process handles it elsewhere)
-    // 2. Snapshot is not already done, AND:
-    // 3. The table is used in sync rules.
-    const shouldSnapshot = snapshot && !result.table.snapshotComplete && result.table.syncAny;
-
-    if (shouldSnapshot) {
-      // Truncate this table, in case a previous snapshot was interrupted.
-      await batch.truncate([result.table]);
-
-      // Start the snapshot inside a transaction.
-      // We use a dedicated connection for this.
-      const db = await this.connections.snapshotConnection();
-      try {
-        const table = await this.snapshotTableInTx(batch, db, result.table);
-        // After the table snapshot, we wait for replication to catch up.
-        // To make sure there is actually something to replicate, we send a keepalive
-        // message.
-        await sendKeepAlive(db);
-        return table;
-      } finally {
-        await db.end();
-      }
-    }
-
     return result.table;
   }
 
diff --git a/modules/module-postgres/src/replication/WalStream.ts b/modules/module-postgres/src/replication/WalStream.ts
@@ -183,7 +183,10 @@ export class WalStream {
     this.relationCache.update(result.table);
 
     // Drop conflicting tables. This includes for example renamed tables.
-    await batch.drop(result.dropTables);
+    if (result.dropTables.length > 0) {
+      this.logger.info(`Dropping conflicting tables: ${result.dropTables.map((t) => t.qualifiedName).join(', ')}`);
+      await batch.drop(result.dropTables);
+    }
 
     // Ensure we have a description for custom types referenced in the table.
     await this.connections.types.fetchTypes(referencedTypeIds);
@@ -195,7 +198,8 @@ export class WalStream {
     const shouldSnapshot = snapshot && !result.table.snapshotComplete && result.table.syncAny;
 
     if (shouldSnapshot) {
-      await this.snapshotter.queueSnapshot(result.table);
+      this.logger.info(`Queuing snapshot for new table ${result.table.qualifiedName}`);
+      await this.snapshotter.queueSnapshot(batch, result.table);
     }
 
     return result.table;
diff --git a/modules/module-postgres/test/src/schema_changes.test.ts b/modules/module-postgres/test/src/schema_changes.test.ts
