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

package/src/storage/implementation/v3/MongoBucketBatchV3.ts

@@ -1,16 +1,24 @@
 import * as lib_mongo from '@powersync/lib-service-mongodb';
-import { ReplicationAssertionError, ServiceAssertionError } from '@powersync/lib-services-framework';
+import {
+  ErrorCode,
+  ReplicationAbortedError,
+  ReplicationAssertionError,
+  ServiceAssertionError,
+  ServiceError
+} from '@powersync/lib-services-framework';
 import { InternalOpId, storage } from '@powersync/service-core';
 import { EventDefinitionId } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
 import { mongoTableId } from '../../../utils/util.js';
 import { canCheckpointState } from '../CheckpointState.js';
 import { MongoBucketBatch, MongoBucketBatchOptions } from '../MongoBucketBatch.js';
+import { OpIdRangeExhausted } from '../MongoIdSequence.js';
 import { MongoParsedSyncConfigSet } from '../MongoParsedSyncConfigSet.js';
 import { MongoWriteBatch } from '../MongoWriteBatch.js';
+import { OperationBatch } from '../OperationBatch.js';
 import { stopReplicationStreamPipeline } from '../SyncRuleStateUpdate.js';
-import { PersistedBatch } from '../common/PersistedBatch.js';
-import { SourceRecordStore } from '../common/SourceRecordStore.js';
+import { PreparedPublication } from '../common/PreparedPublication.js';
+import { LoadedSourceRecord } from '../common/SourceRecordStore.js';
 import { SyncRuleDocumentBase } from '../models.js';
 import { PersistedBatchV3 } from './PersistedBatchV3.js';
 import { SourceRecordStoreV3 } from './SourceRecordStoreV3.js';
@@ -28,7 +36,7 @@ import {
 export class MongoBucketBatchV3 extends MongoBucketBatch {
   declare readonly db: VersionedPowerSyncMongoV3;
 
-  private readonly store: SourceRecordStore;
+  private readonly store: SourceRecordStoreV3;
   /** Avoid repeating idempotent createIndexes calls for every checkpoint flush. */
   private readonly initializedCustomCheckpointEventIds = new Set<EventDefinitionId>();
   // Used only to validate supplied checkpoint ids; it does not initialize collections.
@@ -48,7 +56,7 @@ export class MongoBucketBatchV3 extends MongoBucketBatch {
     this.store = new SourceRecordStoreV3(this.db, this.replicationStreamId, this.mapping);
   }
 
-  protected createPersistedBatch(writtenSize: number): PersistedBatch {
+  protected createPersistedBatch(writtenSize: number): PersistedBatchV3 {
     return new PersistedBatchV3(this.db, this.replicationStreamId, this.mapping, writtenSize, {
       logger: this.logger,
       objectStorage: this.options.objectStorage,
@@ -58,10 +66,374 @@ export class MongoBucketBatchV3 extends MongoBucketBatch {
     });
   }
 
-  protected get sourceRecordStore(): SourceRecordStore {
+  protected get sourceRecordStore(): SourceRecordStoreV3 {
     return this.store;
   }
 
+  protected override async withTransaction<T>(
+    callback: (stream: SyncRuleDocumentBase) => Promise<T>,
+    session = this.session
+  ): Promise<T> {
+    return this.withFencedTransaction(
+      () => this.fence(session, { state: 1, sync_configs: 1 }),
+      async (stream) => {
+        this.assertPreparationContext(stream as ReplicationStreamDocumentV3);
+        return callback(stream);
+      },
+      session
+    );
+  }
+
+  protected override async flushInner(
+    options?: storage.BatchBucketFlushOptions,
+    checkpoint?: (stream: SyncRuleDocumentBase, lastOp: InternalOpId, writes?: MongoWriteBatch) => Promise<void>,
+    projection?: lib_mongo.mongo.Document
+  ): Promise<storage.FlushedResult | null> {
+    using _span = this.tracer.span('storage', 'flush');
+    await this.hooks?.beforeBatchFlush?.(this);
+    const input = this.batch;
+    const checkpoints = this.write_checkpoint_batch.slice();
+    const allocator = this.options.opIdAllocator;
+    await allocator.ensureCapacity();
+    const contextProjection = { ...projection, state: 1, last_persisted_op: 1, sync_configs: 1 };
+    if (checkpoints.length > 0) {
+      await this.prepareCustomWriteCheckpoints();
+    }
+    const prepared = await this.prepareRows(this.session, input);
+    let lastOp = prepared.lastOp;
+    let checkpointOp: bigint | undefined;
+    if (checkpoints.length > 0) {
+      while (true) {
+        try {
+          checkpointOp = allocator.sequence(lastOp).next();
+          lastOp = checkpointOp;
+          break;
+        } catch (error) {
+          if (!(error instanceof OpIdRangeExhausted)) {
+            throw error;
+          }
+          await allocator.reserve();
+        }
+      }
+    }
+    // Reservations are globally disjoint and this lease's allocator consumes IDs
+    // assigned to sealed work, even if publication is abandoned. The writer mutex
+    // excludes other writers under this lease during preparation and publication;
+    // the transaction fence excludes a different lease owner. Thus preparation
+    // does not need to read the persisted head to allocate IDs or reuse records.
+    allocator.consume(lastOp);
+    const clearError = !this.clearedError && prepared.hasWrites;
+    await this.publishPrepared(lastOp, contextProjection, prepared.publications, options, async (stream, writes) => {
+      // An empty flush allocates no ID. Use the durable head from the
+      // fresh fence for its checkpoint and returned flushed operation.
+      if (lastOp === 0n) {
+        lastOp = this.persistedOpHead(stream);
+      }
+      if (clearError) {
+        writes.updateOne(
+          this.db.sync_rules,
+          { _id: this.replicationStreamId },
+          {
+            $set: { last_fatal_error: null, last_fatal_error_ts: null }
+          }
+        );
+      }
+      if (checkpointOp != null) {
+        this.queueCustomWriteCheckpoints(writes, checkpointOp, checkpoints);
+      }
+      if (checkpoint != null && prepared.resumeBatch == null) {
+        // Checkpoint decisions use this attempt's freshly fenced state.
+        await checkpoint(stream, lastOp, writes);
+      } else {
+        this.onReplicationTransactionFlush(writes, lastOp);
+      }
+    });
+    this.batch = prepared.resumeBatch;
+    this.write_checkpoint_batch.splice(0, checkpoints.length);
+    this.clearedError ||= clearError;
+    this.last_flushed_op = lastOp;
+    // markRecordUnavailable only synchronously queues connector-owned repair
+    // keys. It must not mutate input, perform async work, or re-enter storage:
+    // the writer mutex is still held. Delivery must finish before flush() returns;
+    // WalStream drains that queue before commit()/ack().
+    // Once committed, these writes cannot be rolled back if the callback
+    // throws. Invalidate the whole local lease while the mutex is held, so
+    // no writer sharing it (including queued calls) can advance progress
+    // past an undelivered repair. A fresh replication attempt must replay
+    // from the old source position and redetect incomplete rows.
+    try {
+      const notify = this.markRecordUnavailable;
+      if (notify != null) {
+        for (const record of prepared.repairs) {
+          notify(record);
+        }
+      }
+    } catch (error) {
+      this.options.replicationLock.invalidate(
+        new ReplicationAbortedError(
+          'Committed replication repair notification failed; restart replication before advancing source progress',
+          error
+        )
+      );
+      throw error;
+    }
+    await this.hooks?.afterBatchFlush?.(this);
+    return { flushed_op: lastOp };
+  }
+
+  private assertPreparationContext(stream: ReplicationStreamDocumentV3): void {
+    if ([storage.SyncRuleState.STOP, storage.SyncRuleState.TERMINATED].includes(stream.state)) {
+      throw new ReplicationAbortedError('Replication stream has stopped');
+    }
+    this.options.parsedSyncConfig.assertPublicationConfigs(stream.sync_configs);
+  }
+
+  /**
+   * Internal publisher: the public operation holds the mutex through commit.
+   * finish() must only queue writes independent of the final data group, using
+   * the fresh fence and planned IDs; it must not execute the supplied batch or
+   * read results from unpublished data. Activation and cleanup remain separate.
+   */
+  private async publishPrepared(
+    lastOp: InternalOpId,
+    projection: lib_mongo.mongo.Document,
+    publications: PreparedPublication[],
+    options: storage.BatchBucketFlushOptions | undefined,
+    finish: (stream: ReplicationStreamDocumentV3, writes: MongoWriteBatch) => Promise<void>
+  ): Promise<void> {
+    // This gate covers every sealed batch in the prefix, including truncate/drop.
+    // The transaction must never wait for source reads, evaluation or S3 PUTs.
+    // Sequential preparation also avoids leaving another group's uploads running
+    // when one group fails; each group's readiness joins all of its own uploads.
+    for (const publication of publications) {
+      await publication.preparePayloads();
+    }
+    let attempts = 0;
+    const retryDeadline = Date.now() + 90_000;
+    try {
+      await this.withFencedTransaction(
+        () => this.fence(this.session, projection),
+        async (document) => {
+          attempts++;
+          if (attempts > 20 && Date.now() > retryDeadline) {
+            throw new ServiceError(ErrorCode.PSYNC_S1402, 'Max transaction tries exceeded');
+          }
+          const stream = document as ReplicationStreamDocumentV3;
+          this.assertPreparationContext(stream);
+          const head = this.persistedOpHead(stream);
+          if (lastOp !== 0n && lastOp <= head) {
+            throw new ReplicationAssertionError(
+              `Prepared operation ${lastOp} does not advance persisted stream head ${head}`
+            );
+          }
+          // Earlier groups may replace the same source record or bucket state.
+          // Preserve their execution order; only the final group's writes are
+          // independent of the head/checkpoint tail built from this fresh fence.
+          for (const publication of publications.slice(0, -1)) {
+            using _persist = this.tracer.span('storage', 'persist_flush');
+            await publication.publish(this.session, options);
+          }
+          const writes = this.db.createWriteBatch(this.session, { ordered: false });
+          await finish(stream, writes);
+          const finalPublication = publications.at(-1);
+          if (finalPublication != null) {
+            using _persist = this.tracer.span('storage', 'persist_flush');
+            await finalPublication.publish(this.session, options, writes);
+          } else {
+            await writes.execute();
+          }
+        },
+        this.session,
+        'Publishing prepared work'
+      );
+    } catch (error) {
+      // Publication failed after the driver's retries or our attempt limit.
+      // Leave uploaded objects and markers alone: the commit may have succeeded.
+      // Stop every writer on this lease and recover from persisted source progress.
+      this.options.replicationLock.invalidate(
+        new ReplicationAbortedError('Publication failed; restart replication from persisted progress', error)
+      );
+      throw error;
+    }
+  }
+
+  private async prepareRows(session: lib_mongo.mongo.ClientSession, input: OperationBatch | null) {
+    const publications: PreparedPublication[] = [];
+    const repairs: storage.SaveUpdate[] = [];
+    let resumeBatch: OperationBatch | null = null;
+    let lastOp = 0n;
+    let hasWrites = false;
+    if (input == null) {
+      return { publications, repairs, resumeBatch, lastOp, hasWrites };
+    }
+    using _preparation = this.tracer.span('storage', 'prepare');
+    const allocator = this.options.opIdAllocator;
+    let sizes: Map<string, number> | undefined;
+    // Keep the existing size-based read partitioning, including per-table raw
+    // data storage and snapshot skip-existing exceptions. Move these reads;
+    // publication never repeats them as a validation pass.
+    if (
+      this.options.storeCurrentData &&
+      !this.skipExistingRows &&
+      input.batch.some((op) => op.record.sourceTable.storeCurrentData)
+    ) {
+      sizes = await this.sourceRecordStore.loadSizes(
+        session,
+        input.batch
+          .filter((op) => op.record.sourceTable.storeCurrentData)
+          .map((op) => ({ sourceTableId: mongoTableId(op.record.sourceTable.id), replicaId: op.beforeId }))
+      );
+    }
+    // Only this bounded publication's observations/evaluated rows are retained.
+    // No state is shared with another pending publication or another writer.
+    const records = new Map<string, LoadedSourceRecord | null>();
+    let writtenSize = 0;
+    for (const block of input.batched(sizes)) {
+      if (resumeBatch != null) {
+        for (const op of block) {
+          resumeBatch.push(op);
+        }
+        continue;
+      }
+      const missing = block.filter((op) => !records.has(op.internalBeforeKey));
+      if (missing.length > 0) {
+        using _lookup = this.tracer.span('storage', 'lookup');
+        const loaded = await this.sourceRecordStore.loadDocuments(
+          session,
+          missing.map((op) => ({
+            sourceTableId: mongoTableId(op.record.sourceTable.id),
+            replicaId: op.beforeId
+          })),
+          this.skipExistingRows
+        );
+        for (const op of missing) {
+          records.set(op.internalBeforeKey, loaded.get(op.internalBeforeKey) ?? null);
+        }
+      }
+      let group: PersistedBatchV3 | null = this.createPersistedBatch(writtenSize);
+      for (const op of block) {
+        if (resumeBatch != null) {
+          resumeBatch.push(op);
+          continue;
+        }
+        let next: LoadedSourceRecord | null;
+        while (true) {
+          const rowRepairs: storage.SaveUpdate[] = [];
+          const row = this.createPersistedBatch(0);
+          const sequence = allocator.sequence(lastOp);
+          try {
+            using _evaluate = this.tracer.span('evaluate');
+            next = this.saveOperation(
+              row,
+              op,
+              records.get(op.internalBeforeKey) ?? null,
+              sequence,
+              this.markRecordUnavailable == null
+                ? undefined
+                : (record) => {
+                    rowRepairs.push(record);
+                  }
+            );
+            group!.append(row);
+            hasWrites ||= row.currentSize > 0;
+            repairs.push(...rowRepairs);
+            lastOp = sequence.last();
+            break;
+          } catch (error) {
+            if (!(error instanceof OpIdRangeExhausted)) {
+              throw error;
+            }
+            // Retry only the unfinished row; no payload has been uploaded yet.
+            this.options.replicationLock.throwIfAborted();
+            await allocator.reserve();
+          }
+        }
+        if (op.afterId == null || !storage.replicaIdEquals(op.beforeId, op.afterId)) {
+          // Preserve evaluation's existing repeated-row behavior: a deleted
+          // identity no longer supplies raw values or membership to later rows.
+          records.set(op.internalBeforeKey, null);
+        }
+        if (next != null) {
+          records.set(op.internalAfterKey!, next);
+          sizes?.set(op.internalAfterKey!, next.data?.length() ?? 0);
+        }
+        if (group!.shouldFlushTransaction()) {
+          publications.push(group!.seal());
+          group = null;
+          resumeBatch = new OperationBatch();
+        }
+      }
+      if (group != null) {
+        writtenSize = group.currentSize;
+        publications.push(group.seal());
+      }
+    }
+    return { publications, repairs, resumeBatch: resumeBatch?.hasData() ? resumeBatch : null, lastOp, hasWrites };
+  }
+
+  protected override async truncateSingleInternal(sourceTable: storage.SourceTable): Promise<InternalOpId> {
+    const sourceId = mongoTableId(sourceTable.id);
+    const session = this.session;
+    const projection = { state: 1, sync_configs: 1, last_persisted_op: 1 };
+    const allocator = this.options.opIdAllocator;
+    while (true) {
+      await allocator.ensureCapacity();
+      // Preserve the existing bounded table scan. Hard deletion lets the next
+      // scan resume at remaining live records without retaining a table overlay.
+      const records = await this.sourceRecordStore.loadTruncateBatch(session, sourceId, 2000);
+      const group = this.createPersistedBatch(0);
+      let lastOp = 0n;
+      let processed = 0;
+      for (const record of records) {
+        while (true) {
+          const row = this.createPersistedBatch(0);
+          const sequence = allocator.sequence(lastOp);
+          try {
+            row.saveBucketData({
+              op_seq: sequence,
+              before_buckets: record.buckets,
+              evaluated: [],
+              table: sourceTable,
+              sourceKey: record.replicaId
+            });
+            row.saveParameterData({
+              op_seq: sequence,
+              existing_lookups: record.lookups,
+              evaluated: [],
+              sourceTable,
+              sourceKey: record.replicaId
+            });
+            // Since this is not from streaming replication, we can do a hard delete.
+            row.hardDeleteCurrentData(sourceId, record.replicaId);
+            group.append(row);
+            lastOp = sequence.last();
+            break;
+          } catch (error) {
+            if (!(error instanceof OpIdRangeExhausted)) {
+              throw error;
+            }
+            this.options.replicationLock.throwIfAborted();
+            await allocator.reserve();
+          }
+        }
+        processed++;
+        if (group.shouldFlushTransaction()) {
+          break;
+        }
+      }
+      allocator.consume(lastOp);
+      await this.publishPrepared(lastOp, projection, [group.seal()], undefined, async (stream, writes) => {
+        if (lastOp === 0n) {
+          lastOp = this.persistedOpHead(stream);
+        }
+        this.onReplicationTransactionFlush(writes, lastOp);
+      });
+      if (processed === records.length && records.length < 2000) {
+        return lastOp;
+      }
+    }
+  }
+
   protected override persistedOpHead(stream: SyncRuleDocumentBase): InternalOpId {
     return (stream as ReplicationStreamDocumentV3).last_persisted_op ?? 0n;
   }
@@ -151,7 +523,9 @@ export class MongoBucketBatchV3 extends MongoBucketBatch {
     }
   }
 
-  async resolveTables(options: storage.ResolveTablesOptions): Promise<storage.ResolveTablesResult> {
+  protected override async resolveTablesInternal(
+    options: storage.ResolveTablesOptions
+  ): Promise<storage.ResolveTablesResult> {
     // The test-only override is a whole parsed set, so the sync rules and the mapping
     // used below always come from the same parse.
     const parsedOverride = options.parsedSyncConfig as MongoParsedSyncConfigSet | undefined;
@@ -173,8 +547,7 @@ export class MongoBucketBatchV3 extends MongoBucketBatch {
     };
 
     let result: storage.ResolveTablesResult | null = null;
-    const session = this.db.client.startSession();
-    await using _ = { [Symbol.asyncDispose]: () => session.endSession() };
+    const session = this.session;
 
     await this.withTransaction(async () => {
       const col = this.db.sourceTables(this.replicationStreamId);
@@ -251,7 +624,9 @@ export class MongoBucketBatchV3 extends MongoBucketBatch {
     return result!;
   }
 
-  async getSourceTableStatus(table: storage.SourceTable): Promise<storage.SourceTable | null> {
+  protected override async getSourceTableStatusInternal(
+    table: storage.SourceTable
+  ): Promise<storage.SourceTable | null> {
     const doc = (await this.db
       .sourceTables(this.replicationStreamId)
       .findOne({ _id: mongoTableId(table.id) }, { session: this.session })) as SourceTableDocumentV3 | null;
@@ -268,15 +643,38 @@ export class MongoBucketBatchV3 extends MongoBucketBatch {
     );
   }
 
-  async commit(lsn: string, options?: storage.BucketBatchCommitOptions): Promise<storage.CheckpointResult> {
+  protected override async commitInternal(
+    lsn: string,
+    options?: storage.BucketBatchCommitOptions
+  ): Promise<storage.CheckpointResult> {
+    this.options.replicationLock.throwIfAborted();
+    if (this.markRecordUnavailable != null && this.batch?.hasData()) {
+      // A combined data+checkpoint commit cannot return control to the connector
+      // to process repair notifications between those two durable boundaries.
+      // Require its explicit flush -> repair -> commit ordering even when this
+      // particular input might turn out not to need repair. An empty batch does
+      // not prove its external repair queue is drained; that remains the caller's
+      // responsibility. keepalive() delegates here and has the same requirement.
+      throw new ReplicationAssertionError(
+        'Flush pending rows and process markRecordUnavailable repairs before committing source progress'
+      );
+    }
     const { createEmptyCheckpoints } = { ...storage.DEFAULT_BUCKET_BATCH_COMMIT_OPTIONS, ...options };
 
     using _ = this.tracer.span('storage', 'commit');
 
-    const checkpoint = async (stream: SyncRuleDocumentBase, lastOp?: InternalOpId) => {
+    const checkpoint = async (
+      stream: SyncRuleDocumentBase,
+      lastOp?: InternalOpId,
+      additionalWrites?: MongoWriteBatch
+    ) => {
       const now = new Date();
       const preUpdateDocument = stream as ReplicationStreamDocumentV3;
-      const writes = this.db.createWriteBatch(this.session, { ordered: false });
+      this.assertPreparationContext(preUpdateDocument);
+      // Combined commits join the final data group's batch. Empty commits own
+      // and execute their batch here. Decisions depend only on the fresh fence
+      // and planned IDs, so no read of the queued data writes is required.
+      const writes = additionalWrites ?? this.db.createWriteBatch(this.session, { ordered: false });
       writes.updateMany(
         this.db.write_checkpoints,
         { processed_at_lsn: null, 'lsns.1': { $lte: lsn } },
@@ -398,10 +796,12 @@ export class MongoBucketBatchV3 extends MongoBucketBatch {
         );
       }
 
-      await writes.execute();
+      if (additionalWrites == null) {
+        await writes.execute();
+      }
       return { checkpointBlocked, checkpointCreated, checkpointLogState, newCheckpoint };
     };
-    const projection = { sync_configs: 1, last_persisted_op: 1 };
+    const projection = { state: 1, sync_configs: 1, last_persisted_op: 1 };
     const { checkpointBlocked, checkpointCreated, checkpointLogState, newCheckpoint } = await this.flushAndCommit(
       checkpoint,
       () => this.withFencedTransaction(() => this.fence(this.session, projection), checkpoint),
@@ -432,11 +832,7 @@ export class MongoBucketBatchV3 extends MongoBucketBatch {
     };
   }
 
-  async keepalive(lsn: string): Promise<storage.CheckpointResult> {
-    return await this.commit(lsn, { createEmptyCheckpoints: true });
-  }
-
-  async setResumeLsn(lsn: string): Promise<void> {
+  protected override async setResumeLsnInternal(lsn: string): Promise<void> {
     using _ = this.tracer.span('storage', 'set_resume_lsn');
     // 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.
@@ -574,13 +970,16 @@ export class MongoBucketBatchV3 extends MongoBucketBatch {
     });
   }
 
-  async markAllSnapshotDone(no_checkpoint_before_lsn: string): Promise<void> {
+  protected override async markAllSnapshotDoneInternal(no_checkpoint_before_lsn: string): Promise<void> {
     await this.updateStreamMetadata(this.snapshotDoneUpdate(no_checkpoint_before_lsn), {
       'sync_configs._id': { $in: this.syncConfigIds }
     });
   }
 
-  async markSnapshotDone(no_checkpoint_before_lsn: string, options?: { throwOnConflict?: boolean }): Promise<void> {
+  protected override async markSnapshotDoneInternal(
+    no_checkpoint_before_lsn: string,
+    options?: { throwOnConflict?: boolean }
+  ): Promise<void> {
     await this.withTransaction(async () => {
       // Protect against race conditions
       const blockingTables = await this.db
@@ -607,7 +1006,7 @@ export class MongoBucketBatchV3 extends MongoBucketBatch {
     });
   }
 
-  async markTableSnapshotRequired(table: storage.SourceTable): Promise<void> {
+  protected override async markTableSnapshotRequiredInternal(table: storage.SourceTable): Promise<void> {
     const syncConfigIds = this.relevantSyncConfigIds(table);
     if (syncConfigIds.length == 0) {
       return;
@@ -617,7 +1016,7 @@ export class MongoBucketBatchV3 extends MongoBucketBatch {
     });
   }
 
-  async markTableSnapshotDone(
+  protected override async markTableSnapshotDoneInternal(
     tables: storage.SourceTable[],
     no_checkpoint_before_lsn?: string
   ): Promise<storage.SourceTable[]> {
@@ -667,23 +1066,22 @@ export class MongoBucketBatchV3 extends MongoBucketBatch {
     });
   }
 
-  protected override async batchCreateCustomWriteCheckpoints(
-    session: lib_mongo.mongo.ClientSession,
-    opId: InternalOpId
-  ): Promise<void> {
-    if (this.write_checkpoint_batch.length == 0) {
+  private queueCustomWriteCheckpoints(
+    writes: MongoWriteBatch,
+    opId: InternalOpId,
+    checkpoints: storage.CustomWriteCheckpointOptions[]
+  ): void {
+    if (checkpoints.length == 0) {
       return;
     }
 
-    const checkpointsByEvent = Map.groupBy(this.write_checkpoint_batch, (checkpoint) =>
+    const checkpointsByEvent = Map.groupBy(checkpoints, (checkpoint) =>
       this.validateCustomCheckpointEventId(checkpoint.event_id)
     );
 
-    const writes = this.db.createWriteBatch(session, { ordered: false });
     for (const [eventId, checkpoints] of checkpointsByEvent) {
       this.batchCreateEventCustomWriteCheckpoints(writes, opId, eventId, checkpoints);
     }
-    await writes.execute();
   }
 
   protected override async prepareCustomWriteCheckpoints(): Promise<void> {

package/src/storage/implementation/v3/MongoCompactorV3.ts

@@ -645,7 +645,7 @@ export class MongoCompactorV3 extends MongoCompactor implements CompactIntervalC
     bucketContext: BucketDataContextV3,
     objectStorageUsage: ObjectStorageUsage
   ): Promise<BucketStatsWithChecksum> {
-    if (inputs.length == 1) {
+    if (inputs.length == 1 && !this.shouldMoveToObjectStorage(inputs[0])) {
       return statsForDocument(inputs[0]);
     }
 
@@ -1093,7 +1093,7 @@ export class MongoCompactorV3 extends MongoCompactor implements CompactIntervalC
     context: { replicationStreamId: number; definitionId: string },
     objectStorageUsage: ObjectStorageUsage
   ): Promise<CompactionGroupResult> {
-    if (group.inputs.length == 1 && !group.changed) {
+    if (group.inputs.length == 1 && !group.changed && !this.shouldMoveToObjectStorage(group.inputs[0])) {
       return {
         documentId: group.inputs[0]._id,
         stats: statsForDocument(group.inputs[0])
@@ -1588,6 +1588,14 @@ export class MongoCompactorV3 extends MongoCompactor implements CompactIntervalC
     objectStorageUsage.applyDelta(definitionId, newBytes - oldBytes, writes);
   }
 
+  private shouldMoveToObjectStorage(document: BucketDataDocumentV3): boolean {
+    return (
+      this.storage.objectStorage != null &&
+      document.storage_ref == null &&
+      document.size > this.storage.inlineThresholdBytes
+    );
+  }
+
   private async persistBucketData(
     bucket: string,
     chunks: BucketDataDoc[][],
@@ -1607,8 +1615,9 @@ export class MongoCompactorV3 extends MongoCompactor implements CompactIntervalC
     const store = new BucketDataObjectStorage(this.storage.objectStorage);
     const storagePaths = new Set<string>();
     const lifecycle = this.objectStorageLifecycle;
-    // Base placement on the final compacted size. Unchanged documents are not
-    // rewritten, while small MOVE/merge results and CLEAR ops stay inline.
+    // Base placement on the final compacted size. Unchanged documents are only
+    // rewritten to move them to object storage, while small MOVE/merge results
+    // and CLEAR ops stay inline.
     const storedIndexes = serializedChunks.flatMap((document, index) =>
       document.size > this.storage.inlineThresholdBytes ? [index] : []
     );