@powersync/service-module-mongodb-storage 0.20.0 → 0.22.0

package/src/storage/implementation/MongoBucketBatch.ts

@@ -1,5 +1,11 @@
 import { mongo } from '@powersync/lib-service-mongodb';
-import { HydratedSyncConfig, SqlEventDescriptor, SqliteRow, SqliteValue } from '@powersync/service-sync-rules';
+import {
+  EventDefinitionId,
+  HydratedEventDescriptor,
+  HydratedSyncConfig,
+  SqliteRow,
+  SqliteValue
+} from '@powersync/service-sync-rules';
 import * as bson from 'bson';
 
 import {
@@ -30,9 +36,10 @@ import type { VersionedPowerSyncMongo } from './db.js';
 import { MAX_ROW_SIZE } from './MongoBucketBatchShared.js';
 import { MongoIdSequence } from './MongoIdSequence.js';
 import { MongoParsedSyncConfigSet } from './MongoParsedSyncConfigSet.js';
-import { batchCreateCustomWriteCheckpoints } from './MongoWriteCheckpointAPI.js';
+import { MongoWriteBatch } from './MongoWriteBatch.js';
 import { OperationBatch, RecordOperation } from './OperationBatch.js';
 import { ObjectStorage } from './v3/object-storage/ObjectStorage.js';
+import { createObjectStorageUsageWriterId } from './v3/object-storage/ObjectStorageUsage.js';
 
 // Currently, we can only have a single flush() at a time, since it locks the op_id sequence.
 // While the MongoDB transaction retry mechanism handles this okay, using an in-process Mutex
@@ -72,6 +79,10 @@ export interface MongoBucketBatchOptions {
 
   logger: Logger;
   tracer?: PerformanceTracer<'storage' | 'evaluate'>;
+  /**
+   * Aborts in-flight object storage uploads when replication stops.
+   */
+  signal?: AbortSignal;
 
   objectStorage?: ObjectStorage;
   inlineThresholdBytes?: number;
@@ -104,9 +115,10 @@ export abstract class MongoBucketBatch
   private readonly storeCurrentData: boolean;
   public readonly skipExistingRows: boolean;
   protected readonly mapping: BucketDefinitionMapping;
+  protected readonly objectStorageUsageWriterId = createObjectStorageUsageWriterId();
 
   private batch: OperationBatch | null = null;
-  private write_checkpoint_batch: storage.CustomWriteCheckpointOptions[] = [];
+  protected write_checkpoint_batch: storage.CustomWriteCheckpointOptions[] = [];
   private markRecordUnavailable: BucketStorageMarkRecordUnavailable | undefined;
   private hooks: storage.StorageHooks | undefined;
   private clearedError = false;
@@ -185,6 +197,11 @@ export abstract class MongoBucketBatch
     no_checkpoint_before_lsn?: string
   ): Promise<storage.SourceTable[]>;
 
+  protected abstract batchCreateCustomWriteCheckpoints(session: mongo.ClientSession, opId: InternalOpId): Promise<void>;
+
+  /** Perform any version-specific setup required before writing custom checkpoints. */
+  protected async prepareCustomWriteCheckpoints(): Promise<void> {}
+
   async flush(options?: storage.BatchBucketFlushOptions): Promise<storage.FlushedResult | null> {
     let result: storage.FlushedResult | null = null;
     // One flush may be split over multiple transactions.
@@ -207,6 +224,10 @@ export abstract class MongoBucketBatch
     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;
@@ -218,13 +239,15 @@ export abstract class MongoBucketBatch
 
       if (this.write_checkpoint_batch.length > 0) {
         this.logger.info(`Writing ${this.write_checkpoint_batch.length} custom write checkpoints`);
-        await batchCreateCustomWriteCheckpoints(this.db, session, this.write_checkpoint_batch, opSeq.next());
-        this.write_checkpoint_batch = [];
+        await this.batchCreateCustomWriteCheckpoints(session, opSeq.next());
       }
 
       last_op = opSeq.last();
     });
 
+    // Keep checkpoints available if the transaction retries after writing them.
+    this.write_checkpoint_batch = [];
+
     if (clearedError) {
       this.clearedError = true;
     }
@@ -686,7 +709,9 @@ export abstract class MongoBucketBatch
 
       await callback(session, opSeq);
 
-      await this.db.op_id_sequence.updateOne(
+      const writes = this.db.createWriteBatch(session, { ordered: false });
+      writes.updateOne(
+        this.db.op_id_sequence,
         {
           _id: 'main'
         },
@@ -694,13 +719,11 @@ export abstract class MongoBucketBatch
           $set: {
             op_id: opSeq.last()
           }
-        },
-        {
-          session
         }
       );
 
-      await this.db.sync_rules.updateOne(
+      writes.updateOne(
+        this.db.sync_rules,
         {
           _id: this.replicationStreamId
         },
@@ -708,13 +731,13 @@ export abstract class MongoBucketBatch
           $set: {
             last_keepalive_ts: new Date()
           }
-        },
-        { session }
+        }
       );
 
       // Allow subclasses to persist additional flush-time state in the same transaction
       // (e.g. v3 advances the stream-level last_persisted_op).
-      await this.onReplicationTransactionFlush(session, opSeq.last());
+      this.onReplicationTransactionFlush(writes, opSeq.last());
+      await writes.execute();
       // We don't notify checkpoint here - we don't make any checkpoint updates directly
     });
   }
@@ -725,7 +748,7 @@ export abstract class MongoBucketBatch
    * The base implementation does nothing; v3 storage overrides this to `$max` the stream-level
    * `last_persisted_op` durably within the same transaction.
    */
-  protected async onReplicationTransactionFlush(_session: mongo.ClientSession, _lastOp: InternalOpId): Promise<void> {
+  protected onReplicationTransactionFlush(_writes: MongoWriteBatch, _lastOp: InternalOpId): void {
     // No-op by default (v1 behaviour unchanged).
   }
 
@@ -762,11 +785,12 @@ export abstract class MongoBucketBatch
   async save(record: storage.SaveOptions): Promise<storage.FlushedResult | null> {
     const { after, before, sourceTable, tag } = record;
     const storeCurrentData = this.storeCurrentData && sourceTable.storeCurrentData;
-    // syncEvent is the per-table designation from resolveTables. With v3 storage, multiple
-    // SourceTables can exist for the same ref, with a row change saved once per table -
-    // only the designated event carrier may fire events, so each event fires once per row.
+    // V3 source tables own disjoint event-definition ids for each physical table. Multiple
+    // SourceTables may evaluate different events, but each definition is evaluated through
+    // at most one record for this row change.
+    // Legacy storage leaves eventDefinitionIds undefined and selects by table ref.
     if (sourceTable.syncEvent) {
-      for (const event of this.getTableEvents(sourceTable)) {
+      for (const { event, eventId } of this.getTableEvents(sourceTable)) {
         this.iterateListeners((cb) =>
           cb.replicationEvent?.({
             batch: this,
@@ -776,7 +800,8 @@ export abstract class MongoBucketBatch
               after: after && utils.isCompleteRow(storeCurrentData, after) ? after : undefined,
               before: before && utils.isCompleteRow(storeCurrentData, before) ? before : undefined
             },
-            event
+            event,
+            event_id: eventId
           })
         );
       }
@@ -934,11 +959,29 @@ export abstract class MongoBucketBatch
   }
 
   /**
-   * Gets relevant {@link SqlEventDescriptor}s for the given {@link SourceTable}
+   * Gets relevant {@link HydratedEventDescriptor}s for the given {@link SourceTable}
    */
-  protected getTableEvents(table: storage.SourceTable): SqlEventDescriptor[] {
-    return this.sync_rules.eventDescriptors.filter((evt) =>
-      [...evt.getSourceTables()].some((sourceTable) => sourceTable.matches(table.ref))
-    );
+  protected getTableEvents(
+    table: storage.SourceTable
+  ): { event: HydratedEventDescriptor; eventId?: EventDefinitionId }[] {
+    // V3 storage assigns event-definition ids to each source table, so membership is authoritative.
+    // Iterate the table's distinct ids and resolve each through the stream's deduped event map, so a
+    // definition reused across configs has one evaluator for that id. The evaluator may still return
+    // multiple payloads from matching queries. Legacy storage leaves this undefined and selects by table ref.
+    if (table.eventDefinitionIds != null) {
+      const eventById = this.options.parsedSyncConfig.eventById;
+      const events: { event: HydratedEventDescriptor; eventId: EventDefinitionId }[] = [];
+      for (const id of table.eventDefinitionIds) {
+        const event = eventById.get(id);
+        if (event != null && event.tableTriggersEvent(table.ref)) {
+          events.push({ event, eventId: id });
+        }
+      }
+      return events;
+    }
+
+    return this.sync_rules.eventDescriptors
+      .filter((event) => event.tableTriggersEvent(table.ref))
+      .map((event) => ({ event }));
   }
 }

package/src/storage/implementation/MongoCompactor.ts

@@ -1,64 +1,27 @@
 import * as timers from 'node:timers/promises';
 
-import { isMongoServerError, mongo, MONGO_OPERATION_TIMEOUT_MS } from '@powersync/lib-service-mongodb';
-import {
-  logger as defaultLogger,
-  Logger,
-  ReplicationAssertionError,
-  ServiceAssertionError
-} from '@powersync/lib-services-framework';
-import { InternalOpId, isPartialChecksum, PopulateChecksumCacheResults, storage } from '@powersync/service-core';
-import { BucketDefinitionId } from '@powersync/service-sync-rules';
+import { isMongoServerError, mongo } from '@powersync/lib-service-mongodb';
+import { logger as defaultLogger, Logger, ReplicationAssertionError } from '@powersync/lib-services-framework';
+import { storage } from '@powersync/service-core';
 
 import { BucketKey } from './common/BucketDataDoc.js';
 import type { VersionedPowerSyncMongo } from './db.js';
-import { BucketStateDocumentBase } from './models.js';
 import type { MongoSyncBucketStorage } from './MongoSyncBucketStorage.js';
 import { isRetryableObjectStorageError } from './v3/object-storage/ObjectStorage.js';
 
-export interface CurrentBucketState {
-  /** Bucket name */
-  bucket: string;
-  definitionId: BucketDefinitionId;
+export interface MongoCompactOptions extends storage.CompactOptions {
   /**
-   * Rows seen in the bucket, with the last op_id of each.
+   * Only merge adjacent V3 bucket-data chunks. This is used after initial
+   * replication, where reading every operation would defeat the purpose of
+   * the lightweight pass.
    */
-  seen: Map<string, InternalOpId>;
-  /**
-   * Estimated memory usage of the seen Map.
-   */
-  trackingSize: number;
-  /**
-   * Last (lowest) seen op_id that is not a PUT.
-   */
-  lastNotPut: InternalOpId | null;
-  /**
-   * Number of REMOVE/MOVE operations seen since lastNotPut.
-   */
-  opsSincePut: number;
-  /**
-   * Incrementally-updated checksum, up to maxOpId.
-   */
-  checksum: number;
-  /**
-   * Op count for the checksum.
-   */
-  opCount: number;
-  /**
-   * Byte size of ops covered by the checksum.
-   */
-  opBytes: number;
+  compactChunksOnly?: boolean;
 }
 
-export interface MongoCompactOptions extends storage.CompactOptions {}
-
 const DEFAULT_CLEAR_BATCH_LIMIT = 5000;
 const DEFAULT_MOVE_BATCH_LIMIT = 2000;
 const DEFAULT_MOVE_BATCH_QUERY_LIMIT = 10_000;
 const DEFAULT_MOVE_BATCH_BYTE_LIMIT = 16 * 1024 * 1024;
-const DEFAULT_MIN_BUCKET_CHANGES = 10;
-const DEFAULT_MIN_CHANGE_RATIO = 0.1;
-const DIRTY_BUCKET_SCAN_BATCH_SIZE = 2_000;
 /** This default is primarily for tests. */
 const DEFAULT_MEMORY_LIMIT_MB = 64;
 const COMPACTION_RETRY_LIMIT = 3;
@@ -71,28 +34,30 @@ export class ConcurrentCompactionError extends Error {
   }
 }
 
-export interface DirtyBucket {
-  bucket: string;
-  definitionId: BucketDefinitionId | null;
-  estimatedCount: number;
-  dirtyRatio?: number;
+/**
+ * A worker whose bucket lease has been replaced must stop immediately. Unlike
+ * a transactional replacement conflict, retrying it with the same compactor
+ * instance would still use the stale lease and could race the new owner.
+ */
+export class CompactionLeaseLostError extends ConcurrentCompactionError {
+  constructor(message: string) {
+    super(message);
+    this.name = 'CompactionLeaseLostError';
+  }
 }
 
 export abstract class MongoCompactor {
-  protected bucketStateUpdates: mongo.AnyBulkWriteOperation<BucketStateDocumentBase>[] = [];
-
   protected readonly idLimitBytes: number;
   protected readonly moveBatchLimit: number;
   protected readonly moveBatchQueryLimit: number;
   protected readonly moveBatchByteLimit: number;
   protected readonly clearBatchLimit: number;
-  protected readonly minBucketChanges: number;
-  protected readonly minChangeRatio: number;
-  protected readonly maxOpId: bigint;
   protected readonly buckets: string[] | undefined;
   protected readonly deleteCheckpointRequestsBefore: Date | undefined;
   protected readonly signal?: AbortSignal;
   protected readonly group_id: number;
+  protected readonly compactChunksOnly: boolean;
+  protected compactedBucketCount = 0;
 
   protected readonly logger: Logger;
 
@@ -110,35 +75,16 @@ export abstract class MongoCompactor {
     if (this.clearBatchLimit < 2) {
       throw new ReplicationAssertionError('clearBatchLimit must be >= 2');
     }
-    this.minBucketChanges = options.minBucketChanges ?? DEFAULT_MIN_BUCKET_CHANGES;
-    this.minChangeRatio = options.minChangeRatio ?? DEFAULT_MIN_CHANGE_RATIO;
-    this.maxOpId = options.maxOpId ?? 0n;
     this.buckets = options.compactBuckets;
     this.deleteCheckpointRequestsBefore = options.deleteCheckpointRequestsBefore;
     this.signal = options.signal;
+    this.compactChunksOnly = options.compactChunksOnly ?? false;
     this.logger = options.logger ?? defaultLogger;
   }
 
-  /**
-   * Compact buckets by converting operations into MOVE and/or CLEAR operations.
-   *
-   * See /docs/storage/compacting-operations.md for details.
-   */
-  async compact() {
-    await this.deleteOldCheckpointRequests();
-
-    if (this.buckets) {
-      for (const bucket of this.buckets) {
-        // We can make this more efficient later on by iterating through the buckets in a single query.
-        // That makes batching more tricky, so we leave for later.
-        await this.compactSingleBucketRetried(bucket);
-      }
-    } else {
-      await this.compactDirtyBuckets();
-    }
-  }
+  abstract compact(): Promise<number>;
 
-  private async deleteOldCheckpointRequests() {
+  protected async deleteOldCheckpointRequests() {
     if (this.deleteCheckpointRequestsBefore == null) {
       return;
     }
@@ -152,180 +98,17 @@ export abstract class MongoCompactor {
       checkpoint_requested_at: { $exists: true, $lt: this.deleteCheckpointRequestsBefore },
       processed_at_lsn: { $ne: null }
     });
-    await this.db.custom_write_checkpoints.deleteMany({
-      checkpoint_requested_at: { $exists: true, $lt: this.deleteCheckpointRequestsBefore }
-    });
-  }
-
-  /**
-   * Subset of compact, only populating checksums where relevant.
-   */
-  async populateChecksums(options: { minBucketChanges: number }): Promise<PopulateChecksumCacheResults> {
-    let count = 0;
-    // Paginate through dirty buckets in batches until no more buckets meet the criteria.
-    while (true) {
-      this.signal?.throwIfAborted();
-      const buckets = await this.dirtyBucketBatchForChecksums(options);
-      if (buckets.length == 0) {
-        break;
-      }
-      this.signal?.throwIfAborted();
-
-      const start = Date.now();
-      // Filter batch by estimated bucket size, to reduce possibility of timeouts.
-      const checkBuckets: typeof buckets = [];
-      let totalCountEstimate = 0;
-      for (const bucket of buckets) {
-        checkBuckets.push(bucket);
-        totalCountEstimate += bucket.estimatedCount;
-        if (totalCountEstimate > 50_000) {
-          break;
-        }
-      }
-      this.logger.info(
-        `Calculating checksums for batch of ${buckets.length} buckets, estimated count of ${totalCountEstimate}`
-      );
-      await this.updateChecksumsBatch(checkBuckets);
-      this.logger.info(`Updated checksums for batch of ${checkBuckets.length} buckets in ${Date.now() - start}ms`);
-      count += checkBuckets.length;
-    }
-    return { buckets: count };
+    await this.deleteOldCustomCheckpointRequests();
   }
 
-  protected async *dirtyBucketBatchesForCollection<TCollectionBucketState extends BucketStateDocumentBase>(
-    collection: mongo.Collection<TCollectionBucketState>,
-    lastId: TCollectionBucketState['_id'],
-    maxId: TCollectionBucketState['_id'],
-    options: {
-      minBucketChanges: number;
-      minChangeRatio: number;
-    },
-    getDefinitionId: (state: TCollectionBucketState) => BucketDefinitionId | null
-  ): AsyncGenerator<DirtyBucket[]> {
-    // Paginate through the bucket state collection using cursor-based scanning.
-    while (true) {
-      // To avoid timeouts from too many buckets not meeting the minBucketChanges criteria, use an aggregation pipeline
-      // to scan a fixed batch of buckets at a time, but only return buckets that meet the criteria.
-      const [result] = await collection
-        .aggregate<{
-          buckets: TCollectionBucketState[];
-          cursor: Pick<TCollectionBucketState, '_id'>[];
-        }>(
-          [
-            {
-              $match: {
-                _id: { $gt: lastId, $lt: maxId }
-              }
-            },
-            {
-              $sort: { _id: 1 }
-            },
-            {
-              // Scan a fixed number of docs each query so sparse matches don't block progress.
-              $limit: DIRTY_BUCKET_SCAN_BATCH_SIZE
-            },
-            {
-              $facet: {
-                buckets: [
-                  {
-                    $match: {
-                      'estimate_since_compact.count': { $gte: options.minBucketChanges }
-                    }
-                  },
-                  {
-                    $project: {
-                      _id: 1,
-                      estimate_since_compact: 1,
-                      compacted_state: 1
-                    }
-                  }
-                ],
-                // This is used for the next query.
-                cursor: [{ $sort: { _id: -1 } }, { $limit: 1 }, { $project: { _id: 1 } }]
-              }
-            }
-          ],
-          { maxTimeMS: MONGO_OPERATION_TIMEOUT_MS }
-        )
-        .toArray();
-
-      const cursor = result?.cursor?.[0];
-      if (cursor == null) {
-        break;
-      }
-      lastId = cursor._id;
-
-      const mapped = (result?.buckets ?? []).map((bucketState) => {
-        // The numbers, specifically the bytes, could be a bigint. Convert to Number to allow calculating ratios.
-        // BigInt precision is not needed here since this is only an estimate.
-        const updatedCount = bucketState.estimate_since_compact?.count ?? 0;
-        const totalCount = (bucketState.compacted_state?.count ?? 0) + updatedCount;
-        const updatedBytes = Number(bucketState.estimate_since_compact?.bytes ?? 0);
-        const totalBytes = Number(bucketState.compacted_state?.bytes ?? 0) + updatedBytes;
-        const dirtyChangeNumber = totalCount > 0 ? updatedCount / totalCount : 0;
-        const dirtyChangeBytes = totalBytes > 0 ? updatedBytes / totalBytes : 0;
-        return {
-          bucket: bucketState._id.b,
-          definitionId: getDefinitionId(bucketState),
-          estimatedCount: totalCount,
-          dirtyRatio: Math.max(dirtyChangeNumber, dirtyChangeBytes)
-        };
-      });
-
-      yield mapped.filter(
-        (bucket) => bucket.estimatedCount >= options.minBucketChanges && bucket.dirtyRatio >= options.minChangeRatio
-      );
+  protected async deleteOldCustomCheckpointRequests() {
+    if (this.deleteCheckpointRequestsBefore == null) {
+      return;
     }
-  }
-
-  protected async dirtyBucketBatchForChecksumsForCollection<TBucketState extends BucketStateDocumentBase>(
-    collection: mongo.Collection<TBucketState>,
-    filter: mongo.Filter<TBucketState>,
-    getDefinitionId: (state: mongo.WithId<TBucketState>) => BucketDefinitionId | null
-  ): Promise<DirtyBucket[]> {
-    const dirtyBuckets = await collection
-      .find(filter, {
-        projection: {
-          _id: 1,
-          estimate_since_compact: 1,
-          compacted_state: 1
-        },
-        sort: {
-          'estimate_since_compact.count': -1
-        },
-        limit: 200,
-        maxTimeMS: MONGO_OPERATION_TIMEOUT_MS
-      })
-      .toArray();
-
-    return dirtyBuckets.map((bucket) => ({
-      bucket: bucket._id.b,
-      definitionId: getDefinitionId(bucket),
-      estimatedCount: Number(bucket.estimate_since_compact!.count) + Number(bucket.compacted_state?.count ?? 0)
-    }));
-  }
 
-  public abstract dirtyBucketBatches(options: {
-    minBucketChanges: number;
-    minChangeRatio: number;
-  }): AsyncGenerator<DirtyBucket[]>;
-
-  public abstract dirtyBucketBatchForChecksums(options: { minBucketChanges: number }): Promise<DirtyBucket[]>;
-
-  protected async compactDirtyBuckets() {
-    for await (const buckets of this.dirtyBucketBatches({
-      minBucketChanges: this.minBucketChanges,
-      minChangeRatio: this.minChangeRatio
-    })) {
-      this.signal?.throwIfAborted();
-      if (buckets.length == 0) {
-        continue;
-      }
-
-      for (const { bucket, definitionId } of buckets) {
-        await this.compactSingleBucketRetried(bucket, definitionId);
-      }
-    }
+    await this.db.custom_write_checkpoints.deleteMany({
+      checkpoint_requested_at: { $exists: true, $lt: this.deleteCheckpointRequestsBefore }
+    });
   }
 
   /**
@@ -336,14 +119,12 @@ export abstract class MongoCompactor {
    * storage paths are also prepared with lifecycle markers, so a retry can
    * safely overwrite or eventually clean up uploads from the failed attempt.
    */
-  protected async compactSingleBucketRetried(bucket: string, definitionId: BucketDefinitionId | null = null) {
+  protected async retryCompaction(bucket: string, compact: () => Promise<void>) {
     let retryCount = 0;
     while (true) {
       this.signal?.throwIfAborted();
-      // Do not carry queued bucket state writes from a failed attempt into the rescan.
-      this.bucketStateUpdates = [];
       try {
-        await this.compactSingleBucket(bucket, definitionId);
+        await compact();
         return;
       } catch (e) {
         if (this.signal?.aborted) {
@@ -372,101 +153,6 @@ export abstract class MongoCompactor {
       }
     }
   }
-
-  protected abstract compactSingleBucket(bucket: string, definitionId?: BucketDefinitionId | null): Promise<void>;
-
-  protected collectBucketStateUpdates(
-    state: CurrentBucketState,
-    compactedOpId: InternalOpId
-  ): mongo.AnyBulkWriteOperation<BucketStateDocumentBase> {
-    if (state.opCount < 0) {
-      throw new ServiceAssertionError(
-        `Invalid opCount: ${state.opCount} checksum ${state.checksum} opsSincePut: ${state.opsSincePut} maxOpId: ${this.maxOpId}`
-      );
-    }
-    return {
-      updateOne: {
-        filter: this.bucketStateFilter(state.bucket, state.definitionId),
-        update: {
-          $set: {
-            compacted_state: {
-              op_id: compactedOpId,
-              count: state.opCount,
-              checksum: BigInt(state.checksum),
-              bytes: state.opBytes
-            },
-            estimate_since_compact: {
-              // There could have been a whole bunch of new operations added to the bucket while compacting,
-              // which we don't currently cater for. We could potentially query for that, but that adds overhead.
-              count: 0,
-              bytes: 0
-            }
-          }
-        } satisfies mongo.UpdateFilter<BucketStateDocumentBase>,
-        // We generally expect this to have been created before.
-        // We don't create new ones here, to avoid issues with the unique index on bucket_updates.
-        upsert: false
-      }
-    };
-  }
-
-  protected updateBucketChecksums(state: CurrentBucketState, compactedOpId: InternalOpId) {
-    this.bucketStateUpdates.push(this.collectBucketStateUpdates(state, compactedOpId));
-  }
-
-  protected async flushBucketStateUpdates() {
-    if (this.bucketStateUpdates.length > 0) {
-      this.logger.info(`Updating ${this.bucketStateUpdates.length} bucket states`);
-      await this.writeBucketStateUpdates();
-      this.bucketStateUpdates = [];
-    }
-  }
-
-  protected async updateChecksumsBatch(buckets: Pick<DirtyBucket, 'bucket' | 'definitionId'>[]) {
-    const checksums = await this.computeChecksumsForBuckets(buckets);
-    const definitionIdByBucket = new Map(buckets.map((bucket) => [bucket.bucket, bucket.definitionId]));
-
-    for (const bucketChecksum of checksums.values()) {
-      if (isPartialChecksum(bucketChecksum)) {
-        // Should never happen since we don't specify `start`.
-        throw new ServiceAssertionError(`Full checksum expected, got ${JSON.stringify(bucketChecksum)}`);
-      }
-
-      this.bucketStateUpdates.push({
-        updateOne: {
-          filter: this.bucketStateFilter(
-            bucketChecksum.bucket,
-            definitionIdByBucket.get(bucketChecksum.bucket) ?? null
-          ),
-          update: {
-            $set: {
-              compacted_state: {
-                op_id: this.maxOpId,
-                count: bucketChecksum.count,
-                checksum: BigInt(bucketChecksum.checksum),
-                bytes: null
-              },
-              estimate_since_compact: {
-                count: 0,
-                bytes: 0
-              }
-            }
-          } satisfies mongo.UpdateFilter<BucketStateDocumentBase>,
-          // We don't create new ones here - it gets tricky to get the last_op right with the unique index on
-          // bucket_updates.
-          upsert: false
-        }
-      });
-    }
-
-    await this.flushBucketStateUpdates();
-  }
-
-  protected abstract writeBucketStateUpdates(): Promise<void>;
-  protected abstract computeChecksumsForBuckets(
-    buckets: Pick<DirtyBucket, 'bucket' | 'definitionId'>[]
-  ): Promise<storage.PartialChecksumMap>;
-  protected abstract bucketStateFilter(bucket: string, definitionId: BucketDefinitionId | null): mongo.Document;
 }
 
 export interface BucketDataCollectionContext<TBucketData extends mongo.Document> {
@@ -475,6 +161,9 @@ export interface BucketDataCollectionContext<TBucketData extends mongo.Document>
 }
 
 function compactionRetryReason(error: unknown): string | null {
+  if (error instanceof CompactionLeaseLostError) {
+    return null;
+  }
   if (error instanceof ConcurrentCompactionError) {
     return 'concurrent compaction';
   }