@powersync/service-module-mongodb-storage 0.11.0 → 0.12.1

package/src/storage/implementation/MongoCompactor.ts

@@ -1,10 +1,12 @@
-import { mongo } from '@powersync/lib-service-mongodb';
-import { logger, ReplicationAssertionError } from '@powersync/lib-services-framework';
-import { InternalOpId, storage, utils } from '@powersync/service-core';
+import { mongo, MONGO_OPERATION_TIMEOUT_MS } from '@powersync/lib-service-mongodb';
+import { logger, ReplicationAssertionError, ServiceAssertionError } from '@powersync/lib-services-framework';
+import { addChecksums, InternalOpId, isPartialChecksum, storage, utils } from '@powersync/service-core';
 
 import { PowerSyncMongo } from './db.js';
-import { BucketDataDocument, BucketDataKey } from './models.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 */
@@ -27,6 +29,21 @@ interface CurrentBucketState {
    * Number of REMOVE/MOVE operations seen since lastNotPut.
    */
   opsSincePut: number;
+
+  /**
+   * Incrementally-updated checksum, up to maxOpId
+   */
+  checksum: number;
+
+  /**
+   * op count for the checksum
+   */
+  opCount: number;
+
+  /**
+   * Byte size of ops covered by the checksum.
+   */
+  opBytes: number;
 }
 
 /**
@@ -43,25 +60,30 @@ 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 maxOpId: bigint | undefined;
+  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;
     this.clearBatchLimit = options?.clearBatchLimit ?? DEFAULT_CLEAR_BATCH_LIMIT;
-    this.maxOpId = options?.maxOpId;
+    this.maxOpId = options?.maxOpId ?? 0n;
     this.buckets = options?.compactBuckets;
+    this.signal = options?.signal;
   }
 
   /**
@@ -117,31 +139,57 @@ export class MongoCompactor {
       o: new mongo.MaxKey() as any
     };
 
-    while (true) {
+    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 batch = await this.db.bucket_data
-        .find(
+      const cursor = this.db.bucket_data.aggregate<BucketDataDocument & { size: number | bigint }>(
+        [
           {
-            _id: {
-              $gte: lowerBound,
-              $lt: upperBound
+            $match: {
+              _id: {
+                $gte: lowerBound,
+                $lt: upperBound
+              }
             }
           },
+          { $sort: { _id: -1 } },
+          { $limit: this.moveBatchQueryLimit },
           {
-            projection: {
+            $project: {
               _id: 1,
               op: 1,
               table: 1,
               row_id: 1,
               source_table: 1,
-              source_key: 1
-            },
-            limit: this.moveBatchQueryLimit,
-            sort: { _id: -1 },
-            singleBatch: true
+              source_key: 1,
+              checksum: 1,
+              size: { $bsonSize: '$$ROOT' }
+            }
           }
-        )
-        .toArray();
+        ],
+        { 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
@@ -153,34 +201,31 @@ export class MongoCompactor {
 
       for (let doc of batch) {
         if (currentState == null || doc._id.b != currentState.bucket) {
-          if (currentState != null && currentState.lastNotPut != null && currentState.opsSincePut >= 1) {
-            // Important to flush before clearBucket()
-            await this.flush();
-            logger.info(
-              `Inserting CLEAR at ${this.group_id}:${currentState.bucket}:${currentState.lastNotPut} to remove ${currentState.opsSincePut} operations`
-            );
+          await doneWithBucket();
 
-            const bucket = currentState.bucket;
-            const clearOp = currentState.lastNotPut;
-            // Free memory before clearing bucket
-            currentState = null;
-            await this.clearBucket(bucket, clearOp);
-          }
           currentState = {
             bucket: doc._id.b,
             seen: new Map(),
             trackingSize: 0,
             lastNotPut: null,
-            opsSincePut: 0
+            opsSincePut: 0,
+
+            checksum: 0,
+            opCount: 0,
+            opBytes: 0
           };
         }
 
-        if (this.maxOpId != null && doc._id.o > this.maxOpId) {
+        if (doc._id.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);
         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);
@@ -208,6 +253,8 @@ export class MongoCompactor {
                 }
               }
             });
+
+            currentState.opBytes += 200 - Number(doc.size); // TODO: better estimate for this
           } else {
             if (currentState.trackingSize >= idLimitBytes) {
               // Reached memory limit.
@@ -234,24 +281,65 @@ export class MongoCompactor {
           currentState.opsSincePut += 1;
         }
 
-        if (this.updates.length >= this.moveBatchLimit) {
+        if (this.updates.length + this.bucketStateUpdates.length >= this.moveBatchLimit) {
           await this.flush();
         }
       }
+
+      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();
-    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`
+  }
+
+  /**
+   * Call when done with a bucket.
+   */
+  private updateBucketChecksums(state: CurrentBucketState) {
+    if (state.opCount < 0) {
+      throw new ServiceAssertionError(
+        `Invalid opCount: ${state.opCount} checksum ${state.checksum} opsSincePut: ${state.opsSincePut} maxOpId: ${this.maxOpId}`
       );
-      const bucket = currentState.bucket;
-      const clearOp = currentState.lastNotPut;
-      // Free memory before clearing bucket
-      currentState = null;
-      await this.clearBucket(bucket, clearOp);
     }
+    this.bucketStateUpdates.push({
+      updateOne: {
+        filter: {
+          _id: {
+            g: this.group_id,
+            b: state.bucket
+          }
+        },
+        update: {
+          $set: {
+            compacted_state: {
+              op_id: this.maxOpId,
+              count: state.opCount,
+              checksum: BigInt(state.checksum),
+              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.
+              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
+      }
+    });
   }
 
   private async flush() {
@@ -266,15 +354,26 @@ export class MongoCompactor {
       });
       this.updates = [];
     }
+    if (this.bucketStateUpdates.length > 0) {
+      logger.info(`Updating ${this.bucketStateUpdates.length} bucket states`);
+      await this.db.bucket_state.bulkWrite(this.bucketStateUpdates, {
+        ordered: false
+      });
+      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.
    */
-  private async clearBucket(bucket: string, op: InternalOpId) {
+  private async clearBucket(currentState: CurrentBucketState) {
+    const bucket = currentState.bucket;
+    const clearOp = currentState.lastNotPut!;
+
     const opFilter = {
       _id: {
         $gte: {
@@ -285,7 +384,7 @@ export class MongoCompactor {
         $lte: {
           g: this.group_id,
           b: bucket,
-          o: op
+          o: clearOp
         }
       }
     };
@@ -293,7 +392,8 @@ export class MongoCompactor {
     const session = this.db.client.startSession();
     try {
       let done = false;
-      while (!done) {
+      while (!done && !this.signal?.aborted) {
+        let opCountDiff = 0;
         // Do the CLEAR operation in batches, with each batch a separate transaction.
         // The state after each batch is fully consistent.
         // We need a transaction per batch to make sure checksums stay consistent.
@@ -364,15 +464,106 @@ export class MongoCompactor {
               },
               { session }
             );
+
+            opCountDiff = -numberOfOpsToClear + 1;
           },
           {
             writeConcern: { w: 'majority' },
             readConcern: { level: 'snapshot' }
           }
         );
+        // Update _outside_ the transaction, since the transaction can be retried multiple times.
+        currentState.opCount += opCountDiff;
       }
     } finally {
       await session.endSession();
     }
   }
+
+  /**
+   * Subset of compact, only populating checksums where relevant.
+   */
+  async populateChecksums() {
+    let lastId: BucketStateDocument['_id'] | null = null;
+    while (!this.signal?.aborted) {
+      // By filtering buckets, we effectively make this "resumeable".
+      let filter: mongo.Filter<BucketStateDocument> = {
+        compacted_state: { $exists: false }
+      };
+      if (lastId) {
+        filter._id = { $gt: lastId };
+      }
+
+      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));
+
+      lastId = bucketsWithoutChecksums[bucketsWithoutChecksums.length - 1]._id;
+    }
+  }
+
+  private async updateChecksumsBatch(buckets: string[]) {
+    const checksums = await this.storage.checksums.queryPartialChecksums(
+      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
+              }
+            },
+            $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
+        }
+      });
+    }
+
+    await this.flush();
+  }
 }

package/src/storage/implementation/MongoSyncBucketStorage.ts

@@ -10,6 +10,7 @@ import {
   BroadcastIterable,
   CHECKPOINT_INVALIDATE_ALL,
   CheckpointChanges,
+  CompactOptions,
   deserializeParameterLookup,
   GetCheckpointChangesOptions,
   InternalOpId,
@@ -30,21 +31,18 @@ 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 { 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';
-import { MongoParameterCompactor } from './MongoParameterCompactor.js';
 
 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,6 +56,7 @@ export class MongoSyncBucketStorage
   ) {
     super();
     this.db = factory.db;
+    this.checksums = new MongoChecksums(this.db, this.group_id);
     this.writeCheckpointAPI = new MongoWriteCheckpointAPI({
       db: this.db,
       mode: writeCheckpointMode,
@@ -487,76 +486,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);
   }
 
-  private async getChecksumsInternal(batch: storage.FetchPartialBucketChecksum[]): Promise<storage.PartialChecksumMap> {
-    if (batch.length == 0) {
-      return new Map();
-    }
-
-    const filters: any[] = [];
-    for (let request of batch) {
-      filters.push({
-        _id: {
-          $gt: {
-            g: this.group_id,
-            b: request.bucket,
-            o: request.start ? BigInt(request.start) : new bson.MinKey()
-          },
-          $lte: {
-            g: this.group_id,
-            b: request.bucket,
-            o: BigInt(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_OPERATION_TIMEOUT_MS }
-      )
-      .toArray()
-      .catch((e) => {
-        throw lib_mongo.mapQueryError(e, 'while reading checksums');
-      });
-
-    return new Map<string, storage.PartialChecksum>(
-      aggregate.map((doc) => {
-        return [
-          doc._id,
-          {
-            bucket: doc._id,
-            partialCount: doc.count,
-            partialChecksum: Number(BigInt(doc.checksum_total) & 0xffffffffn) & 0xffffffff,
-            isFullChecksum: doc.has_clear_op == 1
-          } satisfies storage.PartialChecksum
-        ];
-      })
-    );
+  clearChecksumCache() {
+    this.checksums.clearCache();
   }
 
   async terminate(options?: storage.TerminateOptions) {
@@ -701,13 +635,34 @@ export class MongoSyncBucketStorage
   }
 
   async compact(options?: storage.CompactOptions) {
-    const checkpoint = await this.getCheckpointInternal();
-    await new MongoCompactor(this.db, this.group_id, options).compact();
-    if (checkpoint != null && options?.compactParameterData) {
-      await new MongoParameterCompactor(this.db, this.group_id, checkpoint.checkpoint, options).compact();
+    let maxOpId = options?.maxOpId;
+    if (maxOpId == null) {
+      const checkpoint = await this.getCheckpointInternal();
+      maxOpId = checkpoint?.checkpoint ?? undefined;
+    }
+    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: Required<Pick<CompactOptions, 'signal' | 'maxOpId'>>): Promise<void> {
+    logger.info(`Populating persistent checksum cache...`);
+    const start = Date.now();
+    // We do a minimal compact here.
+    // We can optimize this in the future.
+    const compactor = new MongoCompactor(this, this.db, {
+      ...options,
+      // 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`);
+  }
+
   /**
    * Instance-wide watch on the latest available checkpoint (op_id + lsn).
    */

package/src/storage/implementation/PersistedBatch.ts

@@ -71,15 +71,17 @@ export class PersistedBatch {
     this.logger = options?.logger ?? defaultLogger;
   }
 
-  private incrementBucket(bucket: string, op_id: InternalOpId) {
+  private incrementBucket(bucket: string, op_id: InternalOpId, bytes: number) {
     let existingState = this.bucketStates.get(bucket);
     if (existingState) {
       existingState.lastOp = op_id;
       existingState.incrementCount += 1;
+      existingState.incrementBytes += bytes;
     } else {
       this.bucketStates.set(bucket, {
         lastOp: op_id,
-        incrementCount: 1
+        incrementCount: 1,
+        incrementBytes: bytes
       });
     }
   }
@@ -115,7 +117,8 @@ export class PersistedBatch {
       }
 
       remaining_buckets.delete(key);
-      this.currentSize += recordData.length + 200;
+      const byteEstimate = recordData.length + 200;
+      this.currentSize += byteEstimate;
 
       const op_id = options.op_seq.next();
       this.debugLastOpId = op_id;
@@ -138,7 +141,7 @@ export class PersistedBatch {
           }
         }
       });
-      this.incrementBucket(k.bucket, op_id);
+      this.incrementBucket(k.bucket, op_id, byteEstimate);
     }
 
     for (let bd of remaining_buckets.values()) {
@@ -166,7 +169,7 @@ export class PersistedBatch {
         }
       });
       this.currentSize += 200;
-      this.incrementBucket(bd.bucket, op_id);
+      this.incrementBucket(bd.bucket, op_id, 200);
     }
   }
 
@@ -369,6 +372,10 @@ export class PersistedBatch {
           update: {
             $set: {
               last_op: state.lastOp
+            },
+            $inc: {
+              'estimate_since_compact.count': state.incrementCount,
+              'estimate_since_compact.bytes': state.incrementBytes
             }
           },
           upsert: true
@@ -381,4 +388,5 @@ export class PersistedBatch {
 interface BucketStateUpdate {
   lastOp: InternalOpId;
   incrementCount: number;
+  incrementBytes: number;
 }

package/src/storage/implementation/models.ts

@@ -98,6 +98,21 @@ export interface BucketStateDocument {
     b: string;
   };
   last_op: bigint;
+  /**
+   * If set, this can be treated as "cache" of a checksum at a specific point.
+   * Can be updated periodically, for example by the compact job.
+   */
+  compacted_state?: {
+    op_id: InternalOpId;
+    count: number;
+    checksum: bigint;
+    bytes: number | null;
+  };
+
+  estimate_since_compact?: {
+    count: number;
+    bytes: number;
+  };
 }
 
 export interface IdSequenceDocument {

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';
@@ -43,7 +43,7 @@ export function generateSlotName(prefix: string, sync_rules_id: number) {
  *
  * For this to be effective, set batchSize = limit in the find command.
  */
-export async function readSingleBatch<T>(cursor: mongo.FindCursor<T>): Promise<{ data: T[]; hasMore: boolean }> {
+export async function readSingleBatch<T>(cursor: mongo.AbstractCursor<T>): Promise<{ data: T[]; hasMore: boolean }> {
   try {
     let data: T[];
     let hasMore = true;