feat(mongodb-storage)!: chunked multi-op bucket documents with range-merging compaction and invariant tests

.changeset/wild-pears-sing.md

@@ -0,0 +1,8 @@
+---
+'@powersync/service-module-mongodb-storage': minor
+'@powersync/service-core': minor
+---
+
+Introduce chunked multi-op bucket documents with invariants and read-filtering tests in MongoDB storage.
+
+Streaming V3 compactor — process buckets incrementally with byte-bounded batches, scoped deletes, and bounded transactions. Add `moveBatchByteLimit` option to `CompactOptions`.

modules/module-mongodb-storage/src/storage/implementation/BucketDefinitionMapping.ts

@@ -8,6 +8,8 @@ import {
 } from '@powersync/service-sync-rules';
 import { SyncConfigDefinition } from '../storage-index.js';
 
+export type { BucketDefinitionId, ParameterIndexId };
+
 export interface SyncConfigWithMapping {
   syncConfigId?: string;
   syncConfig: SyncConfigWithErrors;

modules/module-mongodb-storage/src/storage/implementation/MongoBucketBatch.ts

@@ -70,7 +70,7 @@ export abstract class MongoBucketBatch
   protected logger: Logger;
 
   private readonly client: mongo.MongoClient;
-  public readonly db: VersionedPowerSyncMongo;
+  readonly db: VersionedPowerSyncMongo;
   public readonly session: mongo.ClientSession;
   protected readonly sync_rules: HydratedSyncConfig;
 

modules/module-mongodb-storage/src/storage/implementation/MongoChecksums.ts

@@ -7,19 +7,17 @@ import {
   ChecksumMap,
   FetchPartialBucketChecksum,
   InternalOpId,
-  isPartialChecksum,
   PartialChecksum,
   PartialChecksumMap,
   PartialOrFullChecksum
 } from '@powersync/service-core';
 import type { VersionedPowerSyncMongo } from './db.js';
 
-import * as lib_mongo from '@powersync/lib-service-mongodb';
 import { BucketDefinitionId } from '@powersync/service-sync-rules';
 import { BucketDefinitionMapping } from './BucketDefinitionMapping.js';
-import { BucketDataDocumentBase, StorageConfig } from './models.js';
+import { StorageConfig } from './models.js';
 
-export interface FetchPartialBucketChecksumV3 {
+export interface FetchPartialBucketChecksumByDefinition {
   bucket: string;
   definitionId: BucketDefinitionId;
   start?: InternalOpId;
@@ -51,11 +49,11 @@ export interface MongoChecksumOptions {
 }
 
 const DEFAULT_BUCKET_BATCH_LIMIT = 200;
-const DEFAULT_OPERATION_BATCH_LIMIT = 50_000;
+export const DEFAULT_OPERATION_BATCH_LIMIT = 50_000;
 
 export abstract class MongoChecksums {
   private _cache: ChecksumCache | undefined;
-  private readonly storageConfig: StorageConfig;
+  protected readonly storageConfig: StorageConfig;
 
   constructor(
     protected readonly db: VersionedPowerSyncMongo,
@@ -178,138 +176,10 @@ export abstract class MongoChecksums {
   protected abstract fetchPreStates(
     batch: FetchPartialBucketChecksum[]
   ): Promise<Map<string, { opId: InternalOpId; checksum: BucketChecksum }>>;
-
-  protected async computePartialChecksumsForCollection<
-    TRequest extends FetchPartialBucketChecksumByBucket,
-    TBucketDataDocument extends BucketDataDocumentBase
-  >(
-    batch: TRequest[],
-    collection: lib_mongo.mongo.Collection<TBucketDataDocument>,
-    createFilter: (request: TRequest) => any
-  ): Promise<PartialChecksumMap> {
-    const batchLimit = this.options?.operationBatchLimit ?? DEFAULT_OPERATION_BATCH_LIMIT;
-
-    // Map requests by bucket. We adjust this as we get partial results.
-    let requests = new Map<string, TRequest>();
-    for (let request of batch) {
-      requests.set(request.bucket, request);
-    }
-
-    const partialChecksums = new Map<string, PartialOrFullChecksum>();
-
-    while (requests.size > 0) {
-      const filters = Array.from(requests.values(), createFilter);
-
-      // Historically, checksum may be stored as 'int' or 'double'.
-      // More recently, this should be a 'long'.
-      // $toLong ensures that we always sum it as a long, avoiding inaccuracies in the calculations.
-      const checksumLong = this.storageConfig.longChecksums ? '$checksum' : { $toLong: '$checksum' };
-
-      // Aggregate over a max of `batchLimit` operations at a time.
-      // Let's say we have 3 buckets (A, B, C), each with 10 operations, and our batch limit is 12.
-      // Then we'll do three batches:
-      // 1. Query: A[1-end], B[1-end], C[1-end]
-      //    Returns: A[1-10], B[1-2]
-      // 2. Query: B[3-end], C[1-end]
-      //    Returns: B[3-10], C[1-4]
-      // 3. Query: C[5-end]
-      //    Returns: C[5-10]
-      const aggregate = await collection
-        .aggregate(
-          [
-            {
-              $match: {
-                $or: filters
-              }
-            },
-            // sort and limit _before_ grouping
-            { $sort: { _id: 1 } },
-            { $limit: batchLimit },
-            {
-              $group: {
-                _id: '$_id.b',
-                checksum_total: { $sum: checksumLong },
-                count: { $sum: 1 },
-                has_clear_op: {
-                  $max: {
-                    $cond: [{ $eq: ['$op', 'CLEAR'] }, 1, 0]
-                  }
-                },
-                last_op: { $max: '$_id.o' }
-              }
-            },
-            // Sort the aggregated results (100 max, so should be fast).
-            // This is important to identify which buckets we have partial data for.
-            { $sort: { _id: 1 } }
-          ],
-          { session: undefined, readConcern: 'snapshot', maxTimeMS: lib_mongo.MONGO_CHECKSUM_TIMEOUT_MS }
-        )
-        .toArray()
-        .catch((e) => {
-          throw lib_mongo.mapQueryError(e, 'while reading checksums');
-        });
-
-      let batchCount = 0;
-      let limitReached = false;
-      for (let doc of aggregate) {
-        const bucket = doc._id;
-        const checksum = checksumFromAggregate(doc);
-
-        const existing = partialChecksums.get(bucket);
-        if (existing != null) {
-          partialChecksums.set(bucket, addPartialChecksums(bucket, existing, checksum));
-        } else {
-          partialChecksums.set(bucket, checksum);
-        }
-
-        batchCount += doc.count;
-        if (batchCount == batchLimit) {
-          // Limit reached. Request more in the next batch.
-          // Note that this only affects the _last_ bucket in a batch.
-          limitReached = true;
-          const req = requests.get(bucket);
-          requests.set(bucket, {
-            ...req!,
-            start: doc.last_op
-          });
-        } else {
-          // All done for this bucket
-          requests.delete(bucket);
-        }
-      }
-      if (!limitReached) {
-        break;
-      }
-    }
-
-    return new Map<string, PartialOrFullChecksum>(
-      batch.map((request) => {
-        const bucket = request.bucket;
-        // Could be null if we got no data
-        let partialChecksum = partialChecksums.get(bucket);
-        if (partialChecksum == null) {
-          partialChecksum = {
-            bucket,
-            partialCount: 0,
-            partialChecksum: 0
-          };
-        }
-        if (request.start == null && isPartialChecksum(partialChecksum)) {
-          partialChecksum = {
-            bucket,
-            count: partialChecksum.partialCount,
-            checksum: partialChecksum.partialChecksum
-          };
-        }
-
-        return [bucket, partialChecksum];
-      })
-    );
-  }
 }
 
 export function emptyChecksumForRequest(
-  request: Pick<FetchPartialBucketChecksum | FetchPartialBucketChecksumV3, 'bucket' | 'start'>
+  request: Pick<FetchPartialBucketChecksum | FetchPartialBucketChecksumByDefinition, 'bucket' | 'start'>
 ): PartialOrFullChecksum {
   if (request.start == null) {
     return { bucket: request.bucket, count: 0, checksum: 0 };
@@ -320,7 +190,7 @@ export function emptyChecksumForRequest(
 /**
  * Convert output of the $group stage into a checksum.
  */
-function checksumFromAggregate(doc: bson.Document): PartialOrFullChecksum {
+export function checksumFromAggregate(doc: bson.Document): PartialOrFullChecksum {
   const partialChecksum = Number(BigInt(doc.checksum_total) & 0xffffffffn) & 0xffffffff;
   const bucket = doc._id;
 

modules/module-mongodb-storage/src/storage/implementation/MongoCompactor.ts

@@ -63,6 +63,7 @@ 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 = 64 * 1024 * 1024;
 const DEFAULT_MIN_BUCKET_CHANGES = 10;
 const DEFAULT_MIN_CHANGE_RATIO = 0.1;
 const DIRTY_BUCKET_SCAN_BATCH_SIZE = 2_000;
@@ -83,6 +84,7 @@ export abstract class MongoCompactor {
   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;
@@ -102,6 +104,7 @@ export abstract class MongoCompactor {
     this.idLimitBytes = (options.memoryLimitMB ?? DEFAULT_MEMORY_LIMIT_MB) * 1024 * 1024;
     this.moveBatchLimit = options.moveBatchLimit ?? DEFAULT_MOVE_BATCH_LIMIT;
     this.moveBatchQueryLimit = options.moveBatchQueryLimit ?? DEFAULT_MOVE_BATCH_QUERY_LIMIT;
+    this.moveBatchByteLimit = options.moveBatchByteLimit ?? DEFAULT_MOVE_BATCH_BYTE_LIMIT;
     this.clearBatchLimit = options.clearBatchLimit ?? DEFAULT_CLEAR_BATCH_LIMIT;
     this.minBucketChanges = options.minBucketChanges ?? DEFAULT_MIN_BUCKET_CHANGES;
     this.minChangeRatio = options.minChangeRatio ?? DEFAULT_MIN_CHANGE_RATIO;
@@ -133,6 +136,7 @@ export abstract class MongoCompactor {
    */
   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);
@@ -172,6 +176,7 @@ export abstract class MongoCompactor {
     },
     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.
@@ -304,6 +309,7 @@ export abstract class MongoCompactor {
    */
   protected async compactSingleBucketRetried(bucket: string, definitionId: BucketDefinitionId | null = null) {
     let retryCount = 0;
+    // Retry with exponential backoff up to 3 times on MongoDB errors.
     while (true) {
       try {
         await this.compactSingleBucket(bucket, definitionId);
@@ -343,6 +349,7 @@ export abstract class MongoCompactor {
     // Upper bound is adjusted for each batch.
     let upperBound = bucketContext.maxId;
 
+    // Paginate through bucket data in batches to avoid cursor timeouts.
     while (true) {
       this.signal?.throwIfAborted();
 
@@ -488,13 +495,13 @@ export abstract class MongoCompactor {
     await this.flush(bucketContext);
   }
 
-  protected updateBucketChecksums(state: CurrentBucketState) {
+  protected collectBucketStateUpdates(state: CurrentBucketState): mongo.AnyBulkWriteOperation<BucketStateDocumentBase> {
     if (state.opCount < 0) {
       throw new ServiceAssertionError(
         `Invalid opCount: ${state.opCount} checksum ${state.checksum} opsSincePut: ${state.opsSincePut} maxOpId: ${this.maxOpId}`
       );
     }
-    this.bucketStateUpdates.push({
+    return {
       updateOne: {
         filter: this.bucketStateFilter(state.bucket, state.definitionId),
         update: {
@@ -517,7 +524,11 @@ export abstract class MongoCompactor {
         // We don't create new ones here, to avoid issues with the unique index on bucket_updates.
         upsert: false
       }
-    });
+    };
+  }
+
+  protected updateBucketChecksums(state: CurrentBucketState) {
+    this.bucketStateUpdates.push(this.collectBucketStateUpdates(state));
   }
 
   protected async flush(col: SingleBucketStore) {

modules/module-mongodb-storage/src/storage/implementation/MongoParameterCompactor.ts

@@ -18,12 +18,13 @@ type ParameterCompactionReadDocument = {
  *
  * For background, see the `/docs/parameters-lookups.md` file.
  */
-export abstract class MongoParameterCompactor {
+export class MongoParameterCompactor {
   constructor(
     protected readonly db: VersionedPowerSyncMongo,
     protected readonly group_id: number,
     protected readonly checkpoint: InternalOpId,
-    protected readonly options: CompactOptions
+    protected readonly options: CompactOptions,
+    protected readonly getCollectionsCb?: () => Promise<mongo.Collection<mongo.Document>[]>
   ) {}
 
   async compact() {
@@ -33,11 +34,27 @@ export abstract class MongoParameterCompactor {
     }
   }
 
-  protected abstract getCollections(): Promise<mongo.Collection<mongo.Document>[]>;
+  protected async getCollections(): Promise<mongo.Collection<mongo.Document>[]> {
+    if (this.getCollectionsCb == null) {
+      throw new Error('getCollections callback not provided');
+    }
+    const collections = await this.getCollectionsCb();
+    // Cast from the version-specific collection type to the generic Document type
+    // used by the parameter compactor base class.
+    return collections.map((collection) => collection as unknown as mongo.Collection<mongo.Document>);
+  }
 
-  protected abstract collectionFilter(): mongo.Document;
+  protected collectionFilter(): mongo.Document {
+    return {};
+  }
 
-  protected abstract deleteFilter(doc: mongo.Document): mongo.Document;
+  protected deleteFilter(doc: mongo.Document): mongo.Document {
+    return {
+      lookup: doc.lookup,
+      _id: { $lte: doc._id },
+      key: doc.key
+    };
+  }
 
   protected async compactCollection(collection: mongo.Collection<mongo.Document>) {
     // This is the currently-active checkpoint.
@@ -68,6 +85,7 @@ export abstract class MongoParameterCompactor {
 
     const flush = async (force: boolean) => {
       if (removeIds.length >= 1000 || (force && removeIds.length > 0)) {
+        // MongoDB Filter<T> doesn't fully match our dynamic delete filter shape here.
         const results = await collection.deleteMany({ _id: { $in: removeIds } } as any);
         logger.info(`Removed ${results.deletedCount} (${removeIds.length}) superseded parameter entries`);
         removeIds = [];
@@ -81,6 +99,7 @@ export abstract class MongoParameterCompactor {
     };
 
     while (await cursor.hasNext()) {
+      // readBufferedDocuments returns a generic type; we know the shape from our projection.
       const batch = cursor.readBufferedDocuments() as unknown as ParameterCompactionReadDocument[];
       checkedEntries += batch.length;
       const now = Date.now();

modules/module-mongodb-storage/src/storage/implementation/MongoSyncBucketStorage.ts

@@ -71,6 +71,7 @@ export abstract class MongoSyncBucketStorage
   implements storage.SyncRulesBucketStorage
 {
   readonly db: VersionedPowerSyncMongo;
+
   [DO_NOT_LOG] = true;
 
   readonly checksums: MongoChecksums;
@@ -621,7 +622,7 @@ class EmptyReplicationCheckpoint implements ReplicationCheckpoint {
   readonly checkpoint: InternalOpId = 0n;
   readonly lsn: string | null = null;
 
-  async getParameterSets(_lookups: ScopedParameterLookup[]): Promise<ParameterLookupRows[]> {
+  async getParameterSets(_lookups: ScopedParameterLookup[], _limit: number): Promise<ParameterLookupRows[]> {
     return [];
   }
 }

modules/module-mongodb-storage/src/storage/implementation/common/VersionedPowerSyncMongoBase.ts

@@ -67,7 +67,7 @@ export abstract class BaseVersionedPowerSyncMongo {
   }
 
   protected async listCollectionsByPrefix<T extends mongo.Document>(prefix: string): Promise<mongo.Collection<T>[]> {
-    const collections = await this.db.listCollections({ name: new RegExp(`^${prefix}`) }, { nameOnly: true }).toArray();
+    const collections = await this.db.listCollections({ name: { $regex: `^${prefix}` } }, { nameOnly: true }).toArray();
 
     return collections
       .filter((collection) => collection.name.startsWith(prefix))