@powersync/service-module-mongodb-storage 0.15.4 → 0.16.0

package/src/storage/implementation/MongoCompactor.ts

@@ -1,5 +1,10 @@
 import { isMongoServerError, mongo, MONGO_OPERATION_TIMEOUT_MS } from '@powersync/lib-service-mongodb';
-import { logger, ReplicationAssertionError, ServiceAssertionError } from '@powersync/lib-services-framework';
+import {
+  logger as defaultLogger,
+  Logger,
+  ReplicationAssertionError,
+  ServiceAssertionError
+} from '@powersync/lib-services-framework';
 import {
   addChecksums,
   InternalOpId,
@@ -9,15 +14,18 @@ import {
   utils
 } from '@powersync/service-core';
 
-import { VersionedPowerSyncMongo } from './db.js';
-import { BucketDataDocument, BucketDataKey, BucketStateDocument } from './models.js';
-import { MongoSyncBucketStorage } from './MongoSyncBucketStorage.js';
+import { BucketDefinitionId } from './BucketDefinitionMapping.js';
+import { BucketDataDoc, BucketKey } from './common/BucketDataDoc.js';
+import { BucketDataDocumentGeneric, SingleBucketStore } from './common/SingleBucketStore.js';
+import type { VersionedPowerSyncMongo } from './db.js';
+import { BucketStateDocumentBase } from './models.js';
+import type { MongoSyncBucketStorage } from './MongoSyncBucketStorage.js';
 import { cacheKey } from './OperationBatch.js';
 
 interface CurrentBucketState {
   /** Bucket name */
   bucket: string;
-
+  definitionId: BucketDefinitionId;
   /**
    * Rows seen in the bucket, with the last op_id of each.
    */
@@ -26,36 +34,30 @@ interface CurrentBucketState {
    * 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
+   * Incrementally-updated checksum, up to maxOpId.
    */
   checksum: number;
-
   /**
-   * op count for the checksum
+   * Op count for the checksum.
    */
   opCount: number;
-
   /**
    * Byte size of ops covered by the checksum.
    */
   opBytes: number;
 }
 
-/**
- * Additional options, primarily for testing.
- */
+type CompactClearProperties = 'op' | 'checksum' | 'target_op';
+
 export interface MongoCompactOptions extends storage.CompactOptions {}
 
 const DEFAULT_CLEAR_BATCH_LIMIT = 5000;
@@ -64,28 +66,36 @@ const DEFAULT_MOVE_BATCH_QUERY_LIMIT = 10_000;
 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;
 
-export class MongoCompactor {
-  private updates: mongo.AnyBulkWriteOperation<BucketDataDocument>[] = [];
-  private bucketStateUpdates: mongo.AnyBulkWriteOperation<BucketStateDocument>[] = [];
-
-  private idLimitBytes: number;
-  private moveBatchLimit: number;
-  private moveBatchQueryLimit: number;
-  private clearBatchLimit: number;
-  private minBucketChanges: number;
-  private minChangeRatio: number;
-  private maxOpId: bigint;
-  private buckets: string[] | undefined;
-  private signal?: AbortSignal;
-  private group_id: number;
+export interface DirtyBucket {
+  bucket: string;
+  definitionId: BucketDefinitionId | null;
+  estimatedCount: number;
+  dirtyRatio?: number;
+}
+
+export abstract class MongoCompactor {
+  protected updates: mongo.AnyBulkWriteOperation<BucketDataDocumentGeneric>[] = [];
+  protected bucketStateUpdates: mongo.AnyBulkWriteOperation<BucketStateDocumentBase>[] = [];
+
+  protected readonly idLimitBytes: number;
+  protected readonly moveBatchLimit: number;
+  protected readonly moveBatchQueryLimit: number;
+  protected readonly clearBatchLimit: number;
+  protected readonly minBucketChanges: number;
+  protected readonly minChangeRatio: number;
+  protected readonly maxOpId: bigint;
+  protected readonly buckets: string[] | undefined;
+  protected readonly signal?: AbortSignal;
+  protected readonly group_id: number;
+
+  protected readonly logger: Logger;
 
   constructor(
-    private storage: MongoSyncBucketStorage,
-    private db: VersionedPowerSyncMongo,
+    protected readonly storage: MongoSyncBucketStorage,
+    protected readonly db: VersionedPowerSyncMongo,
     options: MongoCompactOptions
   ) {
     this.group_id = storage.group_id;
@@ -98,6 +108,7 @@ export class MongoCompactor {
     this.maxOpId = options.maxOpId ?? 0n;
     this.buckets = options.compactBuckets;
     this.signal = options.signal;
+    this.logger = options.logger ?? defaultLogger;
   }
 
   /**
@@ -107,9 +118,8 @@ export class MongoCompactor {
    */
   async compact() {
     if (this.buckets) {
-      for (let bucket of this.buckets) {
-        // We can make this more efficient later on by iterating
-        // through the buckets in a single query.
+      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);
       }
@@ -118,8 +128,161 @@ export class MongoCompactor {
     }
   }
 
-  private async compactDirtyBuckets() {
-    for await (let buckets of this.dirtyBucketBatches({
+  /**
+   * Subset of compact, only populating checksums where relevant.
+   */
+  async populateChecksums(options: { minBucketChanges: number }): Promise<PopulateChecksumCacheResults> {
+    let count = 0;
+    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 };
+  }
+
+  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[]> {
+    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 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
     })) {
@@ -128,8 +291,8 @@ export class MongoCompactor {
         continue;
       }
 
-      for (let { bucket } of buckets) {
-        await this.compactSingleBucketRetried(bucket);
+      for (const { bucket, definitionId } of buckets) {
+        await this.compactSingleBucketRetried(bucket, definitionId);
       }
     }
   }
@@ -139,15 +302,15 @@ export class MongoCompactor {
    *
    * This covers against occasional network or other database errors during a long compact job.
    */
-  private async compactSingleBucketRetried(bucket: string) {
+  protected async compactSingleBucketRetried(bucket: string, definitionId: BucketDefinitionId | null = null) {
     let retryCount = 0;
     while (true) {
       try {
-        await this.compactSingleBucket(bucket);
+        await this.compactSingleBucket(bucket, definitionId);
         break;
       } catch (e) {
         if (retryCount < 3 && isMongoServerError(e)) {
-          logger.warn(`Error compacting bucket ${bucket}, retrying...`, e);
+          this.logger.warn(`Error compacting bucket ${bucket}, retrying...`, e);
           retryCount++;
           await new Promise((resolve) => setTimeout(resolve, 1000 * retryCount));
         } else {
@@ -157,64 +320,64 @@ export class MongoCompactor {
     }
   }
 
-  private async compactSingleBucket(bucket: string) {
+  protected async compactSingleBucket(bucket: string, definitionId: BucketDefinitionId | null = null) {
     const idLimitBytes = this.idLimitBytes;
-
-    let currentState: CurrentBucketState = {
+    const bucketContext = await this.getBucketDataContext(bucket, definitionId);
+    if (bucketContext == null) {
+      return;
+    }
+    const currentState: CurrentBucketState = {
       bucket,
+      definitionId: bucketContext.key.definitionId,
       seen: new Map(),
       trackingSize: 0,
       lastNotPut: null,
       opsSincePut: 0,
-
       checksum: 0,
       opCount: 0,
       opBytes: 0
     };
 
-    // Constant lower bound
-    const lowerBound: BucketDataKey = {
-      g: this.group_id,
-      b: bucket,
-      o: new mongo.MinKey() as any
-    };
-
-    // Upper bound is adjusted for each batch
-    let upperBound: BucketDataKey = {
-      g: this.group_id,
-      b: bucket,
-      o: new mongo.MaxKey() as any
-    };
+    // Constant lower bound.
+    const lowerBound = bucketContext.minId;
+    // Upper bound is adjusted for each batch.
+    let upperBound = bucketContext.maxId;
 
     while (true) {
       this.signal?.throwIfAborted();
 
-      // Query one batch at a time, to avoid cursor timeouts
-      const cursor = this.db.bucket_data.aggregate<BucketDataDocument & { size: number | bigint }>(
-        [
-          {
-            $match: {
-              _id: {
-                $gte: lowerBound,
-                $lt: upperBound
-              }
-            }
-          },
-          { $sort: { _id: -1 } },
-          { $limit: this.moveBatchQueryLimit },
-          {
-            $project: {
-              _id: 1,
-              op: 1,
-              table: 1,
-              row_id: 1,
-              source_table: 1,
-              source_key: 1,
-              checksum: 1,
-              size: { $bsonSize: '$$ROOT' }
-            }
+      // Query one batch at a time, to avoid cursor timeouts.
+      const pipeline = [
+        {
+          $match: {
+            _id: {
+              $gte: lowerBound,
+              $lt: upperBound
+            },
+            // Workaround for a clustered collection bug where the $lt operator may include upperBound.
+            // Technically only needed for storage V3.
+            // https://jira.mongodb.org/browse/SERVER-121822
+            '_id.o': { $lt: upperBound.o }
           }
-        ],
+        },
+        { $sort: { _id: -1 } },
+        { $limit: this.moveBatchQueryLimit },
+        {
+          $project: {
+            _id: 1,
+            op: 1,
+            table: 1,
+            row_id: 1,
+            source_table: 1,
+            source_key: 1,
+            checksum: 1,
+            size: { $bsonSize: '$$ROOT' }
+          }
+        }
+      ];
+
+      const cursor = bucketContext.collection.aggregate<BucketDataDocumentGeneric & { size: number | bigint }>(
+        pipeline,
         {
           // batchSize is 1 more than limit to auto-close the cursor.
           // See https://github.com/mongodb/node-mongodb-native/pull/4580
@@ -223,18 +386,25 @@ export class MongoCompactor {
       );
       // We don't limit to a single batch here, since that often causes MongoDB to scan through more than it returns.
       // Instead, we load up to the limit.
-      const batch = await cursor.toArray();
+      const rawBatch = await cursor.toArray();
+      const batch = rawBatch.map((document) => {
+        const { size, ...rest } = document;
+        return {
+          doc: bucketContext.fromPersistedDocument(rest),
+          size
+        };
+      });
 
       if (batch.length == 0) {
-        // We've reached the end
+        // We've reached the end.
         break;
       }
 
-      // Set upperBound for the next batch
-      upperBound = batch[batch.length - 1]._id;
+      // Reuse the exact collection _id value from Mongo for the next bound.
+      upperBound = rawBatch[rawBatch.length - 1]._id;
 
-      for (let doc of batch) {
-        if (doc._id.o > this.maxOpId) {
+      for (const { doc, size } of batch) {
+        if (doc.o > this.maxOpId) {
           continue;
         }
 
@@ -243,19 +413,17 @@ export class MongoCompactor {
 
         let isPersistentPut = doc.op == 'PUT';
 
-        currentState.opBytes += Number(doc.size);
+        currentState.opBytes += Number(size);
         if (doc.op == 'REMOVE' || doc.op == 'PUT') {
           const key = `${doc.table}/${doc.row_id}/${cacheKey(doc.source_table!, doc.source_key!)}`;
           const targetOp = currentState.seen.get(key);
           if (targetOp) {
-            // Will convert to MOVE, so don't count as PUT
+            // Will convert to MOVE, so don't count as PUT.
             isPersistentPut = false;
 
             this.updates.push({
               updateOne: {
-                filter: {
-                  _id: doc._id
-                },
+                filter: { _id: bucketContext.docId(doc.o) },
                 update: {
                   $set: {
                     op: 'MOVE',
@@ -268,24 +436,20 @@ export class MongoCompactor {
                     row_id: 1,
                     data: 1
                   }
-                }
+                } satisfies mongo.UpdateFilter<BucketDataDocumentGeneric>
               }
             });
 
-            currentState.opBytes += 200 - Number(doc.size); // TODO: better estimate for this
-          } else {
-            if (currentState.trackingSize >= idLimitBytes) {
-              // Reached memory limit.
-              // Keep the highest seen values in this case.
-            } else {
-              // flatstr reduces the memory usage by flattening the string
-              currentState.seen.set(utils.flatstr(key), doc._id.o);
-              // length + 16 for the string
-              // 24 for the bigint
-              // 50 for map overhead
-              // 50 for additional overhead
-              currentState.trackingSize += key.length + 140;
-            }
+            // TODO: better estimate for this.
+            currentState.opBytes += 200 - Number(size);
+          } else if (currentState.trackingSize < idLimitBytes) {
+            // flatstr reduces the memory usage by flattening the string.
+            currentState.seen.set(utils.flatstr(key), doc.o);
+            // length + 16 for the string
+            // 24 for the bigint
+            // 50 for map overhead
+            // 50 for additional overhead
+            currentState.trackingSize += key.length + 140;
           }
         }
 
@@ -294,41 +458,37 @@ export class MongoCompactor {
           currentState.opsSincePut = 0;
         } else if (doc.op != 'CLEAR') {
           if (currentState.lastNotPut == null) {
-            currentState.lastNotPut = doc._id.o;
+            currentState.lastNotPut = doc.o;
           }
           currentState.opsSincePut += 1;
         }
 
         if (this.updates.length + this.bucketStateUpdates.length >= this.moveBatchLimit) {
-          await this.flush();
+          await this.flush(bucketContext);
         }
       }
 
-      logger.info(`Processed batch of length ${batch.length} current bucket: ${bucket}`);
+      this.logger.info(`Processed batch of length ${batch.length} current bucket: ${bucket}`);
     }
 
-    // Free memory before clearing bucket
+    // Free memory before clearing the bucket.
     currentState.seen.clear();
     if (currentState.lastNotPut != null && currentState.opsSincePut >= 1) {
-      logger.info(
+      this.logger.info(
         `Inserting CLEAR at ${this.group_id}:${bucket}:${currentState.lastNotPut} to remove ${currentState.opsSincePut} operations`
       );
-      // Need flush() before clear()
-      await this.flush();
-      await this.clearBucket(currentState);
+      // Need flush() before clear().
+      await this.flush(bucketContext);
+      await this.clearBucket(currentState, bucketContext);
     }
 
-    // Do this _after_ clearBucket so that we have accurate counts.
+    // Do this after clearBucket so we have accurate counts.
     this.updateBucketChecksums(currentState);
-
-    // Need another flush after updateBucketChecksums()
-    await this.flush();
+    // Need another flush after updateBucketChecksums().
+    await this.flush(bucketContext);
   }
 
-  /**
-   * Call when done with a bucket.
-   */
-  private updateBucketChecksums(state: CurrentBucketState) {
+  protected updateBucketChecksums(state: CurrentBucketState) {
     if (state.opCount < 0) {
       throw new ServiceAssertionError(
         `Invalid opCount: ${state.opCount} checksum ${state.checksum} opsSincePut: ${state.opsSincePut} maxOpId: ${this.maxOpId}`
@@ -336,12 +496,7 @@ export class MongoCompactor {
     }
     this.bucketStateUpdates.push({
       updateOne: {
-        filter: {
-          _id: {
-            g: this.group_id,
-            b: state.bucket
-          }
-        },
+        filter: this.bucketStateFilter(state.bucket, state.definitionId),
         update: {
           $set: {
             compacted_state: {
@@ -351,14 +506,13 @@ export class MongoCompactor {
               bytes: state.opBytes
             },
             estimate_since_compact: {
-              // Note: 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 could add overhead.
+              // 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
@@ -366,23 +520,24 @@ export class MongoCompactor {
     });
   }
 
-  private async flush() {
+  protected async flush(col: SingleBucketStore) {
     if (this.updates.length > 0) {
-      logger.info(`Compacting ${this.updates.length} ops`);
-      await this.db.bucket_data.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.
+      this.logger.info(`Compacting ${this.updates.length} ops`);
+      await col.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 = [];
     }
+
+    await this.flushBucketStateUpdates();
+  }
+
+  private async flushBucketStateUpdates() {
     if (this.bucketStateUpdates.length > 0) {
-      logger.info(`Updating ${this.bucketStateUpdates.length} bucket states`);
-      await this.db.bucket_state.bulkWrite(this.bucketStateUpdates, {
-        ordered: false
-      });
+      this.logger.info(`Updating ${this.bucketStateUpdates.length} bucket states`);
+      await this.writeBucketStateUpdates();
       this.bucketStateUpdates = [];
     }
   }
@@ -390,26 +545,15 @@ export class MongoCompactor {
   /**
    * Perform a CLEAR compact for a bucket.
    *
-   *
-   * @param bucket bucket name
-   * @param op op_id of the last non-PUT operation, which will be converted to CLEAR.
+   * @param currentState tracks the last non-PUT op, which will be converted to CLEAR.
    */
-  private async clearBucket(currentState: CurrentBucketState) {
-    const bucket = currentState.bucket;
+  protected async clearBucket(currentState: CurrentBucketState, col: SingleBucketStore) {
     const clearOp = currentState.lastNotPut!;
 
     const opFilter = {
       _id: {
-        $gte: {
-          g: this.group_id,
-          b: bucket,
-          o: new mongo.MinKey() as any
-        },
-        $lte: {
-          g: this.group_id,
-          b: bucket,
-          o: clearOp
-        }
+        $gte: col.minId,
+        $lte: col.docId(clearOp)
       }
     };
 
@@ -424,39 +568,40 @@ export class MongoCompactor {
         // We need a transaction per batch to make sure checksums stay consistent.
         await session.withTransaction(
           async () => {
-            const query = this.db.bucket_data.find(opFilter, {
-              session,
-              sort: { _id: 1 },
-              projection: {
-                _id: 1,
-                op: 1,
-                checksum: 1,
-                target_op: 1
-              },
-              limit: this.clearBatchLimit
-            });
+            const query = col.collection.find<Pick<BucketDataDocumentGeneric, '_id' | CompactClearProperties>>(
+              opFilter,
+              {
+                session,
+                sort: { _id: 1 },
+                projection: {
+                  _id: 1,
+                  op: 1,
+                  checksum: 1,
+                  target_op: 1
+                },
+                limit: this.clearBatchLimit
+              }
+            );
             let checksum = 0;
-            let lastOpId: BucketDataKey | null = null;
+            let lastOp: Pick<BucketDataDoc, 'o' | CompactClearProperties> | null = null;
             let targetOp: bigint | null = null;
             let gotAnOp = false;
             let numberOfOpsToClear = 0;
-            for await (let op of query.stream()) {
+            for await (const rawOp of query.stream()) {
+              const op = col.fromPartialPersistedDocument(rawOp);
+
               if (op.op == 'MOVE' || op.op == 'REMOVE' || op.op == 'CLEAR') {
                 checksum = utils.addChecksums(checksum, Number(op.checksum));
-                lastOpId = op._id;
+                lastOp = op;
                 numberOfOpsToClear += 1;
                 if (op.op != 'CLEAR') {
                   gotAnOp = true;
                 }
-                if (op.target_op != null) {
-                  if (targetOp == null || op.target_op > targetOp) {
-                    targetOp = op.target_op;
-                  }
+                if (op.target_op != null && (targetOp == null || op.target_op > targetOp)) {
+                  targetOp = op.target_op;
                 }
               } else {
-                throw new ReplicationAssertionError(
-                  `Unexpected ${op.op} operation at ${op._id.g}:${op._id.b}:${op._id.o}`
-                );
+                throw new ReplicationAssertionError(`Unexpected ${op.op} operation at ${this.formatBucketDataKey(op)}`);
               }
             }
             if (!gotAnOp) {
@@ -464,31 +609,25 @@ export class MongoCompactor {
               return;
             }
 
-            logger.info(`Flushing CLEAR for ${numberOfOpsToClear} ops at ${lastOpId?.o}`);
-            await this.db.bucket_data.deleteMany(
+            this.logger.info(`Flushing CLEAR for ${numberOfOpsToClear} ops at ${lastOp?.o}`);
+            await col.collection.deleteMany(
               {
                 _id: {
-                  $gte: {
-                    g: this.group_id,
-                    b: bucket,
-                    o: new mongo.MinKey() as any
-                  },
-                  $lte: lastOpId!
+                  $gte: col.minId,
+                  $lte: col.docId(lastOp!.o)
                 }
               },
               { session }
             );
 
-            await this.db.bucket_data.insertOne(
-              {
-                _id: lastOpId!,
-                op: 'CLEAR',
-                checksum: BigInt(checksum),
-                data: null,
-                target_op: targetOp
-              },
-              { session }
-            );
+            const op = col.toPersistedDocument({
+              o: lastOp!.o,
+              op: 'CLEAR',
+              checksum: BigInt(checksum),
+              data: null,
+              target_op: targetOp
+            });
+            await col.collection.insertOne(op, { session });
 
             opCountDiff = -numberOfOpsToClear + 1;
           },
@@ -497,7 +636,7 @@ export class MongoCompactor {
             readConcern: { level: 'snapshot' }
           }
         );
-        // Update _outside_ the transaction, since the transaction can be retried multiple times.
+        // Update outside the transaction, since the transaction can be retried multiple times.
         currentState.opCount += opCountDiff;
       }
     } finally {
@@ -505,211 +644,22 @@ export class MongoCompactor {
     }
   }
 
-  /**
-   * Subset of compact, only populating checksums where relevant.
-   */
-  async populateChecksums(options: { minBucketChanges: number }): Promise<PopulateChecksumCacheResults> {
-    let count = 0;
-    while (true) {
-      this.signal?.throwIfAborted();
-      const buckets = await this.dirtyBucketBatchForChecksums(options);
-      if (buckets.length == 0) {
-        // All done
-        break;
-      }
-      this.signal?.throwIfAborted();
-
-      const start = Date.now();
-
-      // Filter batch by estimated bucket size, to reduce possibility of timeouts
-      let checkBuckets: typeof buckets = [];
-      let totalCountEstimate = 0;
-      for (let bucket of buckets) {
-        checkBuckets.push(bucket);
-        totalCountEstimate += bucket.estimatedCount;
-        if (totalCountEstimate > 50_000) {
-          break;
-        }
-      }
-      logger.info(
-        `Calculating checksums for batch of ${buckets.length} buckets, estimated count of ${totalCountEstimate}`
-      );
-      await this.updateChecksumsBatch(checkBuckets.map((b) => b.bucket));
-      logger.info(`Updated checksums for batch of ${checkBuckets.length} buckets in ${Date.now() - start}ms`);
-      count += checkBuckets.length;
-    }
-    return { buckets: count };
-  }
-
-  /**
-   * Return batches of dirty buckets.
-   *
-   * Can be used to iterate through all buckets.
-   *
-   * minBucketChanges: minimum number of changes for a bucket to be included in the results.
-   * minChangeRatio: minimum ratio of changes to total ops for a bucket to be included in the results, number between 0 and 1.
-   */
-  private async *dirtyBucketBatches(options: {
-    minBucketChanges: number;
-    minChangeRatio: number;
-  }): AsyncGenerator<{ bucket: string; estimatedCount: number }[]> {
-    // Previously, we used an index on {_id.g: 1, estimate_since_compact.count: 1} to only buckets with changes.
-    // This works well if there are only a small number of buckets with changes.
-    // However, if buckets are continuosly modified while we are compacting, we get the same buckets over and over again.
-    // This has caused the compact process to re-read the same collection around 5x times in total, which is very inefficient.
-    // To solve this, we now just iterate through all buckets, and filter out the ones with low changes.
-
-    if (options.minBucketChanges <= 0) {
-      throw new ReplicationAssertionError('minBucketChanges must be >= 1');
-    }
-    let lastId = { g: this.group_id, b: new mongo.MinKey() as any };
-    const maxId = { g: this.group_id, b: new mongo.MaxKey() as any };
-    while (true) {
-      // To avoid timeouts from too many buckets not meeting the minBucketChanges criteria, we use an aggregation pipeline
-      // to scan a fixed batch of buckets at a time, but only return buckets that meet the criteria, rather than limiting
-      // on the output number.
-      const [result] = await this.db.bucket_state
-        .aggregate<{
-          buckets: Pick<BucketStateDocument, '_id' | 'estimate_since_compact' | 'compacted_state'>[];
-          cursor: Pick<BucketStateDocument, '_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: {
-                // This is the results for the batch
-                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((b) => {
-        // The numbers, specifically the bytes, could be a bigint. We convert to Number to allow calculating the ratios.
-        // BigInt precision is not needed here since it's just an estimate.
-        const updatedCount = b.estimate_since_compact?.count ?? 0;
-        const totalCount = (b.compacted_state?.count ?? 0) + updatedCount;
-        const updatedBytes = Number(b.estimate_since_compact?.bytes ?? 0);
-        const totalBytes = Number(b.compacted_state?.bytes ?? 0) + updatedBytes;
-        const dirtyChangeNumber = totalCount > 0 ? updatedCount / totalCount : 0;
-        const dirtyChangeBytes = totalBytes > 0 ? updatedBytes / totalBytes : 0;
-        return {
-          bucket: b._id.b,
-          estimatedCount: totalCount,
-          dirtyRatio: Math.max(dirtyChangeNumber, dirtyChangeBytes)
-        };
-      });
-      const filtered = mapped.filter(
-        (b) => b.estimatedCount >= options.minBucketChanges && b.dirtyRatio >= options.minChangeRatio
-      );
-      yield filtered;
-    }
-  }
-
-  /**
-   * Returns a batch of dirty buckets - buckets with most changes first.
-   *
-   * This cannot be used to iterate on its own - the client is expected to process these buckets and
-   * set estimate_since_compact.count: 0 when done, before fetching the next batch.
-   *
-   * Unlike dirtyBucketBatches, used for compacting, this is specifically designed to be resuamble after a restart,
-   * since it is used as the last step for initial replication.
-   *
-   * We currently don't get new data while doing populateChecksums, so we don't need to worry about buckets changing while processing.
-   */
-  private async dirtyBucketBatchForChecksums(options: {
-    minBucketChanges: number;
-  }): Promise<{ bucket: string; estimatedCount: number }[]> {
-    if (options.minBucketChanges <= 0) {
-      throw new ReplicationAssertionError('minBucketChanges must be >= 1');
-    }
-    // We make use of an index on {_id.g: 1, 'estimate_since_compact.count': -1}
-    const dirtyBuckets = await this.db.bucket_state
-      .find(
-        {
-          '_id.g': this.group_id,
-          'estimate_since_compact.count': { $gte: options.minBucketChanges }
-        },
-        {
-          projection: {
-            _id: 1,
-            estimate_since_compact: 1,
-            compacted_state: 1
-          },
-          sort: {
-            'estimate_since_compact.count': -1
-          },
-          limit: 200,
-          maxTimeMS: MONGO_OPERATION_TIMEOUT_MS
-        }
-      )
-      .toArray();
+  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]));
 
-    return dirtyBuckets.map((bucket) => ({
-      bucket: bucket._id.b,
-      estimatedCount: Number(bucket.estimate_since_compact!.count) + Number(bucket.compacted_state?.count ?? 0)
-    }));
-  }
-
-  private async updateChecksumsBatch(buckets: string[]) {
-    const checksums = await this.storage.checksums.computePartialChecksumsDirect(
-      buckets.map((bucket) => {
-        return {
-          bucket,
-          source: {} as any,
-          end: this.maxOpId
-        };
-      })
-    );
-
-    for (let bucketChecksum of checksums.values()) {
+    for (const bucketChecksum of checksums.values()) {
       if (isPartialChecksum(bucketChecksum)) {
-        // Should never happen since we don't specify `start`
+        // Should never happen since we don't specify `start`.
         throw new ServiceAssertionError(`Full checksum expected, got ${JSON.stringify(bucketChecksum)}`);
       }
 
       this.bucketStateUpdates.push({
         updateOne: {
-          filter: {
-            _id: {
-              g: this.group_id,
-              b: bucketChecksum.bucket
-            }
-          },
+          filter: this.bucketStateFilter(
+            bucketChecksum.bucket,
+            definitionIdByBucket.get(bucketChecksum.bucket) ?? null
+          ),
           update: {
             $set: {
               compacted_state: {
@@ -723,14 +673,34 @@ export class MongoCompactor {
                 bytes: 0
               }
             }
-          },
-          // We don't create new ones here - it gets tricky to get the last_op right with the unique index on:
-          // bucket_updates: {'id.g': 1, 'last_op': 1}
+          } 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.flush();
+    await this.flushBucketStateUpdates();
+  }
+
+  protected formatBucketDataKey(doc: Pick<BucketDataDoc, 'bucketKey' | 'o'>) {
+    return `${doc.bucketKey.replicationStreamId}:${doc.bucketKey.bucket}:${doc.o}`;
   }
+
+  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;
+
+  protected abstract getBucketDataContext(
+    bucket: string,
+    definitionId: BucketDefinitionId | null
+  ): Promise<SingleBucketStore | null>;
+}
+
+export interface BucketDataCollectionContext<TBucketData extends mongo.Document> {
+  bucketKey: BucketKey;
+  collection: mongo.Collection<TBucketData>;
 }