@powersync/service-module-mongodb-storage 0.21.0 → 0.23.0

package/src/storage/implementation/v1/MongoBucketBatchV1.ts

@@ -6,8 +6,11 @@ import * as bson from 'bson';
 import { mongoTableId } from '../../../utils/util.js';
 import { calculateCheckpointState } from '../CheckpointState.js';
 import { MongoBucketBatch, MongoBucketBatchOptions } from '../MongoBucketBatch.js';
+import { MongoSyncRulesLock } from '../MongoSyncRulesLock.js';
+import { MongoWriteBatch } from '../MongoWriteBatch.js';
 import { PersistedBatch } from '../common/PersistedBatch.js';
 import { SourceRecordStore } from '../common/SourceRecordStore.js';
+import { SyncRuleDocumentBase } from '../models.js';
 import { PersistedBatchV1 } from './PersistedBatchV1.js';
 import { SourceRecordStoreV1 } from './SourceRecordStoreV1.js';
 import { VersionedPowerSyncMongoV1 } from './VersionedPowerSyncMongoV1.js';
@@ -34,6 +37,27 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
     this.store = new SourceRecordStoreV1(this.db, this.replicationStreamId);
   }
 
+  protected override persistedOpHead(stream: SyncRuleDocumentBase): InternalOpId {
+    const legacy = stream as SyncRuleDocumentV1;
+    const checkpoint = legacy.last_checkpoint ?? 0n;
+    const keepalive = legacy.keepalive_op == null ? 0n : BigInt(legacy.keepalive_op);
+    return checkpoint > keepalive ? checkpoint : keepalive;
+  }
+
+  protected override onReplicationTransactionFlush(writes: MongoWriteBatch, lastOp: InternalOpId): void {
+    // Keep flushed operations recoverable before a checkpoint can be published.
+    // keepalive_op remains a decimal string for compatibility with legacy readers.
+    writes.updateOne(this.db.sync_rules, { _id: this.replicationStreamId }, [
+      {
+        $set: {
+          keepalive_op: {
+            $toString: { $max: [{ $toLong: '$keepalive_op' }, '$last_checkpoint', lastOp] }
+          }
+        }
+      }
+    ]);
+  }
+
   protected override recordPersistedOp(lastOp: InternalOpId): void {
     this.persisted_op = lastOp;
   }
@@ -55,8 +79,15 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
     session: mongo.ClientSession,
     opId: InternalOpId
   ): Promise<void> {
+    // Only the final checkpoint per user/stream is visible after this transaction.
+    const checkpoints = new Map(
+      this.write_checkpoint_batch.map((checkpoint) => [
+        JSON.stringify([checkpoint.sync_rules_id, checkpoint.user_id]),
+        checkpoint
+      ])
+    );
     await this.db.custom_write_checkpoints.bulkWrite(
-      this.write_checkpoint_batch.map((checkpointOptions) => {
+      [...checkpoints.values()].map((checkpointOptions) => {
         const set: Record<string, unknown> = {
           checkpoint: checkpointOptions.checkpoint,
           sync_rules_id: checkpointOptions.sync_rules_id,
@@ -77,7 +108,7 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
           }
         };
       }),
-      { session }
+      { session, ordered: false }
     );
   }
 
@@ -94,7 +125,8 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
     }));
 
     let result: storage.ResolveTablesResult | null = null;
-    await this.db.client.withSession(async (session) => {
+    const session = this.session;
+    await this.withTransaction(async () => {
       const col = this.db.sourceTablesV1(this.replicationStreamId);
 
       // Find records that overlap by name or relation id.
@@ -212,129 +244,136 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
   async commit(lsn: string, options?: storage.BucketBatchCommitOptions): Promise<storage.CheckpointResult> {
     const { createEmptyCheckpoints } = { ...storage.DEFAULT_BUCKET_BATCH_COMMIT_OPTIONS, ...options };
 
-    await this.flush(options);
-
     using _ = this.tracer.span('storage', 'commit');
 
-    const now = new Date();
+    const updateCheckpoint = async (persistedOp = this.persisted_op) => {
+      const now = new Date();
 
-    await this.db.write_checkpoints.updateMany(
-      {
-        processed_at_lsn: null,
-        'lsns.1': { $lte: lsn }
-      },
-      {
-        $set: {
-          processed_at_lsn: lsn
-        }
-      },
-      {
-        session: this.session
-      }
-    );
+      const 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 }] }]
+          }
+        ]
+      };
 
-    const 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 }] }]
-        }
-      ]
-    };
+      const new_keepalive_op = {
+        $cond: [
+          can_checkpoint,
+          { $literal: null },
+          {
+            $toString: {
+              $max: [{ $toLong: '$keepalive_op' }, { $literal: persistedOp }, 0n]
+            }
+          }
+        ]
+      };
 
-    const new_keepalive_op = {
-      $cond: [
-        can_checkpoint,
-        { $literal: null },
+      const new_last_checkpoint = {
+        $cond: [
+          can_checkpoint,
+          {
+            $max: ['$last_checkpoint', { $literal: persistedOp }, { $toLong: '$keepalive_op' }, 0n]
+          },
+          '$last_checkpoint'
+        ]
+      };
+
+      const preUpdateDocument = (await this.db.sync_rules.findOneAndUpdate(
+        MongoSyncRulesLock.ownerFilter(this.replicationStreamId, this.options.replicationLock),
+        [
+          {
+            $set: {
+              _can_checkpoint: can_checkpoint,
+              _not_empty: createEmptyCheckpoints
+                ? true
+                : {
+                    $or: [
+                      { $literal: createEmptyCheckpoints },
+                      { $ne: ['$keepalive_op', new_keepalive_op] },
+                      { $ne: ['$last_checkpoint', new_last_checkpoint] }
+                    ]
+                  }
+            }
+          },
+          {
+            $set: {
+              ...MongoSyncRulesLock.heartbeatUpdate(),
+              last_checkpoint_lsn: {
+                $cond: [{ $and: ['$_can_checkpoint', '$_not_empty'] }, { $literal: lsn }, '$last_checkpoint_lsn']
+              },
+              last_checkpoint_ts: {
+                $cond: [{ $and: ['$_can_checkpoint', '$_not_empty'] }, { $literal: now }, '$last_checkpoint_ts']
+              },
+
+              last_fatal_error: { $literal: null },
+              last_fatal_error_ts: { $literal: null },
+              keepalive_op: new_keepalive_op,
+              last_checkpoint: new_last_checkpoint,
+              snapshot_lsn: {
+                $cond: [{ $and: ['$_can_checkpoint', '$_not_empty'] }, { $literal: null }, '$snapshot_lsn']
+              }
+            }
+          },
+          {
+            $unset: ['_can_checkpoint', '_not_empty']
+          }
+        ],
         {
-          $toString: {
-            $max: [{ $toLong: '$keepalive_op' }, { $literal: this.persisted_op }, 0n]
+          session: this.session,
+          returnDocument: 'before',
+          projection: {
+            snapshot_done: 1,
+            last_checkpoint_lsn: 1,
+            no_checkpoint_before: 1,
+            keepalive_op: 1,
+            last_checkpoint: 1
           }
         }
-      ]
-    };
+      )) as SyncRuleDocumentV1;
 
-    const new_last_checkpoint = {
-      $cond: [
-        can_checkpoint,
-        {
-          $max: ['$last_checkpoint', { $literal: this.persisted_op }, { $toLong: '$keepalive_op' }, 0n]
-        },
-        '$last_checkpoint'
-      ]
+      return MongoSyncRulesLock.assertOwned(preUpdateDocument);
     };
 
-    const preUpdateDocument = (await this.db.sync_rules.findOneAndUpdate(
-      { _id: this.replicationStreamId },
-      [
+    const finishCheckpoint = async (preUpdateDocument: SyncRuleDocumentV1, lastOp?: InternalOpId) => {
+      await this.db.write_checkpoints.updateMany(
         {
-          $set: {
-            _can_checkpoint: can_checkpoint,
-            _not_empty: createEmptyCheckpoints
-              ? true
-              : {
-                  $or: [
-                    { $literal: createEmptyCheckpoints },
-                    { $ne: ['$keepalive_op', new_keepalive_op] },
-                    { $ne: ['$last_checkpoint', new_last_checkpoint] }
-                  ]
-                }
-          }
+          processed_at_lsn: null,
+          'lsns.1': { $lte: lsn }
         },
         {
           $set: {
-            last_checkpoint_lsn: {
-              $cond: [{ $and: ['$_can_checkpoint', '$_not_empty'] }, { $literal: lsn }, '$last_checkpoint_lsn']
-            },
-            last_checkpoint_ts: {
-              $cond: [{ $and: ['$_can_checkpoint', '$_not_empty'] }, { $literal: now }, '$last_checkpoint_ts']
-            },
-            last_keepalive_ts: { $literal: now },
-            last_fatal_error: { $literal: null },
-            last_fatal_error_ts: { $literal: null },
-            keepalive_op: new_keepalive_op,
-            last_checkpoint: new_last_checkpoint,
-            snapshot_lsn: {
-              $cond: [{ $and: ['$_can_checkpoint', '$_not_empty'] }, { $literal: null }, '$snapshot_lsn']
-            }
+            processed_at_lsn: lsn
           }
         },
         {
-          $unset: ['_can_checkpoint', '_not_empty']
+          session: this.session
         }
-      ],
-      {
-        session: this.session,
-        returnDocument: 'before',
-        projection: {
-          snapshot_done: 1,
-          last_checkpoint_lsn: 1,
-          no_checkpoint_before: 1,
-          keepalive_op: 1,
-          last_checkpoint: 1
-        }
-      }
-    )) as SyncRuleDocumentV1;
-
-    if (preUpdateDocument == null) {
-      throw new ReplicationAssertionError(
-        'Failed to update checkpoint - no matching sync_rules document for _id: ' + this.replicationStreamId
       );
-    }
 
-    const checkpointState = calculateCheckpointState({
-      lsn,
-      snapshotDone: preUpdateDocument.snapshot_done === true,
-      lastCheckpointLsn: preUpdateDocument.last_checkpoint_lsn,
-      noCheckpointBefore: preUpdateDocument.no_checkpoint_before,
-      keepaliveOp: preUpdateDocument.keepalive_op == null ? null : BigInt(preUpdateDocument.keepalive_op),
-      lastCheckpoint: preUpdateDocument.last_checkpoint,
-      persistedOp: this.persisted_op,
-      createEmptyCheckpoints
-    });
+      const checkpointState = calculateCheckpointState({
+        lsn,
+        snapshotDone: preUpdateDocument.snapshot_done === true,
+        lastCheckpointLsn: preUpdateDocument.last_checkpoint_lsn,
+        noCheckpointBefore: preUpdateDocument.no_checkpoint_before,
+        keepaliveOp: preUpdateDocument.keepalive_op == null ? null : BigInt(preUpdateDocument.keepalive_op),
+        lastCheckpoint: preUpdateDocument.last_checkpoint,
+        persistedOp: lastOp ?? this.persistedOpHead(preUpdateDocument),
+        createEmptyCheckpoints
+      });
+      return { checkpointState, preUpdateDocument };
+    };
+
+    const { checkpointState, preUpdateDocument } = await this.flushAndCommit(
+      async (_stream, lastOp) => finishCheckpoint(await updateCheckpoint(lastOp), lastOp),
+      // Without a flush, the checkpoint update itself fences the transaction.
+      () => this.withFencedTransaction(updateCheckpoint, finishCheckpoint),
+      options
+    );
     if (checkpointState.checkpointBlocked) {
       if (Date.now() - this.lastWaitingLogThrottled > 5_000) {
         this.logger.info(
@@ -372,41 +411,20 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
 
   async setResumeLsn(lsn: string): Promise<void> {
     using _ = this.tracer.span('storage', 'set_resume_lsn');
-    await this.db.sync_rules.updateOne(
-      {
-        _id: this.replicationStreamId
-      },
-      {
-        $set: {
-          snapshot_lsn: lsn
-        }
-      },
-      {
-        session: this.session,
-        // Losing occasional resume LSN is fine. That may mean reprocessing
-        // some source changes in some edge cases, which is not an issue since
-        // changes are processed in an idempotent way.
-        writeConcern: { w: 1 }
-      }
-    );
+    // Losing occasional resume LSN would only reprocess source changes.
+    // Keep the lease check atomic, but do not wait for majority replication of this resume hint.
+    await this.updateStreamMetadata({ snapshot_lsn: { $literal: lsn } }, {}, { w: 1 });
   }
 
   async markAllSnapshotDone(no_checkpoint_before_lsn: string): Promise<void> {
-    await this.db.sync_rules.updateOne(
-      {
-        _id: this.replicationStreamId
-      },
-      {
-        $set: {
-          snapshot_done: true,
-          last_keepalive_ts: new Date()
-        },
-        $max: {
-          no_checkpoint_before: no_checkpoint_before_lsn
-        }
-      },
-      { session: this.session }
-    );
+    await this.updateStreamMetadata(this.snapshotDoneUpdate(no_checkpoint_before_lsn));
+  }
+
+  private snapshotDoneUpdate(no_checkpoint_before_lsn: string) {
+    return {
+      snapshot_done: true,
+      no_checkpoint_before: { $max: ['$no_checkpoint_before', { $literal: no_checkpoint_before_lsn }] }
+    };
   }
 
   async markSnapshotDone(no_checkpoint_before_lsn: string, options?: { throwOnConflict?: boolean }): Promise<void> {
@@ -436,22 +454,12 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
         }
       }
 
-      await this.markAllSnapshotDone(no_checkpoint_before_lsn);
+      await this.updateStreamMetadataInTransaction(this.snapshotDoneUpdate(no_checkpoint_before_lsn));
     });
   }
 
   async markTableSnapshotRequired(_table: storage.SourceTable): Promise<void> {
-    await this.db.sync_rules.updateOne(
-      {
-        _id: this.replicationStreamId
-      },
-      {
-        $set: {
-          snapshot_done: false
-        }
-      },
-      { session: this.session }
-    );
+    await this.updateStreamMetadata({ snapshot_done: false });
   }
 
   async markTableSnapshotDone(
@@ -462,7 +470,9 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
     const ids = tables.map((table) => mongoTableId(table.id));
 
     await this.withTransaction(async () => {
-      await this.db.sourceTablesV1(this.replicationStreamId).updateMany(
+      const writes = this.db.createWriteBatch(session, { ordered: false });
+      writes.updateMany(
+        this.db.sourceTablesV1(this.replicationStreamId),
         { _id: { $in: ids } },
         {
           $set: {
@@ -471,26 +481,23 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
           $unset: {
             snapshot_status: 1
           }
-        },
-        { session }
+        }
       );
 
       if (no_checkpoint_before_lsn != null) {
-        await this.db.sync_rules.updateOne(
+        writes.updateOne(
+          this.db.sync_rules,
           {
             _id: this.replicationStreamId
           },
           {
-            $set: {
-              last_keepalive_ts: new Date()
-            },
             $max: {
               no_checkpoint_before: no_checkpoint_before_lsn
             }
-          },
-          { session: this.session }
+          }
         );
       }
+      await writes.execute();
     });
     return tables.map((table) => {
       const copy = table.clone();
@@ -507,7 +514,7 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
     const session = this.session;
     let activated = false;
     let needsFutureActivationCheck = true;
-    await session.withTransaction(async () => {
+    await this.withTransaction(async () => {
       // Reset on transaction retries.
       activated = false;
       needsFutureActivationCheck = true;
@@ -545,7 +552,7 @@ export class MongoBucketBatchV1 extends MongoBucketBatch {
       } else if (doc?.state != storage.SyncRuleState.PROCESSING) {
         needsFutureActivationCheck = false;
       }
-    });
+    }, session);
     if (activated) {
       this.logger.info(`Activated new replication stream at ${lsn}`);
       await this.db.notifyCheckpoint();

package/src/storage/implementation/v1/MongoCompactorV1.ts

@@ -13,6 +13,7 @@ import { BucketDefinitionId } from '@powersync/service-sync-rules';
 import { BucketDataDoc } from '../common/BucketDataDoc.js';
 import { BucketStateDocumentBase, LEGACY_BUCKET_DATA_DEFINITION_ID } from '../models.js';
 import { MongoCompactOptions, MongoCompactor } from '../MongoCompactor.js';
+import { MongoWriteBatch } from '../MongoWriteBatch.js';
 import { cacheKey } from '../OperationBatch.js';
 import { BucketDataDocumentV1, BucketStateDocumentV1 } from './models.js';
 import type { MongoSyncBucketStorageV1 } from './MongoSyncBucketStorageV1.js';
@@ -296,11 +297,10 @@ export class MongoCompactorV1 extends MongoCompactor {
   }
 
   private async flushBucketStateUpdates() {
-    if (this.bucketStateUpdates.length > 0) {
-      this.logger.info(`Updating ${this.bucketStateUpdates.length} bucket states`);
-      await this.writeBucketStateUpdates();
-      this.bucketStateUpdates = [];
-    }
+    const writes = this.db.createWriteBatch(undefined, { ordered: false });
+    this.writeBucketStateUpdates(writes);
+    await writes.execute();
+    this.bucketStateUpdates = [];
   }
 
   private async updateChecksumsBatch(buckets: Pick<DirtyBucket, 'bucket' | 'definitionId'>[]) {
@@ -340,10 +340,12 @@ export class MongoCompactorV1 extends MongoCompactor {
     await this.flushBucketStateUpdates();
   }
 
-  protected async writeBucketStateUpdates(): Promise<void> {
-    await this.db.bucketStateV1.bulkWrite(
-      this.bucketStateUpdates as mongo.AnyBulkWriteOperation<BucketStateDocumentV1>[],
-      { ordered: false }
+  protected writeBucketStateUpdates(writes: MongoWriteBatch): void {
+    if (this.bucketStateUpdates.length === 0) return;
+    this.logger.info(`Updating ${this.bucketStateUpdates.length} bucket states`);
+    writes.bulkWriteUnordered(
+      this.db.bucketStateV1,
+      this.bucketStateUpdates as mongo.AnyBulkWriteOperation<BucketStateDocumentV1>[]
     );
   }
 
@@ -541,17 +543,19 @@ export class MongoCompactorV1 extends MongoCompactor {
   }
 
   private async flush(bucketContext: SingleBucketStoreV1) {
+    // Persist data changes before advertising the resulting bucket state.
+    const writes = this.db.createWriteBatch(undefined, { ordered: true });
     if (this.updates.length > 0) {
       this.logger.info(`Compacting ${this.updates.length} ops`);
-      await bucketContext.collection.bulkWrite(this.updates, {
-        // Order is not important. Since checksums are not affected, these operations can happen in any order,
-        // and it's fine if the operations are partially applied. Each individual operation is atomic.
-        ordered: false
-      });
-      this.updates = [];
+      // Order is not important. Since checksums are not affected, these operations can happen in any order,
+      // and it's fine if the operations are partially applied. Each individual operation is atomic.
+      writes.bulkWriteUnordered(bucketContext.collection, this.updates);
     }
 
-    await this.flushBucketStateUpdates();
+    this.writeBucketStateUpdates(writes);
+    await writes.execute();
+    this.updates = [];
+    this.bucketStateUpdates = [];
   }
 
   /**
@@ -622,15 +626,13 @@ export class MongoCompactorV1 extends MongoCompactor {
             }
 
             this.logger.info(`Flushing CLEAR for ${numberOfOpsToClear} ops at ${lastOp?.o}`);
-            await bucketContext.collection.deleteMany(
-              {
-                _id: {
-                  $gte: bucketContext.minId,
-                  $lte: bucketContext.docId(lastOp!.o)
-                }
-              },
-              { session }
-            );
+            const writes = this.db.createWriteBatch(session, { ordered: true });
+            writes.deleteMany(bucketContext.collection, {
+              _id: {
+                $gte: bucketContext.minId,
+                $lte: bucketContext.docId(lastOp!.o)
+              }
+            });
 
             const op = bucketContext.toPersistedDocument({
               o: lastOp!.o,
@@ -639,7 +641,8 @@ export class MongoCompactorV1 extends MongoCompactor {
               data: null,
               target_op: targetOp
             });
-            await bucketContext.collection.insertOne(op, { session });
+            writes.insertOne(bucketContext.collection, op);
+            await writes.execute();
 
             opCountDiff = -numberOfOpsToClear + 1;
           },

package/src/storage/implementation/v1/MongoSyncBucketStorageV1.ts

@@ -90,7 +90,7 @@ export class MongoSyncBucketStorageV1 extends MongoSyncBucketStorage {
     if (keepaliveOp == null && lastCheckpoint == null) {
       return null;
     }
-    return (keepaliveOp ?? 0n) > (lastCheckpoint ?? 0n) ? keepaliveOp : lastCheckpoint;
+    return [keepaliveOp ?? 0n, lastCheckpoint ?? 0n].reduce((a, b) => (a > b ? a : b));
   }
 
   protected async createWriterImpl(options: storage.CreateWriterOptions): Promise<storage.BucketStorageBatch> {

package/src/storage/implementation/v1/PersistedBatchV1.ts

@@ -6,6 +6,7 @@ import * as bson from 'bson';
 import { BucketDataSource, BucketDefinitionId } from '@powersync/service-sync-rules';
 import { mongoTableId } from '../../../utils/util.js';
 import { EMPTY_DATA } from '../MongoBucketBatchShared.js';
+import { MongoWriteBatch } from '../MongoWriteBatch.js';
 import {
   BucketStateUpdate,
   PersistedBatch,
@@ -25,7 +26,9 @@ import {
 export class PersistedBatchV1 extends PersistedBatch {
   declare protected readonly db: VersionedPowerSyncMongoV1;
 
-  currentData: mongo.AnyBulkWriteOperation<CurrentDataDocument>[] = [];
+  // Each upsert supplies the complete source-record state. Only its final
+  // state matters within a flush; bucket and parameter history is kept separately.
+  currentData = new Map<string, mongo.AnyBulkWriteOperation<CurrentDataDocument>>();
 
   protected checkDefinitionId(_definitionId: BucketDefinitionId | null): BucketDefinitionId {
     // V1 storage doesn't persist the id, and we don't use it.
@@ -94,7 +97,7 @@ export class PersistedBatchV1 extends PersistedBatch {
   }
 
   hardDeleteCurrentData(sourceTableId: bson.ObjectId, replicaId: storage.ReplicaId) {
-    this.currentData.push({
+    this.currentData.set(this.currentDataKey(sourceTableId, replicaId), {
       deleteOne: {
         filter: { _id: this.currentDataId(sourceTableId, replicaId) }
       }
@@ -124,7 +127,7 @@ export class PersistedBatchV1 extends PersistedBatch {
       return lookup.lookup;
     });
 
-    this.currentData.push({
+    this.currentData.set(this.currentDataKey(values.sourceTableId, values.replicaId), {
       updateOne: {
         filter: { _id: this.currentDataId(values.sourceTableId, values.replicaId) },
         update: {
@@ -141,57 +144,45 @@ export class PersistedBatchV1 extends PersistedBatch {
   }
 
   protected get currentDataCount() {
-    return this.currentData.length;
+    return this.currentData.size;
   }
 
-  protected async flushBucketData(session: mongo.ClientSession) {
-    await this.db.bucketDataV1.bulkWrite(
+  protected async queueBucketData(writes: MongoWriteBatch) {
+    writes.bulkWriteUnordered(
+      this.db.bucketDataV1,
       this.bucketData.map((document) => ({
         insertOne: {
           document: serializeBucketDataV1(document)
         }
-      })),
-      {
-        session,
-        ordered: false
-      }
+      }))
     );
   }
 
-  protected async flushBucketParameters(session: mongo.ClientSession) {
-    await this.db.parameterIndexV1.bulkWrite(
+  protected queueBucketParameters(writes: MongoWriteBatch): void {
+    writes.bulkWriteUnordered(
+      this.db.parameterIndexV1,
       this.bucketParameters.map((document) => ({
         insertOne: {
           document: taggedBucketParameterDocumentToV1(document)
         }
-      })),
-      {
-        session,
-        ordered: false
-      }
+      }))
     );
   }
 
-  protected async flushCurrentData(session: mongo.ClientSession) {
-    if (this.currentData.length == 0) {
+  protected queueCurrentData(writes: MongoWriteBatch): void {
+    if (this.currentData.size == 0) {
       return;
     }
 
-    await this.db.sourceRecordsV1.bulkWrite(this.currentData, {
-      session,
-      ordered: true
-    });
+    writes.bulkWriteUnordered(this.db.sourceRecordsV1, [...this.currentData.values()]);
   }
 
-  protected async flushBucketStates(session: mongo.ClientSession) {
-    await this.db.bucketStateV1.bulkWrite(this.getBucketStateUpdates(), {
-      session,
-      ordered: false
-    });
+  protected queueBucketStates(writes: MongoWriteBatch): void {
+    writes.bulkWriteUnordered(this.db.bucketStateV1, this.getBucketStateUpdates());
   }
 
   protected resetCurrentData() {
-    this.currentData = [];
+    this.currentData.clear();
   }
 
   private getBucketStateUpdates(): mongo.AnyBulkWriteOperation<BucketStateDocumentV1>[] {
@@ -219,6 +210,10 @@ export class PersistedBatchV1 extends PersistedBatch {
     });
   }
 
+  private currentDataKey(sourceTableId: bson.ObjectId, replicaId: storage.ReplicaId): string {
+    return Buffer.from(bson.serialize({ t: sourceTableId, k: replicaId })).toString('base64');
+  }
+
   private currentDataId(sourceTableId: bson.ObjectId, replicaId: storage.ReplicaId): SourceKey {
     return {
       g: this.group_id,