Split out "snapshot done" check.

Author: rkistner

modules/module-mongodb-storage/src/storage/implementation/MongoBucketBatch.ts

@@ -19,6 +19,7 @@ import {
   isCompleteRow,
   SaveOperationTag,
   storage,
+  SyncRuleState,
   utils
 } from '@powersync/service-core';
 import * as timers from 'node:timers/promises';
@@ -713,7 +714,6 @@ export class MongoBucketBatch
       last_checkpoint_lsn: lsn,
       last_checkpoint_ts: now,
       last_keepalive_ts: now,
-      snapshot_done: true,
       last_fatal_error: null,
       keepalive_op: null
     };
@@ -744,8 +744,7 @@ export class MongoBucketBatch
         _id: this.group_id
       },
       {
-        $set: update,
-        $unset: { snapshot_lsn: 1 }
+        $set: update
       },
       { session: this.session }
     );
@@ -773,7 +772,7 @@ export class MongoBucketBatch
     let activated = false;
     await session.withTransaction(async () => {
       const doc = await this.db.sync_rules.findOne({ _id: this.group_id }, { session });
-      if (doc && doc.state == 'PROCESSING') {
+      if (doc && doc.state == SyncRuleState.PROCESSING && doc.snapshot_done && doc.last_checkpoint != null) {
         await this.db.sync_rules.updateOne(
           {
             _id: this.group_id
@@ -799,13 +798,15 @@ export class MongoBucketBatch
           { session }
         );
         activated = true;
+      } else if (doc?.state != SyncRuleState.PROCESSING) {
+        this.needsActivation = false;
       }
     });
     if (activated) {
       this.logger.info(`Activated new sync rules at ${lsn}`);
       await this.db.notifyCheckpoint();
+      this.needsActivation = false;
     }
-    this.needsActivation = false;
   }
 
   async keepalive(lsn: string): Promise<boolean> {
@@ -857,11 +858,9 @@ export class MongoBucketBatch
       {
         $set: {
           last_checkpoint_lsn: lsn,
-          snapshot_done: true,
           last_fatal_error: null,
           last_keepalive_ts: new Date()
-        },
-        $unset: { snapshot_lsn: 1 }
+        }
       },
       { session: this.session }
     );
@@ -1044,7 +1043,31 @@ export class MongoBucketBatch
     return copy;
   }
 
-  async markSnapshotDone(tables: storage.SourceTable[], no_checkpoint_before_lsn: string) {
+  async markAllSnapshotDone(no_checkpoint_before_lsn: string) {
+    if (this.no_checkpoint_before_lsn == null || no_checkpoint_before_lsn > this.no_checkpoint_before_lsn) {
+      this.no_checkpoint_before_lsn = no_checkpoint_before_lsn;
+    }
+    await this.db.sync_rules.updateOne(
+      {
+        _id: this.group_id
+      },
+      {
+        $set: {
+          snapshot_done: true,
+          last_keepalive_ts: new Date()
+        },
+        $max: {
+          no_checkpoint_before: no_checkpoint_before_lsn
+        },
+        $unset: {
+          snapshot_lsn: 1
+        }
+      },
+      { session: this.session }
+    );
+  }
+
+  async markTableSnapshotDone(tables: storage.SourceTable[], no_checkpoint_before_lsn?: string) {
     const session = this.session;
     const ids = tables.map((table) => table.id);
 
@@ -1062,17 +1085,20 @@ export class MongoBucketBatch
         { session }
       );
 
-      if (no_checkpoint_before_lsn > this.no_checkpoint_before_lsn) {
-        this.no_checkpoint_before_lsn = no_checkpoint_before_lsn;
-
+      if (no_checkpoint_before_lsn != null) {
+        if (this.no_checkpoint_before_lsn == null || no_checkpoint_before_lsn > this.no_checkpoint_before_lsn) {
+          this.no_checkpoint_before_lsn = no_checkpoint_before_lsn;
+        }
         await this.db.sync_rules.updateOne(
           {
             _id: this.group_id
           },
           {
             $set: {
-              no_checkpoint_before: no_checkpoint_before_lsn,
               last_keepalive_ts: new Date()
+            },
+            $max: {
+              no_checkpoint_before: no_checkpoint_before_lsn
             }
           },
           { session: this.session }

modules/module-mongodb/src/replication/ChangeStream.ts

@@ -366,7 +366,7 @@ export class ChangeStream {
 
         for (let table of tablesWithStatus) {
           await this.snapshotTable(batch, table);
-          await batch.markSnapshotDone([table], MongoLSN.ZERO.comparable);
+          await batch.markTableSnapshotDone([table]);
 
           this.touch();
         }
@@ -375,7 +375,7 @@ export class ChangeStream {
         // point before the data can be considered consistent.
         // We could do this for each individual table, but may as well just do it once for the entire snapshot.
         const checkpoint = await createCheckpoint(this.client, this.defaultDb, STANDALONE_CHECKPOINT_ID);
-        await batch.markSnapshotDone([], checkpoint);
+        await batch.markAllSnapshotDone(checkpoint);
 
         // This will not create a consistent checkpoint yet, but will persist the op.
         // Actual checkpoint will be created when streaming replication caught up.
@@ -612,7 +612,7 @@ export class ChangeStream {
       await this.snapshotTable(batch, result.table);
       const no_checkpoint_before_lsn = await createCheckpoint(this.client, this.defaultDb, STANDALONE_CHECKPOINT_ID);
 
-      const [table] = await batch.markSnapshotDone([result.table], no_checkpoint_before_lsn);
+      const [table] = await batch.markTableSnapshotDone([result.table], no_checkpoint_before_lsn);
       return table;
     }
 

modules/module-mysql/src/replication/BinLogStream.ts

@@ -170,7 +170,7 @@ export class BinLogStream {
       } finally {
         connection.release();
       }
-      const [table] = await batch.markSnapshotDone([result.table], gtid.comparable);
+      const [table] = await batch.markTableSnapshotDone([result.table], gtid.comparable);
       return table;
     }
 
@@ -275,10 +275,12 @@ export class BinLogStream {
             const tables = await this.getQualifiedTableNames(batch, tablePattern);
             for (let table of tables) {
               await this.snapshotTable(connection as mysql.Connection, batch, table);
-              await batch.markSnapshotDone([table], headGTID.comparable);
+              await batch.markTableSnapshotDone([table], headGTID.comparable);
               await framework.container.probes.touch();
             }
           }
+          const snapshotDoneGtid = await common.readExecutedGtid(promiseConnection);
+          await batch.markAllSnapshotDone(snapshotDoneGtid.comparable);
           await batch.commit(headGTID.comparable);
         }
       );

modules/module-postgres-storage/src/storage/batch/PostgresBucketBatch.ts

@@ -430,7 +430,7 @@ export class PostgresBucketBatch
     `.execute();
   }
 
-  async markSnapshotDone(
+  async markTableSnapshotDone(
     tables: storage.SourceTable[],
     no_checkpoint_before_lsn: string
   ): Promise<storage.SourceTable[]> {

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

@@ -355,6 +355,9 @@ export class PostgresSnapshotter {
         // This makes sure we don't skip any changes applied before starting this snapshot,
         // in the case of snapshot retries.
         // We could alternatively commit at the replication slot LSN.
+        const rs = await db.query(`select pg_current_wal_lsn() as lsn`);
+        const globalLsnNotBefore = rs.rows[0][0];
+        await batch.markAllSnapshotDone(globalLsnNotBefore);
         await batch.commit(ZERO_LSN);
       }
     );
@@ -416,7 +419,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');
-      const [resultTable] = await batch.markSnapshotDone([table], tableLsnNotBefore);
+      const [resultTable] = await batch.markTableSnapshotDone([table], tableLsnNotBefore);
       this.relationCache.update(resultTable);
       return resultTable;
     } catch (e) {

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

@@ -481,7 +481,7 @@ export class WalStream {
               await this.handleRelation({
                 batch,
                 descriptor: getPgOutputRelation(msg),
-                snapshot: true,
+                snapshot: false,
                 referencedTypeIds: referencedColumnTypeIds(msg)
               });
             } else if (msg.tag == 'begin') {

packages/service-core/src/storage/BucketStorageBatch.ts

@@ -83,7 +83,8 @@ export interface BucketStorageBatch extends ObserverClient<BucketBatchStorageLis
    */
   resumeFromLsn: string | null;
 
-  markSnapshotDone(tables: SourceTable[], no_checkpoint_before_lsn: string): Promise<SourceTable[]>;
+  markTableSnapshotDone(tables: SourceTable[], no_checkpoint_before_lsn?: string): Promise<SourceTable[]>;
+  markAllSnapshotDone(no_checkpoint_before_lsn: string): Promise<void>;
 
   updateTableProgress(table: SourceTable, progress: Partial<TableSnapshotStatus>): Promise<SourceTable>;