[Incremental Reprocessing] [MongoDB] stream while snapshotting (#641)

* [WIP] Concurrent snapshot + streaming for MongoDB.

* Cleanup dead code.

* Use native Promise.withResolvers()

* Fix queue race conditions.

* Rename abortSignal.

* Add failing regression test for resurrected deletes.

* Avoid concurrent snapshots on older storage versions.

* Fix uncaught rejection.

* Fix test.

* Add changeset.

* Add comments back.

* Redo populatePersistentChecksumCache ordering.

* Guard snapshot complete condition.

* Handle non-concurrent inline snapshots.

* Bypass the snapshot guard for tests.

* SQL server: Don't resolve tables we can't snapshot.

* "Capture instance created" is no longer a schema change.

* Handle snapshot queue race conditions.

* Update changeset.

* Refactor.

* Fix resnapshotting tables.

* Remove relationCache from MongoSnapshotter.

* Simplify relationCache usage.

* Clarify replication error handling.

* Move concurrency logic to MongoSnapshotter.

* Use read-only table status when snapshotting.

* Remove dead code.

Author: rkistner

.changeset/slow-boats-eat.md

@@ -0,0 +1,11 @@
+---
+'@powersync/service-module-postgres-storage': minor
+'@powersync/service-module-mongodb-storage': minor
+'@powersync/service-module-postgres': minor
+'@powersync/service-module-mongodb': minor
+'@powersync/service-core': minor
+'@powersync/service-module-mssql': minor
+'@powersync/service-module-mysql': minor
+---
+
+[MongoDB] Support snapshotting concurrently with streaming in storage v3+.

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

@@ -56,6 +56,7 @@ export interface MongoBucketBatchOptions {
   skipExistingRows: boolean;
 
   markRecordUnavailable: BucketStorageMarkRecordUnavailable | undefined;
+  hooks: storage.StorageHooks | undefined;
 
   logger: Logger;
   tracer?: PerformanceTracer<'storage' | 'evaluate'>;
@@ -77,12 +78,13 @@ export abstract class MongoBucketBatch
 
   private readonly slot_name: string;
   private readonly storeCurrentData: boolean;
-  private readonly skipExistingRows: boolean;
+  public readonly skipExistingRows: boolean;
   protected readonly mapping: BucketDefinitionMapping;
 
   private batch: OperationBatch | null = null;
   private write_checkpoint_batch: storage.CustomWriteCheckpointOptions[] = [];
   private markRecordUnavailable: BucketStorageMarkRecordUnavailable | undefined;
+  private hooks: storage.StorageHooks | undefined;
   private clearedError = false;
 
   private tracer: PerformanceTracer<'storage' | 'evaluate'>;
@@ -131,6 +133,7 @@ export abstract class MongoBucketBatch
     this.mapping = options.mapping;
     this.skipExistingRows = options.skipExistingRows;
     this.markRecordUnavailable = options.markRecordUnavailable;
+    this.hooks = options.hooks;
     this.batch = new OperationBatch();
 
     this.persisted_op = options.keepaliveOp ?? null;
@@ -162,8 +165,12 @@ export abstract class MongoBucketBatch
 
   abstract setResumeLsn(lsn: string): Promise<void>;
 
+  abstract getSourceTableStatus(table: storage.SourceTable): Promise<storage.SourceTable | null>;
+
   abstract markAllSnapshotDone(no_checkpoint_before_lsn: string): Promise<void>;
 
+  abstract markSnapshotDone(no_checkpoint_before_lsn: string, options?: { throwOnConflict?: boolean }): Promise<void>;
+
   abstract markTableSnapshotRequired(table: storage.SourceTable): Promise<void>;
 
   abstract markTableSnapshotDone(
@@ -191,6 +198,8 @@ export abstract class MongoBucketBatch
 
     using _ = this.tracer.span('storage', 'flush');
 
+    await this.hooks?.beforeBatchFlush?.(this);
+
     await this.withReplicationTransaction(`Flushing ${batch?.length ?? 0} ops`, async (session, opSeq) => {
       if (batch != null) {
         resumeBatch = await this.replicateBatch(session, batch, opSeq, options);
@@ -214,6 +223,7 @@ export abstract class MongoBucketBatch
 
     this.persisted_op = last_op;
     this.last_flushed_op = last_op;
+    await this.hooks?.afterBatchFlush?.(this);
     return { flushed_op: last_op };
   }
 

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

@@ -77,6 +77,7 @@ export abstract class MongoSyncBucketStorage
   private parsedSyncRulesCache: { parsed: HydratedSyncRules; options: storage.ParseSyncRulesOptions } | undefined;
   private writeCheckpointAPI: MongoWriteCheckpointAPI;
   public readonly logger: Logger;
+  public readonly storageConfig: StorageConfig;
   #storageInitialized = false;
 
   constructor(
@@ -88,7 +89,8 @@ export abstract class MongoSyncBucketStorage
     options: MongoSyncBucketStorageOptions
   ) {
     super();
-    this.db = factory.db.versioned(sync_rules.getStorageConfig());
+    this.storageConfig = options.storageConfig;
+    this.db = factory.db.versioned(this.storageConfig);
     this.checksums = this.createMongoChecksums(options);
     this.writeCheckpointAPI = new MongoWriteCheckpointAPI({
       db: this.db,
@@ -207,6 +209,7 @@ export abstract class MongoSyncBucketStorage
       storeCurrentData: options.storeCurrentData,
       skipExistingRows: options.skipExistingRows ?? false,
       markRecordUnavailable: options.markRecordUnavailable,
+      hooks: options.hooks,
       syncConfigId: state.syncConfigId,
       tracer: options.tracer
     };

modules/module-mongodb-storage/src/storage/implementation/v1/MongoBucketBatchV1.ts

@@ -9,7 +9,7 @@ import { SourceRecordStore } from '../common/SourceRecordStore.js';
 import { PersistedBatchV1 } from './PersistedBatchV1.js';
 import { SourceRecordStoreV1 } from './SourceRecordStoreV1.js';
 import { VersionedPowerSyncMongoV1 } from './VersionedPowerSyncMongoV1.js';
-import { SyncRuleDocumentV1 } from './models.js';
+import { SourceTableDocumentV1, SyncRuleDocumentV1 } from './models.js';
 
 export class MongoBucketBatchV1 extends MongoBucketBatch {
   declare public readonly db: VersionedPowerSyncMongoV1;
@@ -143,6 +143,48 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
     return result!;
   }
 
+  async getSourceTableStatus(table: storage.SourceTable): Promise<storage.SourceTable | null> {
+    const doc = (await this.db.commonSourceTables(this.group_id).findOne(
+      {
+        group_id: this.group_id,
+        _id: mongoTableId(table.id)
+      },
+      { session: this.session }
+    )) as SourceTableDocumentV1 | null;
+    if (doc == null) {
+      return null;
+    }
+
+    const ref = {
+      connectionTag: table.ref.connectionTag,
+      schema: doc.schema_name,
+      name: doc.table_name
+    };
+    const sourceTable = new storage.SourceTable({
+      id: doc._id,
+      ref,
+      objectId: doc.relation_id,
+      replicaIdColumns:
+        doc.replica_id_columns2?.map(
+          (c) => ({ name: c.name, typeId: c.type_oid, type: c.type }) satisfies ColumnDescriptor
+        ) ?? [],
+      snapshotComplete: doc.snapshot_done ?? true,
+      ...this.sync_rules.getMatchingSources(ref)
+    });
+    sourceTable.syncEvent = this.sync_rules.tableTriggersEvent(ref);
+    sourceTable.syncData = sourceTable.bucketDataSources.length > 0;
+    sourceTable.syncParameters = sourceTable.parameterLookupSources.length > 0;
+    sourceTable.snapshotStatus =
+      doc.snapshot_status == null
+        ? undefined
+        : {
+            lastKey: doc.snapshot_status.last_key?.buffer ?? null,
+            totalEstimatedCount: doc.snapshot_status.total_estimated_count,
+            replicatedCount: doc.snapshot_status.replicated_count
+          };
+    return sourceTable;
+  }
+
   async commit(lsn: string, options?: storage.BucketBatchCommitOptions): Promise<storage.CheckpointResult> {
     const { createEmptyCheckpoints } = { ...storage.DEFAULT_BUCKET_BATCH_COMMIT_OPTIONS, ...options };
 
@@ -335,6 +377,30 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
     );
   }
 
+  async markSnapshotDone(no_checkpoint_before_lsn: string, options?: { throwOnConflict?: boolean }): Promise<void> {
+    await this.withTransaction(async () => {
+      // Protect against race conditions
+      const count = await this.db.commonSourceTables(this.group_id).countDocuments(
+        {
+          group_id: this.group_id,
+          snapshot_done: false
+        },
+        { session: this.session }
+      );
+      if (count > 0) {
+        if (options?.throwOnConflict ?? true) {
+          throw new ReplicationAssertionError(
+            `Cannot mark snapshot done while ${count} source table${count == 1 ? '' : 's'} still require snapshotting`
+          );
+        } else {
+          return;
+        }
+      }
+
+      await this.markAllSnapshotDone(no_checkpoint_before_lsn);
+    });
+  }
+
   async markTableSnapshotRequired(_table: storage.SourceTable): Promise<void> {
     await this.db.sync_rules.updateOne(
       {

modules/module-mongodb-storage/src/storage/implementation/v1/MongoSyncBucketStorageV1.ts

@@ -124,7 +124,8 @@ export class MongoSyncBucketStorageV1 extends MongoSyncBucketStorage {
           snapshot_done: 1,
           last_checkpoint_lsn: 1,
           state: 1,
-          snapshot_lsn: 1
+          snapshot_lsn: 1,
+          keepalive_op: 1
         }
       }
     )) as SyncRuleDocumentV1;
@@ -136,7 +137,8 @@ export class MongoSyncBucketStorageV1 extends MongoSyncBucketStorage {
       snapshot_done: doc.snapshot_done,
       snapshot_lsn: doc.snapshot_lsn ?? null,
       active: doc.state == storage.SyncRuleState.ACTIVE,
-      checkpoint_lsn: doc.last_checkpoint_lsn
+      checkpoint_lsn: doc.last_checkpoint_lsn,
+      keepalive_op: doc.keepalive_op == null ? null : BigInt(doc.keepalive_op)
     };
   }
 

modules/module-mongodb-storage/src/storage/implementation/v3/MongoBucketBatchV3.ts

@@ -280,6 +280,17 @@ export class MongoBucketBatchV3 extends MongoBucketBatch {
     };
   }
 
+  async getSourceTableStatus(table: storage.SourceTable): Promise<storage.SourceTable | null> {
+    const doc = (await this.db
+      .commonSourceTables(this.group_id)
+      .findOne({ _id: mongoTableId(table.id) }, { session: this.session })) as SourceTableDocumentV3 | null;
+    if (doc == null) {
+      return null;
+    }
+
+    return this.sourceTableFromDocument(doc, table.ref.connectionTag, this.sync_rules);
+  }
+
   async commit(lsn: string, options?: storage.BucketBatchCommitOptions): Promise<storage.CheckpointResult> {
     const { createEmptyCheckpoints } = { ...storage.DEFAULT_BUCKET_BATCH_COMMIT_OPTIONS, ...options };
 
@@ -507,6 +518,29 @@ export class MongoBucketBatchV3 extends MongoBucketBatch {
     );
   }
 
+  async markSnapshotDone(no_checkpoint_before_lsn: string, options?: { throwOnConflict?: boolean }): Promise<void> {
+    await this.withTransaction(async () => {
+      // Protect against race conditions
+      const count = await this.db.sourceTablesV3(this.group_id).countDocuments(
+        {
+          snapshot_done: false
+        },
+        { session: this.session }
+      );
+      if (count > 0) {
+        if (options?.throwOnConflict ?? true) {
+          throw new ReplicationAssertionError(
+            `Cannot mark snapshot done while ${count} source table${count == 1 ? '' : 's'} still require snapshotting`
+          );
+        } else {
+          return;
+        }
+      }
+
+      await this.markAllSnapshotDone(no_checkpoint_before_lsn);
+    });
+  }
+
   async markTableSnapshotRequired(_table: storage.SourceTable): Promise<void> {
     await this.db.sync_rules.updateOne(
       {

modules/module-mongodb-storage/src/storage/implementation/v3/MongoSyncBucketStorageV3.ts

@@ -215,7 +215,7 @@ export class MongoSyncBucketStorageV3 extends MongoSyncBucketStorage {
 
   protected async getStatusImpl(): Promise<storage.SyncRuleStatus> {
     const doc = await this.syncRulesCollection.findOne(this.syncConfigMatch(), {
-      projection: this.syncConfigProjection({ state: 1, snapshot_lsn: 1 })
+      projection: this.syncConfigProjection({ state: 1, snapshot_lsn: 1, keepalive_op: 1 })
     });
     const syncConfig = this.selectedSyncConfig(doc);
     if (doc == null || syncConfig == null) {
@@ -226,7 +226,8 @@ export class MongoSyncBucketStorageV3 extends MongoSyncBucketStorage {
       snapshot_done: syncConfig.snapshot_done ?? false,
       snapshot_lsn: doc.snapshot_lsn ?? null,
       active: doc.state == storage.SyncRuleState.ACTIVE && syncConfig.state == storage.SyncRuleState.ACTIVE,
-      checkpoint_lsn: syncConfig.last_checkpoint_lsn ?? null
+      checkpoint_lsn: syncConfig.last_checkpoint_lsn ?? null,
+      keepalive_op: syncConfig.keepalive_op ?? null
     };
   }