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

package/dist/storage/implementation/MongoCompactor.js

@@ -1,5 +1,5 @@
-import { isMongoServerError, mongo, MONGO_OPERATION_TIMEOUT_MS } from '@powersync/lib-service-mongodb';
-import { logger, ReplicationAssertionError, ServiceAssertionError } from '@powersync/lib-services-framework';
+import { isMongoServerError, MONGO_OPERATION_TIMEOUT_MS } from '@powersync/lib-service-mongodb';
+import { logger as defaultLogger, ReplicationAssertionError, ServiceAssertionError } from '@powersync/lib-services-framework';
 import { addChecksums, isPartialChecksum, utils } from '@powersync/service-core';
 import { cacheKey } from './OperationBatch.js';
 const DEFAULT_CLEAR_BATCH_LIMIT = 5000;
@@ -25,6 +25,7 @@ export class MongoCompactor {
     buckets;
     signal;
     group_id;
+    logger;
     constructor(storage, db, options) {
         this.storage = storage;
         this.db = db;
@@ -38,6 +39,7 @@ export class MongoCompactor {
         this.maxOpId = options.maxOpId ?? 0n;
         this.buckets = options.compactBuckets;
         this.signal = options.signal;
+        this.logger = options.logger ?? defaultLogger;
     }
     /**
      * Compact buckets by converting operations into MOVE and/or CLEAR operations.
@@ -46,9 +48,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);
             }
@@ -57,8 +58,123 @@ export class MongoCompactor {
             await this.compactDirtyBuckets();
         }
     }
+    /**
+     * Subset of compact, only populating checksums where relevant.
+     */
+    async populateChecksums(options) {
+        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 = [];
+            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 };
+    }
+    async *dirtyBucketBatchesForCollection(collection, lastId, maxId, options, getDefinitionId) {
+        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([
+                {
+                    $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);
+        }
+    }
+    async dirtyBucketBatchForChecksumsForCollection(collection, filter, getDefinitionId) {
+        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)
+        }));
+    }
     async compactDirtyBuckets() {
-        for await (let buckets of this.dirtyBucketBatches({
+        for await (const buckets of this.dirtyBucketBatches({
             minBucketChanges: this.minBucketChanges,
             minChangeRatio: this.minChangeRatio
         })) {
@@ -66,8 +182,8 @@ export class MongoCompactor {
             if (buckets.length == 0) {
                 continue;
             }
-            for (let { bucket } of buckets) {
-                await this.compactSingleBucketRetried(bucket);
+            for (const { bucket, definitionId } of buckets) {
+                await this.compactSingleBucketRetried(bucket, definitionId);
             }
         }
     }
@@ -76,16 +192,16 @@ export class MongoCompactor {
      *
      * This covers against occasional network or other database errors during a long compact job.
      */
-    async compactSingleBucketRetried(bucket) {
+    async compactSingleBucketRetried(bucket, definitionId = 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));
                 }
@@ -95,10 +211,15 @@ export class MongoCompactor {
             }
         }
     }
-    async compactSingleBucket(bucket) {
+    async compactSingleBucket(bucket, definitionId = null) {
         const idLimitBytes = this.idLimitBytes;
-        let currentState = {
+        const bucketContext = await this.getBucketDataContext(bucket, definitionId);
+        if (bucketContext == null) {
+            return;
+        }
+        const currentState = {
             bucket,
+            definitionId: bucketContext.key.definitionId,
             seen: new Map(),
             trackingSize: 0,
             lastNotPut: null,
@@ -107,28 +228,24 @@ export class MongoCompactor {
             opCount: 0,
             opBytes: 0
         };
-        // Constant lower bound
-        const lowerBound = {
-            g: this.group_id,
-            b: bucket,
-            o: new mongo.MinKey()
-        };
-        // Upper bound is adjusted for each batch
-        let upperBound = {
-            g: this.group_id,
-            b: bucket,
-            o: new mongo.MaxKey()
-        };
+        // 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([
+            // 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 } },
@@ -145,39 +262,45 @@ export class MongoCompactor {
                         size: { $bsonSize: '$$ROOT' }
                     }
                 }
-            ], {
+            ];
+            const cursor = bucketContext.collection.aggregate(pipeline, {
                 // batchSize is 1 more than limit to auto-close the cursor.
                 // See https://github.com/mongodb/node-mongodb-native/pull/4580
                 batchSize: this.moveBatchQueryLimit + 1
             });
             // 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;
-            for (let doc of batch) {
-                if (doc._id.o > this.maxOpId) {
+            // Reuse the exact collection _id value from Mongo for the next bound.
+            upperBound = rawBatch[rawBatch.length - 1]._id;
+            for (const { doc, size } of batch) {
+                if (doc.o > this.maxOpId) {
                     continue;
                 }
                 currentState.checksum = addChecksums(currentState.checksum, Number(doc.checksum));
                 currentState.opCount += 1;
                 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',
@@ -193,22 +316,17 @@ export class MongoCompactor {
                                 }
                             }
                         });
-                        currentState.opBytes += 200 - Number(doc.size); // TODO: better estimate for this
+                        // TODO: better estimate for this.
+                        currentState.opBytes += 200 - Number(size);
                     }
-                    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;
-                        }
+                    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;
                     }
                 }
                 if (isPersistentPut) {
@@ -217,44 +335,36 @@ export class MongoCompactor {
                 }
                 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(`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);
+            this.logger.info(`Inserting CLEAR at ${this.group_id}:${bucket}:${currentState.lastNotPut} to remove ${currentState.opsSincePut} operations`);
+            // 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.
-     */
     updateBucketChecksums(state) {
         if (state.opCount < 0) {
             throw new ServiceAssertionError(`Invalid opCount: ${state.opCount} checksum ${state.checksum} opsSincePut: ${state.opsSincePut} maxOpId: ${this.maxOpId}`);
         }
         this.bucketStateUpdates.push({
             updateOne: {
-                filter: {
-                    _id: {
-                        g: this.group_id,
-                        b: state.bucket
-                    }
-                },
+                filter: this.bucketStateFilter(state.bucket, state.definitionId),
                 update: {
                     $set: {
                         compacted_state: {
@@ -264,9 +374,8 @@ 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
                         }
@@ -278,48 +387,36 @@ export class MongoCompactor {
             }
         });
     }
-    async flush() {
+    async flush(col) {
         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();
+    }
+    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 = [];
         }
     }
     /**
      * 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.
      */
-    async clearBucket(currentState) {
-        const bucket = currentState.bucket;
+    async clearBucket(currentState, col) {
         const clearOp = currentState.lastNotPut;
         const opFilter = {
             _id: {
-                $gte: {
-                    g: this.group_id,
-                    b: bucket,
-                    o: new mongo.MinKey()
-                },
-                $lte: {
-                    g: this.group_id,
-                    b: bucket,
-                    o: clearOp
-                }
+                $gte: col.minId,
+                $lte: col.docId(clearOp)
             }
         };
         const session = this.db.client.startSession();
@@ -332,7 +429,7 @@ export class MongoCompactor {
                 // The state after each batch is fully consistent.
                 // We need a transaction per batch to make sure checksums stay consistent.
                 await session.withTransaction(async () => {
-                    const query = this.db.bucket_data.find(opFilter, {
+                    const query = col.collection.find(opFilter, {
                         session,
                         sort: { _id: 1 },
                         projection: {
@@ -344,56 +441,52 @@ export class MongoCompactor {
                         limit: this.clearBatchLimit
                     });
                     let checksum = 0;
-                    let lastOpId = null;
+                    let lastOp = null;
                     let targetOp = 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) {
                         done = true;
                         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()
-                            },
-                            $lte: lastOpId
+                            $gte: col.minId,
+                            $lte: col.docId(lastOp.o)
                         }
                     }, { session });
-                    await this.db.bucket_data.insertOne({
-                        _id: lastOpId,
+                    const op = col.toPersistedDocument({
+                        o: lastOp.o,
                         op: 'CLEAR',
                         checksum: BigInt(checksum),
                         data: null,
                         target_op: targetOp
-                    }, { session });
+                    });
+                    await col.collection.insertOne(op, { session });
                     opCountDiff = -numberOfOpsToClear + 1;
                 }, {
                     writeConcern: { w: 'majority' },
                     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;
             }
         }
@@ -401,180 +494,17 @@ export class MongoCompactor {
             await session.endSession();
         }
     }
-    /**
-     * Subset of compact, only populating checksums where relevant.
-     */
-    async populateChecksums(options) {
-        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 = [];
-            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.
-     */
-    async *dirtyBucketBatches(options) {
-        // 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() };
-        const maxId = { g: this.group_id, b: new mongo.MaxKey() };
-        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([
-                {
-                    $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.
-     */
-    async dirtyBucketBatchForChecksums(options) {
-        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();
-        return dirtyBuckets.map((bucket) => ({
-            bucket: bucket._id.b,
-            estimatedCount: Number(bucket.estimate_since_compact.count) + Number(bucket.compacted_state?.count ?? 0)
-        }));
-    }
     async updateChecksumsBatch(buckets) {
-        const checksums = await this.storage.checksums.computePartialChecksumsDirect(buckets.map((bucket) => {
-            return {
-                bucket,
-                source: {},
-                end: this.maxOpId
-            };
-        }));
-        for (let bucketChecksum of checksums.values()) {
+        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`
+                // 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: {
@@ -589,13 +519,16 @@ export class MongoCompactor {
                             }
                         }
                     },
-                    // 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}
+                    // 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();
+    }
+    formatBucketDataKey(doc) {
+        return `${doc.bucketKey.replicationStreamId}:${doc.bucketKey.bucket}:${doc.o}`;
     }
 }
 //# sourceMappingURL=MongoCompactor.js.map