Track last_op for each bucket.

Author: rkistner

modules/module-mongodb-storage/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();
+  }
+};

modules/module-mongodb-storage/src/storage/implementation/MongoSyncBucketStorage.ts

@@ -32,6 +32,7 @@ import { PowerSyncMongo } from './db.js';
 import {
   BucketDataDocument,
   BucketDataKey,
+  BucketStateDocument,
   SourceKey,
   SourceTableDocument,
   SyncRuleCheckpointState,
@@ -588,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
@@ -870,40 +878,36 @@ export class MongoSyncBucketStorage
   private async getDataBucketChanges(
     options: GetCheckpointChangesOptions
   ): Promise<Pick<CheckpointChanges, 'updatedDataBuckets' | 'invalidateDataBuckets'>> {
-    // The query below can be slow, since we don't have an index on _id.o.
-    // We could try to query the oplog for these, but that is risky.
-    // Or we could store updated buckets in a separate collection, and query those.
-    // For now, we ignore this optimization
-
-    // const dataBucketDocuments = await this.db.bucket_data
-    //   .find(
-    //     {
-    //       '_id.g': this.group_id,
-    //       '_id.o': { $gt: BigInt(options.lastCheckpoint), $lte: BigInt(options.nextCheckpoint) }
-    //     },
-    //     {
-    //       projection: {
-    //         '_id.b': 1
-    //       },
-    //       limit: 1001,
-    //       batchSize: 1001,
-    //       singleBatch: true
-    //     }
-    //   )
-    //   .toArray();
-
-    // const buckets = dataBucketDocuments.map((doc) => doc._id.b);
-    // const invalidateDataBuckets = buckets.length > 1000;
+    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();
 
-    // return {
-    //   invalidateDataBuckets: invalidateDataBuckets,
-    //   updatedDataBuckets: invalidateDataBuckets ? [] : buckets
-    // };
+    const buckets = bucketStateUpdates.map((doc) => doc._id.b);
+    const invalidateDataBuckets = buckets.length > 1000;
 
     return {
-      invalidateDataBuckets: true,
-      updatedDataBuckets: []
+      invalidateDataBuckets: invalidateDataBuckets,
+      updatedDataBuckets: invalidateDataBuckets ? [] : buckets
     };
+
+    // return {
+    //   invalidateDataBuckets: true,
+    //   updatedDataBuckets: []
+    // };
   }
 
   private async getParameterBucketChanges(

modules/module-mongodb-storage/src/storage/implementation/PersistedBatch.ts

@@ -4,13 +4,14 @@ import { EvaluatedParameters, EvaluatedRow } from '@powersync/service-sync-rules
 import * as bson from 'bson';
 
 import { logger } from '@powersync/lib-services-framework';
-import { InternalOpId, storage, utils } from '@powersync/service-core';
+import { BucketState, InternalOpId, storage, utils } from '@powersync/service-core';
 import { currentBucketKey } from './MongoBucketBatch.js';
 import { MongoIdSequence } from './MongoIdSequence.js';
 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;
@@ -111,6 +126,7 @@ export class PersistedBatch {
           }
         }
       });
+      this.incrementBucket(k.bucket, op_id);
     }
 
     for (let bd of remaining_buckets.values()) {
@@ -138,6 +154,7 @@ export class PersistedBatch {
         }
       });
       this.currentSize += 200;
+      this.incrementBucket(bd.bucket, op_id);
     }
   }
 
@@ -268,6 +285,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} + ${
@@ -278,7 +303,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;
 }

modules/module-mongodb-storage/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({});
   }
 
   /**

modules/module-mongodb-storage/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;