@powersync/service-module-mongodb-storage 0.12.0 → 0.12.2

package/src/storage/implementation/MongoCompactor.ts

@@ -1,11 +1,11 @@
-import { mongo } from '@powersync/lib-service-mongodb';
+import { mongo, MONGO_OPERATION_TIMEOUT_MS } from '@powersync/lib-service-mongodb';
 import { logger, ReplicationAssertionError, ServiceAssertionError } from '@powersync/lib-services-framework';
-import { addChecksums, InternalOpId, storage, utils } from '@powersync/service-core';
+import { addChecksums, InternalOpId, isPartialChecksum, storage, utils } from '@powersync/service-core';
 
 import { PowerSyncMongo } from './db.js';
 import { BucketDataDocument, BucketDataKey, BucketStateDocument } from './models.js';
+import { MongoSyncBucketStorage } from './MongoSyncBucketStorage.js';
 import { cacheKey } from './OperationBatch.js';
-import { readSingleBatch } from './util.js';
 
 interface CurrentBucketState {
   /** Bucket name */
@@ -68,12 +68,14 @@ export class MongoCompactor {
   private maxOpId: bigint;
   private buckets: string[] | undefined;
   private signal?: AbortSignal;
+  private group_id: number;
 
   constructor(
+    private storage: MongoSyncBucketStorage,
     private db: PowerSyncMongo,
-    private group_id: number,
     options?: MongoCompactOptions
   ) {
+    this.group_id = storage.group_id;
     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;
@@ -136,33 +138,57 @@ export class MongoCompactor {
       o: new mongo.MaxKey() as any
     };
 
+    const doneWithBucket = async () => {
+      if (currentState == null) {
+        return;
+      }
+      // Free memory before clearing bucket
+      currentState.seen.clear();
+      if (currentState.lastNotPut != null && currentState.opsSincePut >= 1) {
+        logger.info(
+          `Inserting CLEAR at ${this.group_id}:${currentState.bucket}:${currentState.lastNotPut} to remove ${currentState.opsSincePut} operations`
+        );
+        // Need flush() before clear()
+        await this.flush();
+        await this.clearBucket(currentState);
+      }
+
+      // Do this _after_ clearBucket so that we have accurate counts.
+      this.updateBucketChecksums(currentState);
+    };
+
     while (!this.signal?.aborted) {
       // 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
+      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' }
             }
           }
-        },
-        { $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 { data: batch } = await readSingleBatch(cursor);
+        ],
+        { batchSize: this.moveBatchQueryLimit }
+      );
+      // 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();
 
       if (batch.length == 0) {
         // We've reached the end
@@ -174,24 +200,8 @@ export class MongoCompactor {
 
       for (let doc of batch) {
         if (currentState == null || doc._id.b != currentState.bucket) {
-          if (currentState != null) {
-            if (currentState.lastNotPut != null && currentState.opsSincePut >= 1) {
-              // Important to flush before clearBucket()
-              // Does not have to happen before flushBucketChecksums()
-              await this.flush();
-              logger.info(
-                `Inserting CLEAR at ${this.group_id}:${currentState.bucket}:${currentState.lastNotPut} to remove ${currentState.opsSincePut} operations`
-              );
-
-              // Free memory before clearing bucket
-              currentState!.seen.clear();
-
-              await this.clearBucket(currentState);
-            }
+          await doneWithBucket();
 
-            // Should happen after clearBucket() for accurate stats
-            this.updateBucketChecksums(currentState);
-          }
           currentState = {
             bucket: doc._id.b,
             seen: new Map(),
@@ -274,21 +284,14 @@ export class MongoCompactor {
           await this.flush();
         }
       }
-    }
 
-    currentState?.seen.clear();
-    if (currentState?.lastNotPut != null && currentState?.opsSincePut > 1) {
-      logger.info(
-        `Inserting CLEAR at ${this.group_id}:${currentState.bucket}:${currentState.lastNotPut} to remove ${currentState.opsSincePut} operations`
-      );
-      // Need flush() before clear()
-      await this.flush();
-      await this.clearBucket(currentState);
-    }
-    if (currentState != null) {
-      // Do this _after_ clearBucket so that we have accurate counts.
-      this.updateBucketChecksums(currentState);
+      if (currentState != null) {
+        logger.info(`Processed batch of length ${batch.length} current bucket: ${currentState.bucket}`);
+      }
     }
+
+    await doneWithBucket();
+
     // Need another flush after updateBucketChecksums()
     await this.flush();
   }
@@ -325,15 +328,11 @@ export class MongoCompactor {
               count: 0,
               bytes: 0
             }
-          },
-          $setOnInsert: {
-            // Only set this if we're creating the document.
-            // In all other cases, the replication process will have a set a more accurate id.
-            last_op: this.maxOpId
           }
         },
-        // We generally expect this to have been created before, but do handle cases of old unchanged buckets
-        upsert: true
+        // 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
       }
     });
   }
@@ -475,4 +474,97 @@ export class MongoCompactor {
       await session.endSession();
     }
   }
+
+  /**
+   * Subset of compact, only populating checksums where relevant.
+   */
+  async populateChecksums() {
+    // This is updated after each batch
+    let lowerBound: BucketStateDocument['_id'] = {
+      g: this.group_id,
+      b: new mongo.MinKey() as any
+    };
+    // This is static
+    const upperBound: BucketStateDocument['_id'] = {
+      g: this.group_id,
+      b: new mongo.MaxKey() as any
+    };
+    while (!this.signal?.aborted) {
+      // By filtering buckets, we effectively make this "resumeable".
+      const filter: mongo.Filter<BucketStateDocument> = {
+        _id: {
+          $gt: lowerBound,
+          $lt: upperBound
+        },
+        compacted_state: { $exists: false }
+      };
+
+      const bucketsWithoutChecksums = await this.db.bucket_state
+        .find(filter, {
+          projection: {
+            _id: 1
+          },
+          sort: {
+            _id: 1
+          },
+          limit: 5_000,
+          maxTimeMS: MONGO_OPERATION_TIMEOUT_MS
+        })
+        .toArray();
+      if (bucketsWithoutChecksums.length == 0) {
+        // All done
+        break;
+      }
+
+      logger.info(`Calculating checksums for batch of ${bucketsWithoutChecksums.length} buckets`);
+
+      await this.updateChecksumsBatch(bucketsWithoutChecksums.map((b) => b._id.b));
+
+      lowerBound = bucketsWithoutChecksums[bucketsWithoutChecksums.length - 1]._id;
+    }
+  }
+
+  private async updateChecksumsBatch(buckets: string[]) {
+    const checksums = await this.storage.checksums.computePartialChecksumsDirect(
+      buckets.map((bucket) => {
+        return {
+          bucket,
+          end: this.maxOpId
+        };
+      })
+    );
+
+    for (let bucketChecksum of checksums.values()) {
+      if (isPartialChecksum(bucketChecksum)) {
+        // Should never happen since we don't specify `start`
+        throw new ServiceAssertionError(`Full checksum expected, got ${JSON.stringify(bucketChecksum)}`);
+      }
+
+      this.bucketStateUpdates.push({
+        updateOne: {
+          filter: {
+            _id: {
+              g: this.group_id,
+              b: bucketChecksum.bucket
+            }
+          },
+          update: {
+            $set: {
+              compacted_state: {
+                op_id: this.maxOpId,
+                count: bucketChecksum.count,
+                checksum: BigInt(bucketChecksum.checksum),
+                bytes: null
+              }
+            }
+          },
+          // 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}
+          upsert: false
+        }
+      });
+    }
+
+    await this.flush();
+  }
 }

package/src/storage/implementation/MongoSyncBucketStorage.ts

@@ -7,9 +7,7 @@ import {
   ServiceAssertionError
 } from '@powersync/lib-services-framework';
 import {
-  addPartialChecksums,
   BroadcastIterable,
-  BucketChecksum,
   CHECKPOINT_INVALIDATE_ALL,
   CheckpointChanges,
   CompactOptions,
@@ -18,7 +16,6 @@ import {
   InternalOpId,
   internalToExternalOpId,
   maxLsn,
-  PartialChecksum,
   ProtocolOpId,
   ReplicationCheckpoint,
   storage,
@@ -34,21 +31,22 @@ import { MongoBucketStorage } from '../MongoBucketStorage.js';
 import { PowerSyncMongo } from './db.js';
 import { BucketDataDocument, BucketDataKey, BucketStateDocument, SourceKey, SourceTableDocument } from './models.js';
 import { MongoBucketBatch } from './MongoBucketBatch.js';
+import { MongoChecksumOptions, MongoChecksums } from './MongoChecksums.js';
 import { MongoCompactor } from './MongoCompactor.js';
 import { MongoParameterCompactor } from './MongoParameterCompactor.js';
 import { MongoWriteCheckpointAPI } from './MongoWriteCheckpointAPI.js';
 import { idPrefixFilter, mapOpEntry, readSingleBatch, setSessionSnapshotTime } from './util.js';
 
+export interface MongoSyncBucketStorageOptions {
+  checksumOptions?: MongoChecksumOptions;
+}
+
 export class MongoSyncBucketStorage
   extends BaseObserver<storage.SyncRulesBucketStorageListener>
   implements storage.SyncRulesBucketStorage
 {
   private readonly db: PowerSyncMongo;
-  private checksumCache = new storage.ChecksumCache({
-    fetchChecksums: (batch) => {
-      return this.getChecksumsInternal(batch);
-    }
-  });
+  readonly checksums: MongoChecksums;
 
   private parsedSyncRulesCache: { parsed: SqlSyncRules; options: storage.ParseSyncRulesOptions } | undefined;
   private writeCheckpointAPI: MongoWriteCheckpointAPI;
@@ -58,13 +56,15 @@ export class MongoSyncBucketStorage
     public readonly group_id: number,
     private readonly sync_rules: storage.PersistedSyncRulesContent,
     public readonly slot_name: string,
-    writeCheckpointMode: storage.WriteCheckpointMode = storage.WriteCheckpointMode.MANAGED
+    writeCheckpointMode?: storage.WriteCheckpointMode,
+    options?: MongoSyncBucketStorageOptions
   ) {
     super();
     this.db = factory.db;
+    this.checksums = new MongoChecksums(this.db, this.group_id, options?.checksumOptions);
     this.writeCheckpointAPI = new MongoWriteCheckpointAPI({
       db: this.db,
-      mode: writeCheckpointMode,
+      mode: writeCheckpointMode ?? storage.WriteCheckpointMode.MANAGED,
       sync_rules_id: group_id
     });
   }
@@ -491,145 +491,11 @@ export class MongoSyncBucketStorage
   }
 
   async getChecksums(checkpoint: utils.InternalOpId, buckets: string[]): Promise<utils.ChecksumMap> {
-    return this.checksumCache.getChecksumMap(checkpoint, buckets);
+    return this.checksums.getChecksums(checkpoint, buckets);
   }
 
   clearChecksumCache() {
-    this.checksumCache.clear();
-  }
-
-  private async getChecksumsInternal(batch: storage.FetchPartialBucketChecksum[]): Promise<storage.PartialChecksumMap> {
-    if (batch.length == 0) {
-      return new Map();
-    }
-
-    const preFilters: any[] = [];
-    for (let request of batch) {
-      if (request.start == null) {
-        preFilters.push({
-          _id: {
-            g: this.group_id,
-            b: request.bucket
-          },
-          'compacted_state.op_id': { $exists: true, $lte: request.end }
-        });
-      }
-    }
-
-    const preStates = new Map<string, { opId: InternalOpId; checksum: BucketChecksum }>();
-
-    if (preFilters.length > 0) {
-      // For un-cached bucket checksums, attempt to use the compacted state first.
-      const states = await this.db.bucket_state
-        .find({
-          $or: preFilters
-        })
-        .toArray();
-      for (let state of states) {
-        const compactedState = state.compacted_state!;
-        preStates.set(state._id.b, {
-          opId: compactedState.op_id,
-          checksum: {
-            bucket: state._id.b,
-            checksum: Number(compactedState.checksum),
-            count: compactedState.count
-          }
-        });
-      }
-    }
-
-    const filters: any[] = [];
-    for (let request of batch) {
-      let start = request.start;
-      if (start == null) {
-        const preState = preStates.get(request.bucket);
-        if (preState != null) {
-          start = preState.opId;
-        }
-      }
-
-      filters.push({
-        _id: {
-          $gt: {
-            g: this.group_id,
-            b: request.bucket,
-            o: start ?? new bson.MinKey()
-          },
-          $lte: {
-            g: this.group_id,
-            b: request.bucket,
-            o: request.end
-          }
-        }
-      });
-    }
-
-    const aggregate = await this.db.bucket_data
-      .aggregate(
-        [
-          {
-            $match: {
-              $or: filters
-            }
-          },
-          {
-            $group: {
-              _id: '$_id.b',
-              // 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.
-              checksum_total: { $sum: { $toLong: '$checksum' } },
-              count: { $sum: 1 },
-              has_clear_op: {
-                $max: {
-                  $cond: [{ $eq: ['$op', 'CLEAR'] }, 1, 0]
-                }
-              }
-            }
-          }
-        ],
-        { session: undefined, readConcern: 'snapshot', maxTimeMS: lib_mongo.db.MONGO_CHECKSUM_TIMEOUT_MS }
-      )
-      .toArray()
-      .catch((e) => {
-        throw lib_mongo.mapQueryError(e, 'while reading checksums');
-      });
-
-    const partialChecksums = new Map<string, storage.PartialOrFullChecksum>(
-      aggregate.map((doc) => {
-        const partialChecksum = Number(BigInt(doc.checksum_total) & 0xffffffffn) & 0xffffffff;
-        const bucket = doc._id;
-        return [
-          bucket,
-          doc.has_clear_op == 1
-            ? ({
-                // full checksum - replaces any previous one
-                bucket,
-                checksum: partialChecksum,
-                count: doc.count
-              } satisfies BucketChecksum)
-            : ({
-                // partial checksum - is added to a previous one
-                bucket,
-                partialCount: doc.count,
-                partialChecksum
-              } satisfies PartialChecksum)
-        ];
-      })
-    );
-
-    return new Map<string, storage.PartialOrFullChecksum>(
-      batch.map((request) => {
-        const bucket = request.bucket;
-        // Could be null if this is either (1) a partial request, or (2) no compacted checksum was available
-        const preState = preStates.get(bucket);
-        // Could be null if we got no data
-        const partialChecksum = partialChecksums.get(bucket);
-        const merged = addPartialChecksums(bucket, preState?.checksum ?? null, partialChecksum ?? null);
-
-        return [bucket, merged];
-      })
-    );
+    this.checksums.clearCache();
   }
 
   async terminate(options?: storage.TerminateOptions) {
@@ -779,22 +645,25 @@ export class MongoSyncBucketStorage
       const checkpoint = await this.getCheckpointInternal();
       maxOpId = checkpoint?.checkpoint ?? undefined;
     }
-    await new MongoCompactor(this.db, this.group_id, { ...options, maxOpId }).compact();
+    await new MongoCompactor(this, this.db, { ...options, maxOpId }).compact();
+
     if (maxOpId != null && options?.compactParameterData) {
       await new MongoParameterCompactor(this.db, this.group_id, maxOpId, options).compact();
     }
   }
 
-  async populatePersistentChecksumCache(options: Pick<CompactOptions, 'signal' | 'maxOpId'>): Promise<void> {
+  async populatePersistentChecksumCache(options: Required<Pick<CompactOptions, 'signal' | 'maxOpId'>>): Promise<void> {
+    logger.info(`Populating persistent checksum cache...`);
     const start = Date.now();
-    // We do a minimal compact, primarily to populate the checksum cache
-    await this.compact({
+    // We do a minimal compact here.
+    // We can optimize this in the future.
+    const compactor = new MongoCompactor(this, this.db, {
       ...options,
-      // Skip parameter data
-      compactParameterData: false,
       // Don't track updates for MOVE compacting
       memoryLimitMB: 0
     });
+
+    await compactor.populateChecksums();
     const duration = Date.now() - start;
     logger.info(`Populated persistent checksum cache in ${(duration / 1000).toFixed(1)}s`);
   }

package/src/storage/implementation/MongoTestStorageFactoryGenerator.ts

@@ -1,10 +1,12 @@
 import { TestStorageOptions } from '@powersync/service-core';
 import { MongoBucketStorage } from '../MongoBucketStorage.js';
 import { connectMongoForTests } from './util.js';
+import { MongoSyncBucketStorageOptions } from './MongoSyncBucketStorage.js';
 
 export type MongoTestStorageOptions = {
   url: string;
   isCI: boolean;
+  internalOptions?: MongoSyncBucketStorageOptions;
 };
 
 export const MongoTestStorageFactoryGenerator = (factoryOptions: MongoTestStorageOptions) => {
@@ -16,13 +18,14 @@ export const MongoTestStorageFactoryGenerator = (factoryOptions: MongoTestStorag
       await db.db.createCollection('bucket_parameters');
     }
 
-    // Full migrations are not currently run for tests, so we manually create this
-    await db.createCheckpointEventsCollection();
-
     if (!options?.doNotClear) {
       await db.clear();
     }
 
-    return new MongoBucketStorage(db, { slot_name_prefix: 'test_' });
+    // Full migrations are not currently run for tests, so we manually create the important ones
+    await db.createCheckpointEventsCollection();
+    await db.createBucketStateIndex();
+
+    return new MongoBucketStorage(db, { slot_name_prefix: 'test_' }, factoryOptions.internalOptions);
   };
 };

package/src/storage/implementation/db.ts

@@ -127,6 +127,20 @@ export class PowerSyncMongo {
       max: 50 // max number of documents
     });
   }
+
+  /**
+   * Only use in migrations and tests.
+   */
+  async createBucketStateIndex() {
+    // TODO: Implement a better mechanism to use migrations in tests
+    await this.bucket_state.createIndex(
+      {
+        '_id.g': 1,
+        last_op: 1
+      },
+      { name: 'bucket_updates', unique: true }
+    );
+  }
 }
 
 export function createPowerSyncMongo(config: MongoStorageConfig, options?: lib_mongo.MongoConnectionOptions) {

package/src/storage/implementation/models.ts

@@ -97,6 +97,10 @@ export interface BucketStateDocument {
     g: number;
     b: string;
   };
+  /**
+   * Important: There is an unique index on {'_id.g': 1, last_op: 1}.
+   * That means the last_op must match an actual op in the bucket, and not the commit checkpoint.
+   */
   last_op: bigint;
   /**
    * If set, this can be treated as "cache" of a checksum at a specific point.
@@ -106,7 +110,7 @@ export interface BucketStateDocument {
     op_id: InternalOpId;
     count: number;
     checksum: bigint;
-    bytes: number;
+    bytes: number | null;
   };
 
   estimate_since_compact?: {

package/src/storage/implementation/util.ts

@@ -3,7 +3,7 @@ import * as crypto from 'crypto';
 import * as uuid from 'uuid';
 
 import { mongo } from '@powersync/lib-service-mongodb';
-import { storage, utils } from '@powersync/service-core';
+import { BucketChecksum, PartialChecksum, PartialOrFullChecksum, storage, utils } from '@powersync/service-core';
 
 import { PowerSyncMongo } from './db.js';
 import { BucketDataDocument } from './models.js';

package/test/src/__snapshots__/storage.test.ts.snap

@@ -1,6 +1,22 @@
 // Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
 
-exports[`Mongo Sync Bucket Storage > empty storage metrics 1`] = `
+exports[`Mongo Sync Bucket Storage - Data > empty storage metrics 1`] = `
+{
+  "operations_size_bytes": 0,
+  "parameters_size_bytes": 0,
+  "replication_size_bytes": 0,
+}
+`;
+
+exports[`Mongo Sync Bucket Storage - split buckets > empty storage metrics 1`] = `
+{
+  "operations_size_bytes": 0,
+  "parameters_size_bytes": 0,
+  "replication_size_bytes": 0,
+}
+`;
+
+exports[`Mongo Sync Bucket Storage - split operations > empty storage metrics 1`] = `
 {
   "operations_size_bytes": 0,
   "parameters_size_bytes": 0,

package/test/src/storage.test.ts

@@ -1,7 +1,44 @@
 import { register } from '@powersync/service-core-tests';
 import { describe } from 'vitest';
 import { INITIALIZED_MONGO_STORAGE_FACTORY } from './util.js';
+import { env } from './env.js';
+import { MongoTestStorageFactoryGenerator } from '@module/storage/implementation/MongoTestStorageFactoryGenerator.js';
 
-describe('Mongo Sync Bucket Storage', () => register.registerDataStorageTests(INITIALIZED_MONGO_STORAGE_FACTORY));
+describe('Mongo Sync Bucket Storage - Parameters', () =>
+  register.registerDataStorageParameterTests(INITIALIZED_MONGO_STORAGE_FACTORY));
+
+describe('Mongo Sync Bucket Storage - Data', () =>
+  register.registerDataStorageDataTests(INITIALIZED_MONGO_STORAGE_FACTORY));
+
+describe('Mongo Sync Bucket Storage - Checkpoints', () =>
+  register.registerDataStorageCheckpointTests(INITIALIZED_MONGO_STORAGE_FACTORY));
 
 describe('Sync Bucket Validation', register.registerBucketValidationTests);
+
+describe('Mongo Sync Bucket Storage - split operations', () =>
+  register.registerDataStorageDataTests(
+    MongoTestStorageFactoryGenerator({
+      url: env.MONGO_TEST_URL,
+      isCI: env.CI,
+      internalOptions: {
+        checksumOptions: {
+          bucketBatchLimit: 100,
+          operationBatchLimit: 1
+        }
+      }
+    })
+  ));
+
+describe('Mongo Sync Bucket Storage - split buckets', () =>
+  register.registerDataStorageDataTests(
+    MongoTestStorageFactoryGenerator({
+      url: env.MONGO_TEST_URL,
+      isCI: env.CI,
+      internalOptions: {
+        checksumOptions: {
+          bucketBatchLimit: 1,
+          operationBatchLimit: 100
+        }
+      }
+    })
+  ));