@powersync/service-module-mongodb-storage 0.23.2 → 0.24.0

package/src/storage/implementation/MongoBucketBatch.ts

@@ -39,6 +39,7 @@ import { MAX_ROW_SIZE } from './MongoBucketBatchShared.js';
 import { MongoIdSequence, OpIdRangeExhausted } from './MongoIdSequence.js';
 import { MongoOpIdAllocator } from './MongoOpIdAllocator.js';
 import { MongoParsedSyncConfigSet } from './MongoParsedSyncConfigSet.js';
+import { REPLICATION_WRITE_OPTIONS } from './MongoReplicationOptions.js';
 import { MongoSyncRulesLock } from './MongoSyncRulesLock.js';
 import { MongoWriteBatch } from './MongoWriteBatch.js';
 import { OperationBatch, RecordOperation } from './OperationBatch.js';
@@ -48,7 +49,6 @@ import { createObjectStorageUsageWriterId } from './v3/object-storage/ObjectStor
 export interface MongoBucketBatchOptions {
   replicationLock: MongoSyncRulesLock;
   opIdAllocator: MongoOpIdAllocator;
-  db: VersionedPowerSyncMongo;
   /**
    * The parsed sync config set for this batch.
    *
@@ -94,8 +94,8 @@ export abstract class MongoBucketBatch
   protected readonly options: MongoBucketBatchOptions;
   protected logger: Logger;
 
-  private readonly client: mongo.MongoClient;
   readonly db: VersionedPowerSyncMongo;
+  /** Shared by this lease's writers; only use inside a serialized public operation. */
   public readonly session: mongo.ClientSession;
   protected readonly sync_rules: HydratedSyncConfig;
 
@@ -111,16 +111,16 @@ export abstract class MongoBucketBatch
    * via SourceTable.storeCurrentData. The effective per-record value is the conjunction of
    * the two.
    */
-  private readonly storeCurrentData: boolean;
+  protected readonly storeCurrentData: boolean;
   public readonly skipExistingRows: boolean;
   protected readonly mapping: BucketDefinitionMapping;
   protected readonly objectStorageUsageWriterId = createObjectStorageUsageWriterId();
 
-  private batch: OperationBatch | null = null;
+  protected batch: OperationBatch | null = null;
   protected write_checkpoint_batch: storage.CustomWriteCheckpointOptions[] = [];
-  private markRecordUnavailable: BucketStorageMarkRecordUnavailable | undefined;
-  private hooks: storage.StorageHooks | undefined;
-  private clearedError = false;
+  protected readonly markRecordUnavailable: BucketStorageMarkRecordUnavailable | undefined;
+  protected hooks: storage.StorageHooks | undefined;
+  protected clearedError = false;
 
   protected tracer: PerformanceTracer<'storage' | 'evaluate'>;
 
@@ -151,11 +151,12 @@ export abstract class MongoBucketBatch
           ? options.replicationLock.signal
           : AbortSignal.any([options.signal, options.replicationLock.signal])
     };
-    this.client = options.db.client;
-    this.db = options.db;
+    // A MongoDB session can only be used with its originating client. Writers
+    // created through another factory must still use this lease's database.
+    this.db = options.replicationLock.db;
     this.replicationStreamId = options.replicationStreamId;
     this.resumeFromLsn = options.resumeFromLsn;
-    this.session = this.client.startSession();
+    this.session = options.replicationLock.session;
     this.replicationStreamName = options.replicationStreamName;
     this.sync_rules = options.parsedSyncConfig.hydratedSyncConfig;
     this.storeCurrentData = options.storeCurrentData;
@@ -175,7 +176,11 @@ export abstract class MongoBucketBatch
     });
   }
 
-  abstract resolveTables(options: storage.ResolveTablesOptions): Promise<storage.ResolveTablesResult>;
+  async resolveTables(options: storage.ResolveTablesOptions): Promise<storage.ResolveTablesResult> {
+    return this.withWriterLock(() => this.resolveTablesInternal(options));
+  }
+
+  protected abstract resolveTablesInternal(options: storage.ResolveTablesOptions): Promise<storage.ResolveTablesResult>;
 
   protected abstract createPersistedBatch(writtenSize: number): PersistedBatch;
 
@@ -183,31 +188,73 @@ export abstract class MongoBucketBatch
 
   protected abstract cleanupDroppedSourceTables(sourceTables: storage.SourceTable[]): Promise<void>;
 
-  abstract commit(lsn: string, options?: storage.BucketBatchCommitOptions): Promise<storage.CheckpointResult>;
+  async commit(lsn: string, options?: storage.BucketBatchCommitOptions): Promise<storage.CheckpointResult> {
+    return this.withWriterLock(() => this.commitInternal(lsn, options));
+  }
+
+  protected abstract commitInternal(
+    lsn: string,
+    options?: storage.BucketBatchCommitOptions
+  ): Promise<storage.CheckpointResult>;
+
+  async keepalive(lsn: string): Promise<storage.CheckpointResult> {
+    return this.commit(lsn, { createEmptyCheckpoints: true });
+  }
+
+  async setResumeLsn(lsn: string): Promise<void> {
+    return this.withWriterLock(() => this.setResumeLsnInternal(lsn));
+  }
+
+  protected abstract setResumeLsnInternal(lsn: string): Promise<void>;
 
-  abstract keepalive(lsn: string): Promise<storage.CheckpointResult>;
+  async getSourceTableStatus(table: storage.SourceTable): Promise<storage.SourceTable | null> {
+    return this.withWriterLock(() => this.getSourceTableStatusInternal(table));
+  }
 
-  abstract setResumeLsn(lsn: string): Promise<void>;
+  protected abstract getSourceTableStatusInternal(table: storage.SourceTable): Promise<storage.SourceTable | null>;
 
-  abstract getSourceTableStatus(table: storage.SourceTable): Promise<storage.SourceTable | null>;
+  async markAllSnapshotDone(no_checkpoint_before_lsn: string): Promise<void> {
+    return this.withWriterLock(() => this.markAllSnapshotDoneInternal(no_checkpoint_before_lsn));
+  }
 
-  abstract markAllSnapshotDone(no_checkpoint_before_lsn: string): Promise<void>;
+  protected abstract markAllSnapshotDoneInternal(no_checkpoint_before_lsn: string): Promise<void>;
 
-  abstract markSnapshotDone(no_checkpoint_before_lsn: string, options?: { throwOnConflict?: boolean }): Promise<void>;
+  async markSnapshotDone(no_checkpoint_before_lsn: string, options?: { throwOnConflict?: boolean }): Promise<void> {
+    return this.withWriterLock(() => this.markSnapshotDoneInternal(no_checkpoint_before_lsn, options));
+  }
 
-  abstract markTableSnapshotRequired(table: storage.SourceTable): Promise<void>;
+  protected abstract markSnapshotDoneInternal(
+    no_checkpoint_before_lsn: string,
+    options?: { throwOnConflict?: boolean }
+  ): Promise<void>;
 
-  abstract markTableSnapshotDone(
+  async markTableSnapshotRequired(table: storage.SourceTable): Promise<void> {
+    return this.withWriterLock(() => this.markTableSnapshotRequiredInternal(table));
+  }
+
+  protected abstract markTableSnapshotRequiredInternal(table: storage.SourceTable): Promise<void>;
+
+  async markTableSnapshotDone(
     tables: storage.SourceTable[],
     no_checkpoint_before_lsn?: string
-  ): Promise<storage.SourceTable[]>;
+  ): Promise<storage.SourceTable[]> {
+    return this.withWriterLock(() => this.markTableSnapshotDoneInternal(tables, no_checkpoint_before_lsn));
+  }
 
-  protected abstract batchCreateCustomWriteCheckpoints(session: mongo.ClientSession, opId: InternalOpId): Promise<void>;
+  protected abstract markTableSnapshotDoneInternal(
+    tables: storage.SourceTable[],
+    no_checkpoint_before_lsn?: string
+  ): Promise<storage.SourceTable[]>;
 
   /** Perform any version-specific setup required before writing custom checkpoints. */
   protected async prepareCustomWriteCheckpoints(): Promise<void> {}
 
   async flush(options?: storage.BatchBucketFlushOptions): Promise<storage.FlushedResult | null> {
+    return this.withWriterLock(() => this.flushInternal(options));
+  }
+
+  /** Internal helpers require the public operation to hold the writer mutex. */
+  protected async flushInternal(options?: storage.BatchBucketFlushOptions): Promise<storage.FlushedResult | null> {
     let result: storage.FlushedResult | null = null;
     // One flush may be split over multiple transactions.
     // Each flushInner() is one transaction.
@@ -222,7 +269,7 @@ export abstract class MongoBucketBatch
 
   /** Publish the checkpoint in the final flush transaction, or use the empty-commit path. */
   protected async flushAndCommit<T>(
-    checkpoint: (stream: SyncRuleDocumentBase, lastOp: InternalOpId) => Promise<T>,
+    checkpoint: (stream: SyncRuleDocumentBase, lastOp: InternalOpId, writes?: MongoWriteBatch) => Promise<T>,
     commitWithoutFlush: () => Promise<T>,
     options?: storage.BatchBucketFlushOptions,
     projection?: mongo.Document
@@ -234,8 +281,8 @@ export abstract class MongoBucketBatch
     while (this.batch != null || this.write_checkpoint_batch.length > 0) {
       await this.flushInner(
         options,
-        async (stream, lastOp) => {
-          result = await checkpoint(stream, lastOp);
+        async (stream, lastOp, writes) => {
+          result = await checkpoint(stream, lastOp, writes);
         },
         projection
       );
@@ -243,200 +290,26 @@ export abstract class MongoBucketBatch
     return result;
   }
 
-  private async flushInner(
+  protected abstract flushInner(
     options?: storage.BatchBucketFlushOptions,
-    checkpoint?: (stream: SyncRuleDocumentBase, lastOp: InternalOpId) => Promise<void>,
+    checkpoint?: (stream: SyncRuleDocumentBase, lastOp: InternalOpId, writes?: MongoWriteBatch) => Promise<void>,
     projection?: mongo.Document
-  ): Promise<storage.FlushedResult | null> {
-    const batch = this.batch;
-    let last_op: InternalOpId | null = null;
-    let resumeBatch: OperationBatch | null = null;
-    let clearedError = false;
-
-    using _ = this.tracer.span('storage', 'flush');
-
-    await this.hooks?.beforeBatchFlush?.(this);
-    if (this.write_checkpoint_batch.length > 0) {
-      // Collection/index creation cannot run inside the replication transaction.
-      await this.prepareCustomWriteCheckpoints();
-    }
-
-    await this.withReplicationTransaction(
-      `Flushing ${batch?.length ?? 0} ops`,
-      async (session, opSeq) => {
-        clearedError = false;
-        if (batch != null) {
-          const result = await this.replicateBatch(session, batch, opSeq, options);
-          resumeBatch = result.resumeBatch;
-          clearedError ||= result.clearedError;
-        }
-
-        if (this.write_checkpoint_batch.length > 0) {
-          this.logger.info(`Writing ${this.write_checkpoint_batch.length} custom write checkpoints`);
-          await this.batchCreateCustomWriteCheckpoints(session, opSeq.next());
-        }
-
-        last_op = opSeq.last();
-      },
-      async (stream, lastOp) => {
-        if (checkpoint != null && resumeBatch == null) {
-          // The checkpoint must also persist the head, including when visibility is blocked.
-          await checkpoint(stream, lastOp);
-          return true;
-        }
-        return false;
-      },
-      projection
-    );
-
-    // Keep checkpoints available if the transaction retries after writing them.
-    this.write_checkpoint_batch = [];
-
-    if (clearedError) {
-      this.clearedError = true;
-    }
-
-    // null if done, set if we need another flush
-    this.batch = resumeBatch;
-
-    if (last_op == null) {
-      throw new ReplicationAssertionError('Unexpected last_op == null');
-    }
-
-    this.recordPersistedOp(last_op);
-    this.last_flushed_op = last_op;
-    await this.hooks?.afterBatchFlush?.(this);
-    return { flushed_op: last_op };
-  }
-
-  private async replicateBatch(
-    session: mongo.ClientSession,
-    batch: OperationBatch,
-    op_seq: MongoIdSequence,
-    options?: storage.BucketBatchCommitOptions
-  ): Promise<{ resumeBatch: OperationBatch | null; clearedError: boolean }> {
-    let sizes: Map<string, number> | undefined = undefined;
-    using _ = this.tracer.span('storage', 'replicate_batch');
-    // Only look up current_data sizes if the batch stores current_data and at least one
-    // table in it does too (per-table can disable it, e.g. Postgres REPLICA IDENTITY FULL).
-    const anyTableStoresCurrentData =
-      this.storeCurrentData && batch.batch.some((r) => r.record.sourceTable.storeCurrentData);
-    if (anyTableStoresCurrentData && !this.skipExistingRows) {
-      // We skip this step if no tables store current_data, since the sizes will
-      // always be small in that case.
-
-      // With skipExistingRows, we don't load the full documents into memory,
-      // so we can also skip the size lookup step.
-
-      // Find sizes of current_data documents, to assist in intelligent batching without
-      // exceeding memory limits.
-      //
-      // A previous attempt tried to do batching by the results of the current_data query
-      // (automatically limited to 48MB(?) per batch by MongoDB). The issue is that it changes
-      // the order of processing, which then becomes really tricky to manage.
-      // This now takes 2+ queries, but doesn't have any issues with order of operations.
-      // Within this branch this.storeCurrentData is true, so the per-table flag is the
-      // effective value - only look up sizes for tables that actually store current_data.
-      const sizeLookups = batch.batch
-        .filter((r) => r.record.sourceTable.storeCurrentData)
-        .map((r) => ({
-          sourceTableId: mongoTableId(r.record.sourceTable.id),
-          replicaId: r.beforeId
-        }));
-
-      sizes = await this.sourceRecordStore.loadSizes(session, sizeLookups);
-    }
-
-    // If set, we need to start a new transaction with this batch.
-    let resumeBatch: OperationBatch | null = null;
-    let transactionSize = 0;
+  ): Promise<storage.FlushedResult | null>;
 
-    let didFlush = false;
-
-    // Now batch according to the sizes
-    // This is a single batch if storeCurrentData == false
-    for await (let b of batch.batched(sizes)) {
-      if (resumeBatch) {
-        for (let op of b) {
-          resumeBatch.push(op);
-        }
-        continue;
-      }
-      using lookupSpan = this.tracer.span('storage', 'lookup');
-      const lookups = b.map((r) => ({
-        sourceTableId: mongoTableId(r.record.sourceTable.id),
-        replicaId: r.beforeId
-      }));
-      let sourceRecordLookup = await this.sourceRecordStore.loadDocuments(session, lookups, this.skipExistingRows);
-      lookupSpan.end();
-
-      let persistedBatch: PersistedBatch | null = this.createPersistedBatch(transactionSize);
-
-      // The current code structure makes it tricky to cleanly split this span from the one
-      // where fluhsing. So we manually end and re-create this span whenever we flush.
-      let evalSpan = this.tracer.span('evaluate');
-      for (let op of b) {
-        if (resumeBatch) {
-          resumeBatch.push(op);
-          continue;
-        }
-        const sourceRecord = sourceRecordLookup.get(op.internalBeforeKey) ?? null;
-        if (sourceRecord != null) {
-          // If it will be used again later, it will be set again using nextData below
-          sourceRecordLookup.delete(op.internalBeforeKey);
-        }
-        const nextData = this.saveOperation(persistedBatch!, op, sourceRecord, op_seq);
-        if (nextData != null) {
-          // Update our current_data and size cache
-          sourceRecordLookup.set(op.internalAfterKey!, nextData);
-          sizes?.set(op.internalAfterKey!, nextData.data?.length() ?? 0);
-        }
-
-        if (persistedBatch!.shouldFlushTransaction()) {
-          evalSpan.end();
-          // Transaction is getting big.
-          // Flush, and resume in a new transaction.
-          using persistSpan = this.tracer.span('storage', 'persist_flush');
-          const { flushedAny } = await persistedBatch!.flush(this.session, options);
-
-          didFlush ||= flushedAny;
-          persistedBatch = null;
-          // Computing our current progress is a little tricky here, since
-          // we're stopping in the middle of a batch.
-          // We create a new batch, and push any remaining operations to it.
-          resumeBatch = new OperationBatch();
-          persistSpan.end();
-          evalSpan = this.tracer.span('evaluate');
-        }
-      }
-      evalSpan.end();
-
-      if (persistedBatch) {
-        transactionSize = persistedBatch.currentSize;
-        using _ = this.tracer.span('storage', 'persist_flush');
-        const { flushedAny } = await persistedBatch.flush(this.session, options);
-        didFlush ||= flushedAny;
-      }
-    }
-
-    const clearedError = didFlush && !this.clearedError;
-    if (clearedError) {
-      // No need to clear an error more than once per batch, since an error would always result in restarting the batch.
-      using _ = this.tracer.span('storage', 'clear_error');
-      await this.clearError(this.session);
-    }
-
-    return {
-      resumeBatch: resumeBatch?.hasData() ? resumeBatch : null,
-      clearedError
-    };
-  }
-
-  private saveOperation(
+  /**
+   * Evaluate one source operation and append its storage writes to batch.
+   *
+   * Diagnostics run during evaluation. V1 calls markRecordUnavailable immediately;
+   * v3/v4 supply a collector and deliver repair notifications only after a
+   * confirmed publication. A failed row attempt may have registered a repair,
+   * so its caller must discard that row's collected notification on failure.
+   */
+  protected saveOperation(
     batch: PersistedBatch,
     operation: RecordOperation,
     sourceRecord: LoadedSourceRecord | null,
-    opSeq: MongoIdSequence
+    opSeq: MongoIdSequence,
+    onRecordUnavailable: BucketStorageMarkRecordUnavailable | undefined
   ) {
     const record = operation.record;
     const beforeId = operation.beforeId;
@@ -474,19 +347,6 @@ export abstract class MongoBucketBatch
         // Not an error if we re-apply a transaction
         existing_buckets = [];
         existing_lookups = [];
-        if (!isCompleteRow(storeCurrentData, after!)) {
-          if (this.markRecordUnavailable != null) {
-            // This will trigger a "resnapshot" of the record.
-            // This is not relevant if storeCurrentData is false, since we'll get the full row
-            // directly in the replication stream.
-            this.markRecordUnavailable(record);
-          } else {
-            // Log to help with debugging if there was a consistency issue
-            this.logger.warn(
-              `Cannot find previous record for update on ${record.sourceTable.qualifiedName}: ${beforeId} / ${record.before?.id}`
-            );
-          }
-        }
       } else {
         existing_buckets = result.buckets;
         existing_lookups = result.lookups;
@@ -495,6 +355,24 @@ export abstract class MongoBucketBatch
           after = storage.mergeToast<SqliteValue>(after!, data);
         }
       }
+      if (!isCompleteRow(storeCurrentData, after!)) {
+        if (onRecordUnavailable != null) {
+          // This will trigger a "resnapshot" of the record.
+          // This is not relevant if storeCurrentData is false, since we'll get the full row
+          // directly in the replication stream.
+          // Check after merging, including existing incomplete rows/tombstones.
+          // A flush can commit before a process dies or notification delivery
+          // fails. Its source position has not advanced, so WAL replay must
+          // request repair again rather than treating mere record existence as
+          // proof that the missing TOAST values have been recovered.
+          onRecordUnavailable(record);
+        } else if (result == null) {
+          // Log to help with debugging if there was a consistency issue
+          this.logger.warn(
+            `Cannot find previous record for update on ${record.sourceTable.qualifiedName}: ${beforeId} / ${record.before?.id}`
+          );
+        }
+      }
     } else if (record.tag == SaveOperationTag.DELETE) {
       const result = sourceRecord;
       if (result == null) {
@@ -504,7 +382,7 @@ export abstract class MongoBucketBatch
         // Log to help with debugging if there was a consistency issue.
         // Gate on the batch-level flag: FULL tables (per-record flag false) still get a
         // current_data entry, so a missing record on DELETE is meaningful for them too.
-        if (this.storeCurrentData && this.markRecordUnavailable == null) {
+        if (this.storeCurrentData && onRecordUnavailable == null) {
           this.logger.warn(
             `Cannot find previous record for delete on ${record.sourceTable.qualifiedName}: ${beforeId} / ${record.before?.id}`
           );
@@ -707,42 +585,73 @@ export abstract class MongoBucketBatch
     return this.withFencedTransaction(() => this.fence(session), cb, session);
   }
 
-  /** The first operation must modify the stream document conditional on lease ownership. */
+  /**
+   * Internal transaction helper: the public writer operation already holds the
+   * mutex. The first operation must modify the stream document conditional on
+   * lease ownership. This database fence is required on every retry even though
+   * local writers are serialized; another process may have taken over the lease.
+   */
   protected async withFencedTransaction<S, T>(
     acquireFence: () => Promise<S>,
     cb: (stream: S) => Promise<T>,
-    session = this.session
+    session = this.session,
+    description?: string
   ): Promise<T> {
-    return this.withWriterLock(() =>
-      session.withTransaction(
-        async () => {
+    this.options.replicationLock.throwIfAborted();
+    let attempts = 0;
+    // Includes driver commit retries and the delays between transaction attempts.
+    using _transactionSpan = this.tracer.span('storage', 'transaction');
+    const result = await session.withTransaction(
+      async () => {
+        attempts++;
+        // One callback invocation per transaction attempt. Commit happens after
+        // this span; command monitoring is needed to measure the full wire lifetime.
+        using attemptSpan = this.tracer.span('storage', 'transaction_body');
+        try {
           const stream = await acquireFence();
-          try {
-            const result = await cb(stream);
-            this.options.replicationLock.throwIfAborted();
-            return result;
-          } catch (e: unknown) {
-            if (e instanceof OpIdRangeExhausted) {
-              throw e;
-            }
-            if (e instanceof mongo.MongoError && e.hasErrorLabel('TransientTransactionError')) {
-              // Likely write conflict caused by concurrent writes to this replication stream.
-            } else {
-              this.logger.warn('Transaction error', e as Error);
-            }
-            const delay = Math.random() * 50;
-            using _ = this.tracer.span('storage', 'retry_delay');
-            await timers.setTimeout(delay);
+          const result = await cb(stream);
+          this.options.replicationLock.throwIfAborted();
+          return result;
+        } catch (e: unknown) {
+          attemptSpan.end();
+          if (e instanceof OpIdRangeExhausted) {
             throw e;
           }
-        },
-        { maxCommitTimeMS: 10000, writeConcern: { w: 'majority' } }
-      )
+          if (e instanceof mongo.MongoError && e.hasErrorLabel('TransientTransactionError')) {
+            // Log every tenth attempt, including failures of the fence, without
+            // logging a stack or bulk result for every driver retry.
+            if ((attempts + 1) % 10 === 0) {
+              this.logger.info(`${description ?? 'Replication transaction'} - try ${attempts + 1}`, {
+                error: e.message,
+                code: e.code,
+                codeName: e instanceof mongo.MongoServerError ? e.codeName : undefined,
+                errorLabels: e.errorLabels
+              });
+            }
+          } else {
+            this.logger.warn('Transaction error', e as Error);
+          }
+          const delay = Math.random() * 50;
+          using _ = this.tracer.span('storage', 'retry_delay');
+          await timers.setTimeout(delay);
+          throw e;
+        }
+      },
+      { maxCommitTimeMS: 10000, ...REPLICATION_WRITE_OPTIONS }
     );
+    return result;
   }
 
+  /**
+   * Public database operations acquire the lease's mutex once, then call only
+   * internal helpers until they finish. Those helpers must never call a public
+   * database operation or reacquire this mutex. All session use, including reads,
+   * belongs inside this boundary so it cannot join another writer's transaction.
+   * Transaction helpers still fence database ownership on every attempt.
+   */
   private async withWriterLock<T>(callback: () => Promise<T>): Promise<T> {
     using lockSpan = this.tracer.span('storage', 'internal_lock');
+    this.options.replicationLock.throwIfAborted();
     return this.options.replicationLock.writerMutex.exclusiveLock(async () => {
       lockSpan.end();
       this.options.replicationLock.throwIfAborted();
@@ -750,13 +659,13 @@ export abstract class MongoBucketBatch
     });
   }
 
-  /** Single-document updates enforce ownership atomically, without a separate transaction. */
+  /** Internal single-document update; the public operation already holds the mutex. */
   protected async updateStreamMetadata(
     set: mongo.Document,
     filter: mongo.Document = {},
-    writeConcern: mongo.WriteConcernSettings = { w: 'majority' }
+    writeConcern: mongo.WriteConcernSettings = REPLICATION_WRITE_OPTIONS.writeConcern
   ): Promise<void> {
-    await this.withWriterLock(() => this.writeStreamMetadata(set, filter, writeConcern));
+    await this.writeStreamMetadata(set, filter, writeConcern);
   }
 
   /** Caller must already be inside a fenced transaction holding the writer mutex. */
@@ -775,63 +684,7 @@ export abstract class MongoBucketBatch
       { session: this.session, ...(writeConcern == null ? {} : { writeConcern }) }
     );
     if (result.matchedCount === 0) {
-      throw new ReplicationAbortedError('Replication writer no longer owns the stream');
-    }
-  }
-
-  private async withReplicationTransaction(
-    description: string,
-    callback: (session: mongo.ClientSession, opSeq: MongoIdSequence) => Promise<void>,
-    finish?: (stream: SyncRuleDocumentBase, lastOp: InternalOpId) => Promise<boolean>,
-    projection?: mongo.Document
-  ): Promise<void> {
-    let flushTry = 0;
-    const lastTry = Date.now() + 90000;
-    const allocator = this.options.opIdAllocator;
-    let lastOp = 0n;
-    // Refill outside the publication transaction, before evaluating any rows.
-    // Exhaustion remains a fallback for large batches or a newer stream head.
-    this.options.replicationLock.throwIfAborted();
-    await allocator.ensureCapacity();
-    for (;;) {
-      try {
-        await this.withFencedTransaction(
-          () => this.fence(this.session, projection),
-          async (stream) => {
-            flushTry += 1;
-            if (flushTry % 10 == 0) {
-              this.logger.info(`${description} - try ${flushTry}`);
-            }
-            if (flushTry > 20 && Date.now() > lastTry) {
-              throw new ServiceError(ErrorCode.PSYNC_S1402, 'Max transaction tries exceeded');
-            }
-            // The fence already holds this stream document for the transaction.
-            // A different writer may have used higher IDs since our last flush.
-            const opSeq = allocator.sequence(this.persistedOpHead(stream));
-            await callback(this.session, opSeq);
-            lastOp = opSeq.last();
-
-            if (!(await finish?.(stream, lastOp))) {
-              const writes = this.db.createWriteBatch(this.session, { ordered: false });
-              // Allow subclasses to persist additional flush-time state in the same transaction.
-              this.onReplicationTransactionFlush(writes, lastOp);
-              await writes.execute();
-            }
-            // Notifications and cleanup must wait until the transaction commits.
-          }
-        );
-        allocator.committed(lastOp);
-        return;
-      } catch (error) {
-        if (!(error instanceof OpIdRangeExhausted)) {
-          throw error;
-        }
-        // withTransaction has aborted every write. Reserve outside that transaction,
-        // then replay evaluation from the same committed stream head. Existing
-        // ranges can be reused on retry because no operation from this attempt committed.
-        this.options.replicationLock.throwIfAborted();
-        await allocator.reserve();
-      }
+      throw new ReplicationAbortedError('Replication writer lost its lease or metadata target changed');
     }
   }
 
@@ -841,21 +694,6 @@ export abstract class MongoBucketBatch
   /** Advance the persisted head in the same transaction as the operations. */
   protected abstract onReplicationTransactionFlush(writes: MongoWriteBatch, lastOp: InternalOpId): void;
 
-  /**
-   * Called after a replication transaction has successfully committed, with the last persisted op id.
-   *
-   * v1 storage tracks this in memory to fold into the next checkpoint. v3 storage does not need it:
-   * the stream-level `last_persisted_op` is already advanced durably by
-   * {@link onReplicationTransactionFlush} or the combined checkpoint update within the same
-   * transaction, and empty commits read it from the document.
-   *
-   * Keep calls to this adjacent to {@link withReplicationTransaction} usage - both must observe
-   * every path that persists ops.
-   */
-  protected recordPersistedOp(_lastOp: InternalOpId): void {
-    // No-op by default.
-  }
-
   async [Symbol.asyncDispose]() {
     if (this.batch?.hasData() || this.write_checkpoint_batch.length > 0) {
       // We don't error here, since:
@@ -863,7 +701,7 @@ export abstract class MongoBucketBatch
       // 2. SuppressedError is messy to deal with.
       this.logger.warn('Disposing writer with unflushed changes');
     }
-    await this.session.endSession();
+    // The lease owns the shared session and ends it after all writers stop.
     super.clearListeners();
   }
 
@@ -921,8 +759,11 @@ export abstract class MongoBucketBatch
    * Drop is equivalent to TRUNCATE, plus removing our record of the table.
    */
   async drop(sourceTables: storage.SourceTable[]): Promise<storage.FlushedResult | null> {
-    await this.truncate(sourceTables);
-    const result = await this.flush();
+    return this.withWriterLock(() => this.dropInternal(sourceTables));
+  }
+
+  protected async dropInternal(sourceTables: storage.SourceTable[]): Promise<storage.FlushedResult | null> {
+    await this.truncateInternal(sourceTables);
 
     await this.withTransaction(async () => {
       for (let table of sourceTables) {
@@ -933,19 +774,22 @@ export abstract class MongoBucketBatch
     });
 
     await this.cleanupDroppedSourceTables(sourceTables);
-    return result;
+    return null;
   }
 
   async truncate(sourceTables: storage.SourceTable[]): Promise<storage.FlushedResult | null> {
-    await this.flush();
+    return this.withWriterLock(() => this.truncateInternal(sourceTables));
+  }
+
+  protected async truncateInternal(sourceTables: storage.SourceTable[]): Promise<storage.FlushedResult | null> {
+    await this.flushInternal();
 
     let last_op: InternalOpId | null = null;
     for (let table of sourceTables) {
-      last_op = await this.truncateSingle(table);
+      last_op = await this.truncateSingleInternal(table);
     }
 
     if (last_op) {
-      this.recordPersistedOp(last_op);
       return {
         flushed_op: last_op
       };
@@ -955,53 +799,19 @@ export abstract class MongoBucketBatch
   }
 
   async truncateSingle(sourceTable: storage.SourceTable): Promise<InternalOpId> {
-    let last_op: InternalOpId | null = null;
-
-    // To avoid too large transactions, we limit the amount of data we delete per transaction.
-    // Since we don't use the record data here, we don't have explicit size limits per batch.
-    const BATCH_LIMIT = 2000;
-
-    let lastBatchCount = BATCH_LIMIT;
-    while (lastBatchCount == BATCH_LIMIT) {
-      await this.withReplicationTransaction(`Truncate ${sourceTable.qualifiedName}`, async (session, opSeq) => {
-        using evalSpan = this.tracer.span('evaluate');
-        const sourceTableId = mongoTableId(sourceTable.id);
-        const batch = await this.sourceRecordStore.loadTruncateBatch(session, sourceTableId, BATCH_LIMIT);
-        const persistedBatch = this.createPersistedBatch(0);
-
-        for (let value of batch) {
-          persistedBatch.saveBucketData({
-            op_seq: opSeq,
-            before_buckets: value.buckets,
-            evaluated: [],
-            table: sourceTable,
-            sourceKey: value.replicaId
-          });
-          persistedBatch.saveParameterData({
-            op_seq: opSeq,
-            existing_lookups: value.lookups,
-            evaluated: [],
-            sourceTable: sourceTable,
-            sourceKey: value.replicaId
-          });
-
-          // Since this is not from streaming replication, we can do a hard delete
-          persistedBatch.hardDeleteCurrentData(sourceTableId, value.replicaId);
-        }
-        evalSpan.end();
-
-        using _ = this.tracer.span('storage', 'persist_flush');
-        await persistedBatch.flush(session);
-        lastBatchCount = batch.length;
+    return this.withWriterLock(() => this.truncateSingleInternal(sourceTable));
+  }
 
-        last_op = opSeq.last();
-      });
-    }
+  protected abstract truncateSingleInternal(sourceTable: storage.SourceTable): Promise<InternalOpId>;
 
-    return last_op!;
+  async updateTableProgress(
+    table: storage.SourceTable,
+    progress: Partial<storage.TableSnapshotStatus>
+  ): Promise<storage.SourceTable> {
+    return this.withWriterLock(() => this.updateTableProgressInternal(table, progress));
   }
 
-  async updateTableProgress(
+  private async updateTableProgressInternal(
     table: storage.SourceTable,
     progress: Partial<storage.TableSnapshotStatus>
   ): Promise<storage.SourceTable> {
@@ -1032,23 +842,6 @@ export abstract class MongoBucketBatch
     return copy;
   }
 
-  protected async clearError(session: mongo.ClientSession): Promise<void> {
-    await this.db.sync_rules.updateOne(
-      {
-        _id: this.replicationStreamId
-      },
-      {
-        $set: {
-          last_fatal_error: null,
-          last_fatal_error_ts: null
-        }
-      },
-      {
-        session
-      }
-    );
-  }
-
   /**
    * Gets relevant {@link HydratedEventDescriptor}s for the given {@link SourceTable}
    */