@powersync/service-module-mongodb-storage 0.0.0-dev-20251015143910 → 0.0.0-dev-20251110113516

package/src/storage/implementation/MongoCompactor.ts

@@ -1,6 +1,13 @@
 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 {
+  addChecksums,
+  InternalOpId,
+  isPartialChecksum,
+  PopulateChecksumCacheResults,
+  storage,
+  utils
+} from '@powersync/service-core';
 
 import { PowerSyncMongo } from './db.js';
 import { BucketDataDocument, BucketDataKey, BucketStateDocument } from './models.js';
@@ -10,6 +17,7 @@ import { cacheKey } from './OperationBatch.js';
 interface CurrentBucketState {
   /** Bucket name */
   bucket: string;
+
   /**
    * Rows seen in the bucket, with the last op_id of each.
    */
@@ -96,67 +104,56 @@ export class MongoCompactor {
         // We can make this more efficient later on by iterating
         // through the buckets in a single query.
         // That makes batching more tricky, so we leave for later.
-        await this.compactInternal(bucket);
+        await this.compactSingleBucket(bucket);
       }
     } else {
-      await this.compactInternal(undefined);
+      await this.compactDirtyBuckets();
     }
   }
 
-  async compactInternal(bucket: string | undefined) {
-    const idLimitBytes = this.idLimitBytes;
+  private async compactDirtyBuckets() {
+    while (!this.signal?.aborted) {
+      // Process all buckets with 1 or more changes since last time
+      const buckets = await this.dirtyBucketBatch({ minBucketChanges: 1 });
+      if (buckets.length == 0) {
+        // All done
+        break;
+      }
+      for (let bucket of buckets) {
+        await this.compactSingleBucket(bucket);
+      }
+    }
+  }
 
-    let currentState: CurrentBucketState | null = null;
+  private async compactSingleBucket(bucket: string) {
+    const idLimitBytes = this.idLimitBytes;
 
-    let bucketLower: string | mongo.MinKey;
-    let bucketUpper: string | mongo.MaxKey;
+    let currentState: CurrentBucketState = {
+      bucket,
+      seen: new Map(),
+      trackingSize: 0,
+      lastNotPut: null,
+      opsSincePut: 0,
 
-    if (bucket == null) {
-      bucketLower = new mongo.MinKey();
-      bucketUpper = new mongo.MaxKey();
-    } else if (bucket.includes('[')) {
-      // Exact bucket name
-      bucketLower = bucket;
-      bucketUpper = bucket;
-    } else {
-      // Bucket definition name
-      bucketLower = `${bucket}[`;
-      bucketUpper = `${bucket}[\uFFFF`;
-    }
+      checksum: 0,
+      opCount: 0,
+      opBytes: 0
+    };
 
     // Constant lower bound
     const lowerBound: BucketDataKey = {
       g: this.group_id,
-      b: bucketLower as string,
+      b: bucket,
       o: new mongo.MinKey() as any
     };
 
     // Upper bound is adjusted for each batch
     let upperBound: BucketDataKey = {
       g: this.group_id,
-      b: bucketUpper as string,
+      b: bucket,
       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 }>(
@@ -184,7 +181,11 @@ export class MongoCompactor {
             }
           }
         ],
-        { batchSize: this.moveBatchQueryLimit }
+        {
+          // batchSize is 1 more than limit to auto-close the cursor.
+          // See https://github.com/mongodb/node-mongodb-native/pull/4580
+          batchSize: this.moveBatchQueryLimit + 1
+        }
       );
       // 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.
@@ -199,22 +200,6 @@ export class MongoCompactor {
       upperBound = batch[batch.length - 1]._id;
 
       for (let doc of batch) {
-        if (currentState == null || doc._id.b != currentState.bucket) {
-          await doneWithBucket();
-
-          currentState = {
-            bucket: doc._id.b,
-            seen: new Map(),
-            trackingSize: 0,
-            lastNotPut: null,
-            opsSincePut: 0,
-
-            checksum: 0,
-            opCount: 0,
-            opBytes: 0
-          };
-        }
-
         if (doc._id.o > this.maxOpId) {
           continue;
         }
@@ -285,12 +270,22 @@ export class MongoCompactor {
         }
       }
 
-      if (currentState != null) {
-        logger.info(`Processed batch of length ${batch.length} current bucket: ${currentState.bucket}`);
-      }
+      logger.info(`Processed batch of length ${batch.length} current bucket: ${bucket}`);
+    }
+
+    // Free memory before clearing bucket
+    currentState.seen.clear();
+    if (currentState.lastNotPut != null && currentState.opsSincePut >= 1) {
+      logger.info(
+        `Inserting CLEAR at ${this.group_id}:${bucket}:${currentState.lastNotPut} to remove ${currentState.opsSincePut} operations`
+      );
+      // Need flush() before clear()
+      await this.flush();
+      await this.clearBucket(currentState);
     }
 
-    await doneWithBucket();
+    // Do this _after_ clearBucket so that we have accurate counts.
+    this.updateBucketChecksums(currentState);
 
     // Need another flush after updateBucketChecksums()
     await this.flush();
@@ -478,50 +473,55 @@ export class MongoCompactor {
   /**
    * 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
-    };
+  async populateChecksums(options: { minBucketChanges: number }): Promise<PopulateChecksumCacheResults> {
+    let count = 0;
     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 buckets = await this.dirtyBucketBatch(options);
+      if (buckets.length == 0) {
+        // All done
+        break;
+      }
+      const start = Date.now();
+      logger.info(`Calculating checksums for batch of ${buckets.length} buckets, starting at ${buckets[0]}`);
+
+      await this.updateChecksumsBatch(buckets);
+      logger.info(`Updated checksums for batch of ${buckets.length} buckets in ${Date.now() - start}ms`);
+      count += buckets.length;
+    }
+    return { buckets: count };
+  }
 
-      const bucketsWithoutChecksums = await this.db.bucket_state
-        .find(filter, {
+  /**
+   * Returns a batch of dirty buckets - buckets with most changes first.
+   *
+   * This cannot be used to iterate on its own - the client is expected to process these buckets and
+   * set estimate_since_compact.count: 0 when done, before fetching the next batch.
+   */
+  private async dirtyBucketBatch(options: { minBucketChanges: number }): Promise<string[]> {
+    if (options.minBucketChanges <= 0) {
+      throw new ReplicationAssertionError('minBucketChanges must be >= 1');
+    }
+    // We make use of an index on {_id.g: 1, 'estimate_since_compact.count': -1}
+    const dirtyBuckets = await this.db.bucket_state
+      .find(
+        {
+          '_id.g': this.group_id,
+          'estimate_since_compact.count': { $gte: options.minBucketChanges }
+        },
+        {
           projection: {
             _id: 1
           },
           sort: {
-            _id: 1
+            'estimate_since_compact.count': -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));
+        }
+      )
+      .toArray();
 
-      lowerBound = bucketsWithoutChecksums[bucketsWithoutChecksums.length - 1]._id;
-    }
+    return dirtyBuckets.map((bucket) => bucket._id.b);
   }
 
   private async updateChecksumsBatch(buckets: string[]) {
@@ -555,6 +555,10 @@ export class MongoCompactor {
                 count: bucketChecksum.count,
                 checksum: BigInt(bucketChecksum.checksum),
                 bytes: null
+              },
+              estimate_since_compact: {
+                count: 0,
+                bytes: 0
               }
             }
           },

package/src/storage/implementation/MongoSyncBucketStorage.ts

@@ -16,6 +16,8 @@ import {
   InternalOpId,
   internalToExternalOpId,
   maxLsn,
+  PopulateChecksumCacheOptions,
+  PopulateChecksumCacheResults,
   ProtocolOpId,
   ReplicationCheckpoint,
   storage,
@@ -404,7 +406,9 @@ export class MongoSyncBucketStorage
         limit: batchLimit,
         // Increase batch size above the default 101, so that we can fill an entire batch in
         // one go.
-        batchSize: batchLimit,
+        // batchSize is 1 more than limit to auto-close the cursor.
+        // See https://github.com/mongodb/node-mongodb-native/pull/4580
+        batchSize: batchLimit + 1,
         // Raw mode is returns an array of Buffer instead of parsed documents.
         // We use it so that:
         // 1. We can calculate the document size accurately without serializing again.
@@ -664,7 +668,7 @@ export class MongoSyncBucketStorage
     }
   }
 
-  async populatePersistentChecksumCache(options: Required<Pick<CompactOptions, 'signal' | 'maxOpId'>>): Promise<void> {
+  async populatePersistentChecksumCache(options: PopulateChecksumCacheOptions): Promise<PopulateChecksumCacheResults> {
     logger.info(`Populating persistent checksum cache...`);
     const start = Date.now();
     // We do a minimal compact here.
@@ -675,9 +679,14 @@ export class MongoSyncBucketStorage
       memoryLimitMB: 0
     });
 
-    await compactor.populateChecksums();
+    const result = await compactor.populateChecksums({
+      // There are cases with millions of small buckets, in which case it can take very long to
+      // populate the checksums, with minimal benefit. We skip the small buckets here.
+      minBucketChanges: options.minBucketChanges ?? 10
+    });
     const duration = Date.now() - start;
     logger.info(`Populated persistent checksum cache in ${(duration / 1000).toFixed(1)}s`);
+    return result;
   }
 
   /**
@@ -906,7 +915,9 @@ export class MongoSyncBucketStorage
             '_id.b': 1
           },
           limit: limit + 1,
-          batchSize: limit + 1,
+          // batchSize is 1 more than limit to auto-close the cursor.
+          // See https://github.com/mongodb/node-mongodb-native/pull/4580
+          batchSize: limit + 2,
           singleBatch: true
         }
       )
@@ -936,7 +947,9 @@ export class MongoSyncBucketStorage
             lookup: 1
           },
           limit: limit + 1,
-          batchSize: limit + 1,
+          // batchSize is 1 more than limit to auto-close the cursor.
+          // See https://github.com/mongodb/node-mongodb-native/pull/4580
+          batchSize: limit + 2,
           singleBatch: true
         }
       )

package/src/storage/implementation/MongoWriteCheckpointAPI.ts

@@ -111,7 +111,9 @@ export class MongoWriteCheckpointAPI implements storage.WriteCheckpointAPI {
         },
         {
           limit: limit + 1,
-          batchSize: limit + 1,
+          // batchSize is 1 more than limit to auto-close the cursor.
+          // See https://github.com/mongodb/node-mongodb-native/pull/4580
+          batchSize: limit + 2,
           singleBatch: true
         }
       )
@@ -140,7 +142,9 @@ export class MongoWriteCheckpointAPI implements storage.WriteCheckpointAPI {
         },
         {
           limit: limit + 1,
-          batchSize: limit + 1,
+          // batchSize is 1 more than limit to auto-close the cursor.
+          // See https://github.com/mongodb/node-mongodb-native/pull/4580
+          batchSize: limit + 2,
           singleBatch: true
         }
       )

package/src/storage/implementation/db.ts

@@ -158,6 +158,19 @@ export class PowerSyncMongo {
       { name: 'bucket_updates', unique: true }
     );
   }
+  /**
+   * Only use in migrations and tests.
+   */
+  async createBucketStateIndex2() {
+    // TODO: Implement a better mechanism to use migrations in tests
+    await this.bucket_state.createIndex(
+      {
+        '_id.g': 1,
+        'estimate_since_compact.count': -1
+      },
+      { name: 'dirty_count' }
+    );
+  }
 }
 
 export function createPowerSyncMongo(config: MongoStorageConfig, options?: lib_mongo.MongoConnectionOptions) {

package/src/utils/util.ts

@@ -39,7 +39,7 @@ export function generateSlotName(prefix: string, sync_rules_id: number) {
  * However, that makes `has_more` detection very difficult, since the cursor is always closed
  * after the first batch. Instead, we do a workaround to only fetch a single batch below.
  *
- * For this to be effective, set batchSize = limit in the find command.
+ * For this to be effective, set batchSize = limit + 1 in the find command.
  */
 export async function readSingleBatch<T>(cursor: mongo.AbstractCursor<T>): Promise<{ data: T[]; hasMore: boolean }> {
   try {
@@ -114,3 +114,50 @@ export function setSessionSnapshotTime(session: mongo.ClientSession, time: bson.
     throw new ServiceAssertionError(`Session snapshotTime is already set`);
   }
 }
+
+export const createPaginatedConnectionQuery = async <T extends mongo.Document>(
+  query: mongo.Filter<T>,
+  collection: mongo.Collection<T>,
+  limit: number,
+  cursor?: string
+) => {
+  const createQuery = (cursor?: string) => {
+    if (!cursor) {
+      return query;
+    }
+    return {
+      $and: [
+        query,
+        {
+          /** We are using the connected at date as the cursor so that the functionality works the same on Postgres implementation
+           * The id field in postgres is an uuid, this will work similarly to the ObjectId in Mongodb
+           * */
+          connected_at: {
+            $lt: new Date(cursor)
+          }
+        }
+      ]
+    } as mongo.Filter<T>;
+  };
+
+  /** cursor.count() deprecated */
+  const total = await collection.countDocuments(query);
+
+  const findCursor = collection.find(createQuery(cursor), {
+    sort: {
+      /** We are sorting by connected at date descending to match cursor Postgres implementation */
+      connected_at: -1
+    }
+  });
+
+  const items = await findCursor.limit(limit).toArray();
+  const count = items.length;
+  return {
+    items,
+    total,
+    count,
+    /** Setting the cursor to the connected at date of the last item in the list */
+    cursor: count === limit ? items[items.length - 1].connected_at.toISOString() : undefined,
+    more: count < total
+  };
+};

package/test/src/__snapshots__/connection-report-storage.test.ts.snap

@@ -3,13 +3,13 @@
 exports[`Connection reporting storage > Should create a connection report if its after a day 1`] = `
 [
   {
-    "client_id": "client_week",
+    "client_id": "client_one",
     "sdk": "powersync-js/1.24.5",
     "user_agent": "powersync-js/1.21.0 powersync-web Firefox/141 linux",
     "user_id": "user_week",
   },
   {
-    "client_id": "client_week",
+    "client_id": "client_one",
     "sdk": "powersync-js/1.24.5",
     "user_agent": "powersync-js/1.21.0 powersync-web Firefox/141 linux",
     "user_id": "user_week",
@@ -213,3 +213,160 @@ exports[`Report storage tests > Should show currently connected users 1`] = `
   "users": 2,
 }
 `;
+
+exports[`Report storage tests > Should show paginated response of all connections of specified client_id 1`] = `
+{
+  "count": 1,
+  "cursor": undefined,
+  "items": [
+    {
+      "client_id": "client_two",
+      "sdk": "powersync-js/1.21.1",
+      "user_agent": "powersync-js/1.21.0 powersync-web Chromium/138 linux",
+      "user_id": "user_two",
+    },
+  ],
+  "more": false,
+  "total": 1,
+}
+`;
+
+exports[`Report storage tests > Should show paginated response of all connections with a limit 1`] = `
+{
+  "count": 4,
+  "cursor": "<removed-for-snapshot>",
+  "items": [
+    {
+      "client_id": "client_one",
+      "sdk": "powersync-dart/1.6.4",
+      "user_agent": "powersync-dart/1.6.4 Dart (flutter-web) Chrome/128 android",
+      "user_id": "user_one",
+    },
+    {
+      "client_id": "client_four",
+      "sdk": "powersync-js/1.21.4",
+      "user_agent": "powersync-js/1.21.0 powersync-web Firefox/141 linux",
+      "user_id": "user_four",
+    },
+    {
+      "client_id": "",
+      "sdk": "unknown",
+      "user_agent": "Dart (flutter-web) Chrome/128 android",
+      "user_id": "user_one",
+    },
+    {
+      "client_id": "client_two",
+      "sdk": "powersync-js/1.21.1",
+      "user_agent": "powersync-js/1.21.0 powersync-web Chromium/138 linux",
+      "user_id": "user_two",
+    },
+  ],
+  "more": true,
+  "total": 8,
+}
+`;
+
+exports[`Report storage tests > Should show paginated response of all connections with a limit 2`] = `
+{
+  "count": 4,
+  "cursor": undefined,
+  "items": [
+    {
+      "client_id": "client_one",
+      "sdk": "powersync-dart/1.6.4",
+      "user_agent": "powersync-dart/1.6.4 Dart (flutter-web) Chrome/128 android",
+      "user_id": "user_one",
+    },
+    {
+      "client_id": "client_four",
+      "sdk": "powersync-js/1.21.4",
+      "user_agent": "powersync-js/1.21.0 powersync-web Firefox/141 linux",
+      "user_id": "user_four",
+    },
+    {
+      "client_id": "",
+      "sdk": "unknown",
+      "user_agent": "Dart (flutter-web) Chrome/128 android",
+      "user_id": "user_one",
+    },
+    {
+      "client_id": "client_two",
+      "sdk": "powersync-js/1.21.1",
+      "user_agent": "powersync-js/1.21.0 powersync-web Chromium/138 linux",
+      "user_id": "user_two",
+    },
+  ],
+  "more": true,
+  "total": 8,
+}
+`;
+
+exports[`Report storage tests > Should show paginated response of connections of specified user_id 1`] = `
+{
+  "count": 2,
+  "cursor": undefined,
+  "items": [
+    {
+      "client_id": "client_one",
+      "sdk": "powersync-dart/1.6.4",
+      "user_agent": "powersync-dart/1.6.4 Dart (flutter-web) Chrome/128 android",
+      "user_id": "user_one",
+    },
+    {
+      "client_id": "",
+      "sdk": "unknown",
+      "user_agent": "Dart (flutter-web) Chrome/128 android",
+      "user_id": "user_one",
+    },
+  ],
+  "more": false,
+  "total": 2,
+}
+`;
+
+exports[`Report storage tests > Should show paginated response of connections over a date range 1`] = `
+{
+  "count": 6,
+  "cursor": undefined,
+  "items": [
+    {
+      "client_id": "client_one",
+      "sdk": "powersync-dart/1.6.4",
+      "user_agent": "powersync-dart/1.6.4 Dart (flutter-web) Chrome/128 android",
+      "user_id": "user_one",
+    },
+    {
+      "client_id": "client_four",
+      "sdk": "powersync-js/1.21.4",
+      "user_agent": "powersync-js/1.21.0 powersync-web Firefox/141 linux",
+      "user_id": "user_four",
+    },
+    {
+      "client_id": "client_two",
+      "sdk": "powersync-js/1.21.1",
+      "user_agent": "powersync-js/1.21.0 powersync-web Chromium/138 linux",
+      "user_id": "user_two",
+    },
+    {
+      "client_id": "",
+      "sdk": "unknown",
+      "user_agent": "Dart (flutter-web) Chrome/128 android",
+      "user_id": "user_one",
+    },
+    {
+      "client_id": "client_three",
+      "sdk": "powersync-js/1.21.2",
+      "user_agent": "powersync-js/1.21.0 powersync-web Firefox/141 linux",
+      "user_id": "user_three",
+    },
+    {
+      "client_id": "client_one",
+      "sdk": "powersync-js/1.24.5",
+      "user_agent": "powersync-js/1.21.0 powersync-web Firefox/141 linux",
+      "user_id": "user_week",
+    },
+  ],
+  "more": false,
+  "total": 6,
+}
+`;

package/test/src/storage_compacting.test.ts

@@ -97,10 +97,25 @@ bucket_definitions:
       await populate(bucketStorage);
       const { checkpoint } = await bucketStorage.getCheckpoint();
 
-      await bucketStorage.populatePersistentChecksumCache({
+      // Default is to small small numbers - should be a no-op
+      const result0 = await bucketStorage.populatePersistentChecksumCache({
+        maxOpId: checkpoint
+      });
+      expect(result0.buckets).toEqual(0);
+
+      // This should cache the checksums for the two buckets
+      const result1 = await bucketStorage.populatePersistentChecksumCache({
+        maxOpId: checkpoint,
+        minBucketChanges: 1
+      });
+      expect(result1.buckets).toEqual(2);
+
+      // This should be a no-op, as the checksums are already cached
+      const result2 = await bucketStorage.populatePersistentChecksumCache({
         maxOpId: checkpoint,
-        signal: new AbortController().signal
+        minBucketChanges: 1
       });
+      expect(result2.buckets).toEqual(0);
 
       const checksumAfter = await bucketStorage.getChecksums(checkpoint, ['by_user2["u1"]', 'by_user2["u2"]']);
       expect(checksumAfter.get('by_user2["u1"]')).toEqual({