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

package/src/storage/implementation/MongoChecksums.ts

@@ -0,0 +1,320 @@
+import {
+  addPartialChecksums,
+  bson,
+  BucketChecksum,
+  ChecksumCache,
+  ChecksumMap,
+  FetchPartialBucketChecksum,
+  InternalOpId,
+  isPartialChecksum,
+  PartialChecksum,
+  PartialChecksumMap,
+  PartialOrFullChecksum
+} from '@powersync/service-core';
+import * as lib_mongo from '@powersync/lib-service-mongodb';
+import { logger } from '@powersync/lib-services-framework';
+import { PowerSyncMongo } from './db.js';
+
+/**
+ * Checksum query implementation.
+ */
+export class MongoChecksums {
+  private cache = new ChecksumCache({
+    fetchChecksums: (batch) => {
+      return this.getChecksumsInternal(batch);
+    }
+  });
+
+  constructor(
+    private db: PowerSyncMongo,
+    private group_id: number
+  ) {}
+
+  /**
+   * Calculate checksums, utilizing the cache.
+   */
+  async getChecksums(checkpoint: InternalOpId, buckets: string[]): Promise<ChecksumMap> {
+    return this.cache.getChecksumMap(checkpoint, buckets);
+  }
+
+  clearCache() {
+    this.cache.clear();
+  }
+
+  /**
+   * Calculate (partial) checksums from bucket_state and the data collection.
+   *
+   * Results are not cached.
+   */
+  private async getChecksumsInternal(batch: FetchPartialBucketChecksum[]): Promise<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 mappedRequests = batch.map((request) => {
+      let start = request.start;
+      if (start == null) {
+        const preState = preStates.get(request.bucket);
+        if (preState != null) {
+          start = preState.opId;
+        }
+      }
+      return {
+        ...request,
+        start
+      };
+    });
+
+    const queriedChecksums = await this.queryPartialChecksums(mappedRequests);
+
+    return new Map<string, 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 = queriedChecksums.get(bucket);
+        const merged = addPartialChecksums(bucket, preState?.checksum ?? null, partialChecksum ?? null);
+
+        return [bucket, merged];
+      })
+    );
+  }
+
+  /**
+   * Calculate (partial) checksums from the data collection directly.
+   */
+  async queryPartialChecksums(batch: FetchPartialBucketChecksum[]): Promise<PartialChecksumMap> {
+    try {
+      return await this.queryPartialChecksumsInternal(batch);
+    } catch (e) {
+      if (e.codeName == 'MaxTimeMSExpired') {
+        logger.warn(`Checksum query timed out; falling back to slower version`, e);
+        // Timeout - try the slower but more robust version
+        return await this.queryPartialChecksumsFallback(batch);
+      }
+      throw lib_mongo.mapQueryError(e, 'while reading checksums');
+    }
+  }
+
+  private async queryPartialChecksumsInternal(batch: FetchPartialBucketChecksum[]): Promise<PartialChecksumMap> {
+    const filters: any[] = [];
+    for (let request of batch) {
+      filters.push({
+        _id: {
+          $gt: {
+            g: this.group_id,
+            b: request.bucket,
+            o: request.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
+            }
+          },
+          CHECKSUM_QUERY_GROUP_STAGE
+        ],
+        { session: undefined, readConcern: 'snapshot', maxTimeMS: lib_mongo.MONGO_CHECKSUM_TIMEOUT_MS }
+      )
+      // Don't map the error here - we want to keep timeout errors as-is
+      .toArray();
+
+    const partialChecksums = new Map<string, PartialOrFullChecksum>(
+      aggregate.map((doc) => {
+        const bucket = doc._id;
+        return [bucket, checksumFromAggregate(doc)];
+      })
+    );
+
+    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];
+      })
+    );
+  }
+
+  /**
+   * Checksums for large buckets can run over the query timeout.
+   * To avoid this, we query in batches.
+   * This version can handle larger amounts of data, but is slower, especially for many buckets.
+   */
+  async queryPartialChecksumsFallback(batch: FetchPartialBucketChecksum[]): Promise<PartialChecksumMap> {
+    const partialChecksums = new Map<string, PartialOrFullChecksum>();
+    for (let request of batch) {
+      const checksum = await this.slowChecksum(request);
+      partialChecksums.set(request.bucket, checksum);
+    }
+
+    return partialChecksums;
+  }
+
+  private async slowChecksum(request: FetchPartialBucketChecksum): Promise<PartialOrFullChecksum> {
+    const batchLimit = 50_000;
+
+    let lowerBound = 0n;
+    const bucket = request.bucket;
+
+    let runningChecksum: PartialOrFullChecksum = {
+      bucket,
+      partialCount: 0,
+      partialChecksum: 0
+    };
+    if (request.start == null) {
+      runningChecksum = {
+        bucket,
+        count: 0,
+        checksum: 0
+      };
+    }
+
+    while (true) {
+      const filter = {
+        _id: {
+          $gt: {
+            g: this.group_id,
+            b: bucket,
+            o: lowerBound
+          },
+          $lte: {
+            g: this.group_id,
+            b: bucket,
+            o: request.end
+          }
+        }
+      };
+      const docs = await this.db.bucket_data
+        .aggregate(
+          [
+            {
+              $match: filter
+            },
+            // sort and limit _before_ grouping
+            { $sort: { _id: 1 } },
+            { $limit: batchLimit },
+            CHECKSUM_QUERY_GROUP_STAGE
+          ],
+          { session: undefined, readConcern: 'snapshot', maxTimeMS: lib_mongo.MONGO_CHECKSUM_TIMEOUT_MS }
+        )
+        .toArray();
+      const doc = docs[0];
+      if (doc == null) {
+        return runningChecksum;
+      }
+      const partial = checksumFromAggregate(doc);
+      runningChecksum = addPartialChecksums(bucket, runningChecksum, partial);
+      const isFinal = doc.count != batchLimit;
+      if (isFinal) {
+        break;
+      } else {
+        lowerBound = doc.last_op;
+      }
+    }
+    return runningChecksum;
+  }
+}
+
+const CHECKSUM_QUERY_GROUP_STAGE = {
+  $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]
+      }
+    },
+    last_op: { $max: '$_id.o' }
+  }
+};
+
+/**
+ * Convert output of CHECKSUM_QUERY_GROUP_STAGE into a checksum.
+ */
+function checksumFromAggregate(doc: bson.Document): PartialOrFullChecksum {
+  const partialChecksum = Number(BigInt(doc.checksum_total) & 0xffffffffn) & 0xffffffff;
+  const bucket = doc._id;
+
+  if (doc.has_clear_op == 1) {
+    return {
+      // full checksum - replaces any previous one
+      bucket,
+      checksum: partialChecksum,
+      count: doc.count
+    } satisfies BucketChecksum;
+  } else {
+    return {
+      // partial checksum - is added to a previous one
+      bucket,
+      partialCount: doc.count,
+      partialChecksum
+    } satisfies PartialChecksum;
+  }
+}

package/src/storage/implementation/MongoCompactor.ts

@@ -1,9 +1,10 @@
-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';
 
@@ -68,12 +69,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 +139,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 +201,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 +285,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();
   }
@@ -475,4 +479,91 @@ export class MongoCompactor {
       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();
+  }
 }