@powersync/service-module-mongodb-storage 0.0.0-dev-20250310210938 → 0.0.0-dev-20250312090341

package/package.json

@@ -2,7 +2,7 @@
   "name": "@powersync/service-module-mongodb-storage",
   "repository": "https://github.com/powersync-ja/powersync-service",
   "types": "dist/index.d.ts",
-  "version": "0.0.0-dev-20250310210938",
+  "version": "0.0.0-dev-20250312090341",
   "main": "dist/index.js",
   "license": "FSL-1.1-Apache-2.0",
   "type": "module",
@@ -28,15 +28,15 @@
     "lru-cache": "^10.2.2",
     "uuid": "^9.0.1",
     "@powersync/lib-services-framework": "0.5.3",
-    "@powersync/service-core": "0.0.0-dev-20250310210938",
+    "@powersync/service-core": "0.0.0-dev-20250312090341",
     "@powersync/service-jsonbig": "0.17.10",
     "@powersync/service-sync-rules": "0.24.1",
-    "@powersync/service-types": "0.0.0-dev-20250310210938",
-    "@powersync/lib-service-mongodb": "0.0.0-dev-20250310210938"
+    "@powersync/service-types": "0.9.0",
+    "@powersync/lib-service-mongodb": "0.0.0-dev-20250312090341"
   },
   "devDependencies": {
     "@types/uuid": "^9.0.4",
-    "@powersync/service-core-tests": "0.0.0-dev-20250310210938"
+    "@powersync/service-core-tests": "0.0.0-dev-20250312090341"
   },
   "scripts": {
     "build": "tsc -b",

package/src/migrations/db/migrations/1741697235857-bucket-state-index.ts

@@ -0,0 +1,40 @@
+import { migrations } from '@powersync/service-core';
+import * as storage from '../../../storage/storage-index.js';
+import { MongoStorageConfig } from '../../../types/types.js';
+
+const INDEX_NAME = 'bucket_updates';
+
+export const up: migrations.PowerSyncMigrationFunction = async (context) => {
+  const {
+    service_context: { configuration }
+  } = context;
+  const db = storage.createPowerSyncMongo(configuration.storage as MongoStorageConfig);
+
+  try {
+    await db.bucket_state.createIndex(
+      {
+        '_id.g': 1,
+        last_op: 1
+      },
+      { name: INDEX_NAME, unique: true }
+    );
+  } finally {
+    await db.client.close();
+  }
+};
+
+export const down: migrations.PowerSyncMigrationFunction = async (context) => {
+  const {
+    service_context: { configuration }
+  } = context;
+
+  const db = storage.createPowerSyncMongo(configuration.storage as MongoStorageConfig);
+
+  try {
+    if (await db.bucket_state.indexExists(INDEX_NAME)) {
+      await db.bucket_state.dropIndex(INDEX_NAME);
+    }
+  } finally {
+    await db.client.close();
+  }
+};

package/src/storage/implementation/MongoCompactor.ts

@@ -314,10 +314,12 @@ export class MongoCompactor {
             let lastOpId: BucketDataKey | null = null;
             let targetOp: bigint | null = null;
             let gotAnOp = false;
+            let numberOfOpsToClear = 0;
             for await (let op of query.stream()) {
               if (op.op == 'MOVE' || op.op == 'REMOVE' || op.op == 'CLEAR') {
                 checksum = utils.addChecksums(checksum, op.checksum);
                 lastOpId = op._id;
+                numberOfOpsToClear += 1;
                 if (op.op != 'CLEAR') {
                   gotAnOp = true;
                 }
@@ -337,7 +339,7 @@ export class MongoCompactor {
               return;
             }
 
-            logger.info(`Flushing CLEAR at ${lastOpId?.o}`);
+            logger.info(`Flushing CLEAR for ${numberOfOpsToClear} ops at ${lastOpId?.o}`);
             await this.db.bucket_data.deleteMany(
               {
                 _id: {
@@ -362,6 +364,22 @@ export class MongoCompactor {
               },
               { session }
             );
+
+            // Note: This does not update anything if there is no existing state
+            await this.db.bucket_state.updateOne(
+              {
+                _id: {
+                  g: this.group_id,
+                  b: bucket
+                }
+              },
+              {
+                $inc: {
+                  op_count: 1 - numberOfOpsToClear
+                }
+              },
+              { session }
+            );
           },
           {
             writeConcern: { w: 'majority' },

package/src/storage/implementation/MongoSyncBucketStorage.ts

@@ -9,27 +9,30 @@ import {
 } from '@powersync/lib-services-framework';
 import {
   BroadcastIterable,
-  CHECKPOINT_INVALIDATE_ALL,
   CheckpointChanges,
   GetCheckpointChangesOptions,
   InternalOpId,
   internalToExternalOpId,
   ProtocolOpId,
+  getLookupBucketDefinitionName,
   ReplicationCheckpoint,
-  SourceTable,
   storage,
   utils,
-  WatchWriteCheckpointOptions
+  WatchWriteCheckpointOptions,
+  CHECKPOINT_INVALIDATE_ALL,
+  deserializeParameterLookup
 } from '@powersync/service-core';
 import { SqliteJsonRow, SqliteJsonValue, SqlSyncRules } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
 import { wrapWithAbort } from 'ix/asynciterable/operators/withabort.js';
+import { LRUCache } from 'lru-cache';
 import * as timers from 'timers/promises';
 import { MongoBucketStorage } from '../MongoBucketStorage.js';
 import { PowerSyncMongo } from './db.js';
 import {
   BucketDataDocument,
   BucketDataKey,
+  BucketStateDocument,
   SourceKey,
   SourceTableDocument,
   SyncRuleCheckpointState,
@@ -39,6 +42,7 @@ import { MongoBucketBatch } from './MongoBucketBatch.js';
 import { MongoCompactor } from './MongoCompactor.js';
 import { MongoWriteCheckpointAPI } from './MongoWriteCheckpointAPI.js';
 import { idPrefixFilter, mapOpEntry, readSingleBatch } from './util.js';
+import { JSONBig } from '@powersync/service-jsonbig';
 
 export class MongoSyncBucketStorage
   extends BaseObserver<storage.SyncRulesBucketStorageListener>
@@ -585,6 +589,13 @@ export class MongoSyncBucketStorage
       { maxTimeMS: lib_mongo.db.MONGO_CLEAR_OPERATION_TIMEOUT_MS }
     );
 
+    await this.db.bucket_state.deleteMany(
+      {
+        _id: idPrefixFilter<BucketStateDocument['_id']>({ g: this.group_id }, ['b'])
+      },
+      { maxTimeMS: lib_mongo.db.MONGO_CLEAR_OPERATION_TIMEOUT_MS }
+    );
+
     await this.db.source_tables.deleteMany(
       {
         group_id: this.group_id
@@ -795,12 +806,7 @@ export class MongoSyncBucketStorage
 
       const updates: CheckpointChanges =
         lastCheckpoint == null
-          ? {
-              invalidateDataBuckets: true,
-              invalidateParameterBuckets: true,
-              updatedDataBuckets: [],
-              updatedParameterBucketDefinitions: []
-            }
+          ? CHECKPOINT_INVALIDATE_ALL
           : await this.getCheckpointChanges({
               lastCheckpoint: lastCheckpoint,
               nextCheckpoint: checkpoint
@@ -869,7 +875,119 @@ export class MongoSyncBucketStorage
     return pipeline;
   }
 
+  private async getDataBucketChanges(
+    options: GetCheckpointChangesOptions
+  ): Promise<Pick<CheckpointChanges, 'updatedDataBuckets' | 'invalidateDataBuckets'>> {
+    const bucketStateUpdates = await this.db.bucket_state
+      .find(
+        {
+          // We have an index on (_id.g, last_op).
+          '_id.g': this.group_id,
+          last_op: { $gt: BigInt(options.lastCheckpoint) }
+        },
+        {
+          projection: {
+            '_id.b': 1
+          },
+          limit: 1001,
+          batchSize: 1001,
+          singleBatch: true
+        }
+      )
+      .toArray();
+
+    const buckets = bucketStateUpdates.map((doc) => doc._id.b);
+    const invalidateDataBuckets = buckets.length > 1000;
+
+    return {
+      invalidateDataBuckets: invalidateDataBuckets,
+      updatedDataBuckets: invalidateDataBuckets ? [] : buckets
+    };
+  }
+
+  private async getParameterBucketChanges(
+    options: GetCheckpointChangesOptions
+  ): Promise<Pick<CheckpointChanges, 'updatedParameterLookups' | 'invalidateParameterBuckets'>> {
+    // TODO: limit max query running time
+    const parameterUpdates = await this.db.bucket_parameters
+      .find(
+        {
+          _id: { $gt: BigInt(options.lastCheckpoint), $lt: BigInt(options.nextCheckpoint) },
+          'key.g': this.group_id
+        },
+        {
+          projection: {
+            lookup: 1
+          },
+          limit: 1001,
+          batchSize: 1001,
+          singleBatch: true
+        }
+      )
+      .toArray();
+    const invalidateParameterUpdates = parameterUpdates.length > 1000;
+
+    return {
+      invalidateParameterBuckets: invalidateParameterUpdates,
+      updatedParameterLookups: invalidateParameterUpdates
+        ? new Set<string>()
+        : new Set<string>(parameterUpdates.map((p) => JSONBig.stringify(deserializeParameterLookup(p.lookup))))
+    };
+  }
+
+  // TODO:
+  // We can optimize this by implementing it like ChecksumCache: We can use partial cache results to do
+  // more efficient lookups in some cases.
+  private checkpointChangesCache = new LRUCache<string, CheckpointChanges, { options: GetCheckpointChangesOptions }>({
+    max: 50,
+    maxSize: 10 * 1024 * 1024,
+    sizeCalculation: (value: CheckpointChanges) => {
+      const paramSize = [...value.updatedParameterLookups].reduce<number>((a, b) => a + b.length, 0);
+      const bucketSize = [...value.updatedDataBuckets].reduce<number>((a, b) => a + b.length, 0);
+      return 100 + paramSize + bucketSize;
+    },
+    fetchMethod: async (_key, _staleValue, options) => {
+      return this.getCheckpointChangesInternal(options.context.options);
+    }
+  });
+
+  private _hasDynamicBucketsCached: boolean | undefined = undefined;
+
+  private hasDynamicBucketQueries(): boolean {
+    if (this._hasDynamicBucketsCached != null) {
+      return this._hasDynamicBucketsCached;
+    }
+    const syncRules = this.getParsedSyncRules({
+      defaultSchema: 'default' // n/a
+    });
+    const hasDynamicBuckets = syncRules.hasDynamicBucketQueries();
+    this._hasDynamicBucketsCached = hasDynamicBuckets;
+    return hasDynamicBuckets;
+  }
+
   async getCheckpointChanges(options: GetCheckpointChangesOptions): Promise<CheckpointChanges> {
-    return CHECKPOINT_INVALIDATE_ALL;
+    if (!this.hasDynamicBucketQueries()) {
+      // Special case when we have no dynamic parameter queries.
+      // In this case, we can avoid doing any queries.
+      return {
+        invalidateDataBuckets: true,
+        updatedDataBuckets: [],
+        invalidateParameterBuckets: false,
+        updatedParameterLookups: new Set<string>()
+      };
+    }
+    const key = `${options.lastCheckpoint}_${options.nextCheckpoint}`;
+    const result = await this.checkpointChangesCache.fetch(key, { context: { options } });
+    return result!;
+  }
+
+  private async getCheckpointChangesInternal(options: GetCheckpointChangesOptions): Promise<CheckpointChanges> {
+    const dataUpdates = await this.getDataBucketChanges(options);
+    const parameterUpdates = await this.getParameterBucketChanges(options);
+
+    return {
+      ...dataUpdates,
+      ...parameterUpdates
+    };
   }
 }

package/src/storage/implementation/PersistedBatch.ts

@@ -11,6 +11,7 @@ import { PowerSyncMongo } from './db.js';
 import {
   BucketDataDocument,
   BucketParameterDocument,
+  BucketStateDocument,
   CurrentBucket,
   CurrentDataDocument,
   SourceKey
@@ -48,6 +49,7 @@ export class PersistedBatch {
   bucketData: mongo.AnyBulkWriteOperation<BucketDataDocument>[] = [];
   bucketParameters: mongo.AnyBulkWriteOperation<BucketParameterDocument>[] = [];
   currentData: mongo.AnyBulkWriteOperation<CurrentDataDocument>[] = [];
+  bucketStates: Map<string, BucketStateUpdate> = new Map();
 
   /**
    * For debug logging only.
@@ -66,6 +68,19 @@ export class PersistedBatch {
     this.currentSize = writtenSize;
   }
 
+  private incrementBucket(bucket: string, op_id: InternalOpId) {
+    let existingState = this.bucketStates.get(bucket);
+    if (existingState) {
+      existingState.lastOp = op_id;
+      existingState.incrementCount += 1;
+    } else {
+      this.bucketStates.set(bucket, {
+        lastOp: op_id,
+        incrementCount: 1
+      });
+    }
+  }
+
   saveBucketData(options: {
     op_seq: MongoIdSequence;
     sourceKey: storage.ReplicaId;
@@ -120,6 +135,7 @@ export class PersistedBatch {
           }
         }
       });
+      this.incrementBucket(k.bucket, op_id);
     }
 
     for (let bd of remaining_buckets.values()) {
@@ -147,6 +163,7 @@ export class PersistedBatch {
         }
       });
       this.currentSize += 200;
+      this.incrementBucket(bd.bucket, op_id);
     }
   }
 
@@ -277,6 +294,14 @@ export class PersistedBatch {
       });
     }
 
+    if (this.bucketStates.size > 0) {
+      await db.bucket_state.bulkWrite(this.getBucketStateUpdates(), {
+        session,
+        // Per-bucket operation - order doesn't matter
+        ordered: false
+      });
+    }
+
     const duration = performance.now() - startAt;
     logger.info(
       `powersync_${this.group_id} Flushed ${this.bucketData.length} + ${this.bucketParameters.length} + ${
@@ -287,7 +312,37 @@ export class PersistedBatch {
     this.bucketData = [];
     this.bucketParameters = [];
     this.currentData = [];
+    this.bucketStates.clear();
     this.currentSize = 0;
     this.debugLastOpId = null;
   }
+
+  private getBucketStateUpdates(): mongo.AnyBulkWriteOperation<BucketStateDocument>[] {
+    return Array.from(this.bucketStates.entries()).map(([bucket, state]) => {
+      return {
+        updateOne: {
+          filter: {
+            _id: {
+              g: this.group_id,
+              b: bucket
+            }
+          },
+          update: {
+            $set: {
+              last_op: state.lastOp
+            },
+            $inc: {
+              op_count: state.incrementCount
+            }
+          },
+          upsert: true
+        }
+      } satisfies mongo.AnyBulkWriteOperation<BucketStateDocument>;
+    });
+  }
+}
+
+interface BucketStateUpdate {
+  lastOp: InternalOpId;
+  incrementCount: number;
 }

package/src/storage/implementation/db.ts

@@ -6,6 +6,7 @@ import { MongoStorageConfig } from '../../types/types.js';
 import {
   BucketDataDocument,
   BucketParameterDocument,
+  BucketStateDocument,
   CurrentDataDocument,
   CustomWriteCheckpointDocument,
   IdSequenceDocument,
@@ -33,6 +34,7 @@ export class PowerSyncMongo {
   readonly write_checkpoints: mongo.Collection<WriteCheckpointDocument>;
   readonly instance: mongo.Collection<InstanceDocument>;
   readonly locks: mongo.Collection<lib_mongo.locks.Lock>;
+  readonly bucket_state: mongo.Collection<BucketStateDocument>;
 
   readonly client: mongo.MongoClient;
   readonly db: mongo.Db;
@@ -55,6 +57,7 @@ export class PowerSyncMongo {
     this.write_checkpoints = db.collection('write_checkpoints');
     this.instance = db.collection('instance');
     this.locks = this.db.collection('locks');
+    this.bucket_state = this.db.collection('bucket_state');
   }
 
   /**
@@ -70,6 +73,7 @@ export class PowerSyncMongo {
     await this.write_checkpoints.deleteMany({});
     await this.instance.deleteOne({});
     await this.locks.deleteMany({});
+    await this.bucket_state.deleteMany({});
   }
 
   /**

package/src/storage/implementation/models.ts

@@ -75,6 +75,15 @@ export interface SourceTableDocument {
   snapshot_done: boolean | undefined;
 }
 
+export interface BucketStateDocument {
+  _id: {
+    g: number;
+    b: string;
+  };
+  last_op: bigint;
+  op_count: number;
+}
+
 export interface IdSequenceDocument {
   _id: string;
   op_id: bigint;

package/test/src/setup.ts

@@ -1,12 +1,13 @@
 import { container } from '@powersync/lib-services-framework';
+import { test_utils } from '@powersync/service-core-tests';
 import { beforeAll, beforeEach } from 'vitest';
-import { METRICS_HELPER } from '@powersync/service-core-tests';
 
 beforeAll(async () => {
   // Executes for every test file
   container.registerDefaults();
+  await test_utils.initMetrics();
 });
 
 beforeEach(async () => {
-  METRICS_HELPER.resetMetrics();
+  await test_utils.resetMetrics();
 });