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/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

@@ -1,3 +1,4 @@
+import * as lib_mongo from '@powersync/lib-service-mongodb';
 import { mongo } from '@powersync/lib-service-mongodb';
 import { HydratedSyncConfig, SqlEventDescriptor, SqliteRow, SqliteValue } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
@@ -60,6 +61,8 @@ export interface MongoBucketBatchOptions {
 
   logger: Logger;
   tracer?: PerformanceTracer<'storage' | 'evaluate'>;
+
+  listSourceRecordCollections?: (groupId: number) => Promise<mongo.Collection<any>[]>;
 }
 
 export abstract class MongoBucketBatch
@@ -70,7 +73,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;
 
@@ -95,6 +98,7 @@ export abstract class MongoBucketBatch
   private markRecordUnavailable: BucketStorageMarkRecordUnavailable | undefined;
   private hooks: storage.StorageHooks | undefined;
   private clearedError = false;
+  private listSourceRecordCollections?: (groupId: number) => Promise<mongo.Collection<any>[]>;
 
   private tracer: PerformanceTracer<'storage' | 'evaluate'>;
 
@@ -142,6 +146,7 @@ export abstract class MongoBucketBatch
     this.mapping = options.mapping;
     this.skipExistingRows = options.skipExistingRows;
     this.markRecordUnavailable = options.markRecordUnavailable;
+    this.listSourceRecordCollections = options.listSourceRecordCollections;
     this.hooks = options.hooks;
     this.batch = new OperationBatch();
 
@@ -166,7 +171,28 @@ export abstract class MongoBucketBatch
 
   protected abstract get sourceRecordStore(): SourceRecordStore;
 
-  protected abstract cleanupDroppedSourceTables(sourceTables: storage.SourceTable[]): Promise<void>;
+  protected async cleanupDroppedSourceTables(sourceTables: storage.SourceTable[]) {
+    if (this.listSourceRecordCollections == null) {
+      return;
+    }
+    const collections = await this.listSourceRecordCollections(this.replicationStreamId);
+    const tableIds = new Set(sourceTables.map((t) => mongoTableId(t.id).toHexString()));
+    for (const collection of collections) {
+      const name = collection.collectionName;
+      const prefix = `source_records_${this.replicationStreamId}_`;
+      if (name.startsWith(prefix)) {
+        const tableId = name.slice(prefix.length);
+        if (tableIds.has(tableId)) {
+          await collection.drop().catch((error) => {
+            if (lib_mongo.isMongoServerError(error) && error.codeName === 'NamespaceNotFound') {
+              return;
+            }
+            throw error;
+          });
+        }
+      }
+    }
+  }
 
   abstract commit(lsn: string, options?: storage.BucketBatchCommitOptions): Promise<storage.CheckpointResult>;
 

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) {