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

package/src/storage/implementation/MongoSyncBucketStorage.ts

@@ -2,20 +2,23 @@ import * as lib_mongo from '@powersync/lib-service-mongodb';
 import { mongo } from '@powersync/lib-service-mongodb';
 import {
   BaseObserver,
-  ErrorCode,
   logger,
   ReplicationAbortedError,
-  ServiceAssertionError,
-  ServiceError
+  ServiceAssertionError
 } from '@powersync/lib-services-framework';
 import {
+  addPartialChecksums,
   BroadcastIterable,
+  BucketChecksum,
   CHECKPOINT_INVALIDATE_ALL,
   CheckpointChanges,
+  CompactOptions,
   deserializeParameterLookup,
   GetCheckpointChangesOptions,
   InternalOpId,
   internalToExternalOpId,
+  maxLsn,
+  PartialChecksum,
   ProtocolOpId,
   ReplicationCheckpoint,
   storage,
@@ -29,18 +32,12 @@ 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
-} from './models.js';
+import { BucketDataDocument, BucketDataKey, BucketStateDocument, SourceKey, SourceTableDocument } from './models.js';
 import { MongoBucketBatch } from './MongoBucketBatch.js';
 import { MongoCompactor } from './MongoCompactor.js';
+import { MongoParameterCompactor } from './MongoParameterCompactor.js';
 import { MongoWriteCheckpointAPI } from './MongoWriteCheckpointAPI.js';
-import { idPrefixFilter, mapOpEntry, readSingleBatch } from './util.js';
+import { idPrefixFilter, mapOpEntry, readSingleBatch, setSessionSnapshotTime } from './util.js';
 
 export class MongoSyncBucketStorage
   extends BaseObserver<storage.SyncRulesBucketStorageListener>
@@ -105,22 +102,44 @@ export class MongoSyncBucketStorage
   }
 
   async getCheckpoint(): Promise<storage.ReplicationCheckpoint> {
-    const doc = await this.db.sync_rules.findOne(
-      { _id: this.group_id },
-      {
-        projection: { last_checkpoint: 1, last_checkpoint_lsn: 1, snapshot_done: 1 }
+    return (await this.getCheckpointInternal()) ?? new EmptyReplicationCheckpoint();
+  }
+
+  async getCheckpointInternal(): Promise<storage.ReplicationCheckpoint | null> {
+    return await this.db.client.withSession({ snapshot: true }, async (session) => {
+      const doc = await this.db.sync_rules.findOne(
+        { _id: this.group_id },
+        {
+          session,
+          projection: { _id: 1, state: 1, last_checkpoint: 1, last_checkpoint_lsn: 1, snapshot_done: 1 }
+        }
+      );
+      if (!doc?.snapshot_done || !['ACTIVE', 'ERRORED'].includes(doc.state)) {
+        // Sync rules not active - return null
+        return null;
       }
-    );
-    if (!doc?.snapshot_done) {
-      return {
-        checkpoint: 0n,
-        lsn: null
-      };
-    }
-    return {
-      checkpoint: doc?.last_checkpoint ?? 0n,
-      lsn: doc?.last_checkpoint_lsn ?? null
-    };
+
+      // Specifically using operationTime instead of clusterTime
+      // There are 3 fields in the response:
+      // 1. operationTime, not exposed for snapshot sessions (used for causal consistency)
+      // 2. clusterTime (used for connection management)
+      // 3. atClusterTime, which is session.snapshotTime
+      // We use atClusterTime, to match the driver's internal snapshot handling.
+      // There are cases where clusterTime > operationTime and atClusterTime,
+      // which could cause snapshot queries using this as the snapshotTime to timeout.
+      // This was specifically observed on MongoDB 6.0 and 7.0.
+      const snapshotTime = (session as any).snapshotTime as bson.Timestamp | undefined;
+      if (snapshotTime == null) {
+        throw new ServiceAssertionError('Missing snapshotTime in getCheckpoint()');
+      }
+      return new MongoReplicationCheckpoint(
+        this,
+        // null/0n is a valid checkpoint in some cases, for example if the initial snapshot was empty
+        doc.last_checkpoint ?? 0n,
+        doc.last_checkpoint_lsn ?? null,
+        snapshotTime
+      );
+    });
   }
 
   async startBatch(
@@ -131,7 +150,7 @@ export class MongoSyncBucketStorage
       {
         _id: this.group_id
       },
-      { projection: { last_checkpoint_lsn: 1, no_checkpoint_before: 1, keepalive_op: 1 } }
+      { projection: { last_checkpoint_lsn: 1, no_checkpoint_before: 1, keepalive_op: 1, snapshot_lsn: 1 } }
     );
     const checkpoint_lsn = doc?.last_checkpoint_lsn ?? null;
 
@@ -142,6 +161,7 @@ export class MongoSyncBucketStorage
       groupId: this.group_id,
       slotName: this.slot_name,
       lastCheckpointLsn: checkpoint_lsn,
+      resumeFromLsn: maxLsn(checkpoint_lsn, doc?.snapshot_lsn),
       noCheckpointBeforeLsn: doc?.no_checkpoint_before ?? options.zeroLSN,
       keepaliveOp: doc?.keepalive_op ? BigInt(doc.keepalive_op) : null,
       storeCurrentData: options.storeCurrentData,
@@ -162,9 +182,9 @@ export class MongoSyncBucketStorage
   async resolveTable(options: storage.ResolveTableOptions): Promise<storage.ResolveTableResult> {
     const { group_id, connection_id, connection_tag, entity_descriptor } = options;
 
-    const { schema, name: table, objectId, replicationColumns } = entity_descriptor;
+    const { schema, name, objectId, replicaIdColumns } = entity_descriptor;
 
-    const columns = replicationColumns.map((column) => ({
+    const normalizedReplicaIdColumns = replicaIdColumns.map((column) => ({
       name: column.name,
       type: column.type,
       type_oid: column.typeId
@@ -176,8 +196,8 @@ export class MongoSyncBucketStorage
         group_id: group_id,
         connection_id: connection_id,
         schema_name: schema,
-        table_name: table,
-        replica_id_columns2: columns
+        table_name: name,
+        replica_id_columns2: normalizedReplicaIdColumns
       };
       if (objectId != null) {
         filter.relation_id = objectId;
@@ -190,24 +210,24 @@ export class MongoSyncBucketStorage
           connection_id: connection_id,
           relation_id: objectId,
           schema_name: schema,
-          table_name: table,
+          table_name: name,
           replica_id_columns: null,
-          replica_id_columns2: columns,
+          replica_id_columns2: normalizedReplicaIdColumns,
           snapshot_done: false,
           snapshot_status: undefined
         };
 
         await col.insertOne(doc, { session });
       }
-      const sourceTable = new storage.SourceTable(
-        doc._id,
-        connection_tag,
-        objectId,
-        schema,
-        table,
-        replicationColumns,
-        doc.snapshot_done ?? true
-      );
+      const sourceTable = new storage.SourceTable({
+        id: doc._id,
+        connectionTag: connection_tag,
+        objectId: objectId,
+        schema: schema,
+        name: name,
+        replicaIdColumns: replicaIdColumns,
+        snapshotComplete: doc.snapshot_done ?? true
+      });
       sourceTable.syncEvent = options.sync_rules.tableTriggersEvent(sourceTable);
       sourceTable.syncData = options.sync_rules.tableSyncsData(sourceTable);
       sourceTable.syncParameters = options.sync_rules.tableSyncsParameters(sourceTable);
@@ -222,7 +242,7 @@ export class MongoSyncBucketStorage
 
       let dropTables: storage.SourceTable[] = [];
       // Detect tables that are either renamed, or have different replica_id_columns
-      let truncateFilter = [{ schema_name: schema, table_name: table }] as any[];
+      let truncateFilter = [{ schema_name: schema, table_name: name }] as any[];
       if (objectId != null) {
         // Only detect renames if the source uses relation ids.
         truncateFilter.push({ relation_id: objectId });
@@ -240,15 +260,16 @@ export class MongoSyncBucketStorage
         .toArray();
       dropTables = truncate.map(
         (doc) =>
-          new storage.SourceTable(
-            doc._id,
-            connection_tag,
-            doc.relation_id,
-            doc.schema_name,
-            doc.table_name,
-            doc.replica_id_columns2?.map((c) => ({ name: c.name, typeOid: c.type_oid, type: c.type })) ?? [],
-            doc.snapshot_done ?? true
-          )
+          new storage.SourceTable({
+            id: doc._id,
+            connectionTag: connection_tag,
+            objectId: doc.relation_id,
+            schema: doc.schema_name,
+            name: doc.table_name,
+            replicaIdColumns:
+              doc.replica_id_columns2?.map((c) => ({ name: c.name, typeOid: c.type_oid, type: c.type })) ?? [],
+            snapshotComplete: doc.snapshot_done ?? true
+          })
       );
 
       result = {
@@ -259,38 +280,67 @@ export class MongoSyncBucketStorage
     return result!;
   }
 
-  async getParameterSets(checkpoint: utils.InternalOpId, lookups: ParameterLookup[]): Promise<SqliteJsonRow[]> {
-    const lookupFilter = lookups.map((lookup) => {
-      return storage.serializeLookup(lookup);
-    });
-    const rows = await this.db.bucket_parameters
-      .aggregate([
-        {
-          $match: {
-            'key.g': this.group_id,
-            lookup: { $in: lookupFilter },
-            _id: { $lte: checkpoint }
-          }
-        },
-        {
-          $sort: {
-            _id: -1
-          }
-        },
-        {
-          $group: {
-            _id: { key: '$key', lookup: '$lookup' },
-            bucket_parameters: {
-              $first: '$bucket_parameters'
+  async getParameterSets(checkpoint: MongoReplicationCheckpoint, lookups: ParameterLookup[]): Promise<SqliteJsonRow[]> {
+    return this.db.client.withSession({ snapshot: true }, async (session) => {
+      // Set the session's snapshot time to the checkpoint's snapshot time.
+      // An alternative would be to create the session when the checkpoint is created, but managing
+      // the session lifetime would become more complex.
+      // Starting and ending sessions are cheap (synchronous when no transactions are used),
+      // so this should be fine.
+      // This is a roundabout way of setting {readConcern: {atClusterTime: clusterTime}}, since
+      // that is not exposed directly by the driver.
+      // Future versions of the driver may change the snapshotTime behavior, so we need tests to
+      // validate that this works as expected. We test this in the compacting tests.
+      setSessionSnapshotTime(session, checkpoint.snapshotTime);
+      const lookupFilter = lookups.map((lookup) => {
+        return storage.serializeLookup(lookup);
+      });
+      // This query does not use indexes super efficiently, apart from the lookup filter.
+      // From some experimentation I could do individual lookups more efficient using an index
+      // on {'key.g': 1, lookup: 1, 'key.t': 1, 'key.k': 1, _id: -1},
+      // but could not do the same using $group.
+      // For now, just rely on compacting to remove extraneous data.
+      // For a description of the data format, see the `/docs/parameters-lookups.md` file.
+      const rows = await this.db.bucket_parameters
+        .aggregate(
+          [
+            {
+              $match: {
+                'key.g': this.group_id,
+                lookup: { $in: lookupFilter },
+                _id: { $lte: checkpoint.checkpoint }
+              }
+            },
+            {
+              $sort: {
+                _id: -1
+              }
+            },
+            {
+              $group: {
+                _id: { key: '$key', lookup: '$lookup' },
+                bucket_parameters: {
+                  $first: '$bucket_parameters'
+                }
+              }
             }
+          ],
+          {
+            session,
+            readConcern: 'snapshot',
+            // Limit the time for the operation to complete, to avoid getting connection timeouts
+            maxTimeMS: lib_mongo.db.MONGO_OPERATION_TIMEOUT_MS
           }
-        }
-      ])
-      .toArray();
-    const groupedParameters = rows.map((row) => {
-      return row.bucket_parameters;
+        )
+        .toArray()
+        .catch((e) => {
+          throw lib_mongo.mapQueryError(e, 'while evaluating parameter queries');
+        });
+      const groupedParameters = rows.map((row) => {
+        return row.bucket_parameters;
+      });
+      return groupedParameters.flat();
     });
-    return groupedParameters.flat();
   }
 
   async *getBucketDataBatch(
@@ -444,24 +494,71 @@ export class MongoSyncBucketStorage
     return this.checksumCache.getChecksumMap(checkpoint, buckets);
   }
 
+  clearChecksumCache() {
+    this.checksumCache.clear();
+  }
+
   private async getChecksumsInternal(batch: storage.FetchPartialBucketChecksum[]): Promise<storage.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 filters: any[] = [];
     for (let request of batch) {
+      let start = request.start;
+      if (start == null) {
+        const preState = preStates.get(request.bucket);
+        if (preState != null) {
+          start = preState.opId;
+        }
+      }
+
       filters.push({
         _id: {
           $gt: {
             g: this.group_id,
             b: request.bucket,
-            o: request.start ? BigInt(request.start) : new bson.MinKey()
+            o: start ?? new bson.MinKey()
           },
           $lte: {
             g: this.group_id,
             b: request.bucket,
-            o: BigInt(request.end)
+            o: request.end
           }
         }
       });
@@ -491,26 +588,48 @@ export class MongoSyncBucketStorage
             }
           }
         ],
-        { session: undefined, readConcern: 'snapshot', maxTimeMS: lib_mongo.db.MONGO_OPERATION_TIMEOUT_MS }
+        { session: undefined, readConcern: 'snapshot', maxTimeMS: lib_mongo.db.MONGO_CHECKSUM_TIMEOUT_MS }
       )
       .toArray()
       .catch((e) => {
         throw lib_mongo.mapQueryError(e, 'while reading checksums');
       });
 
-    return new Map<string, storage.PartialChecksum>(
+    const partialChecksums = new Map<string, storage.PartialOrFullChecksum>(
       aggregate.map((doc) => {
+        const partialChecksum = Number(BigInt(doc.checksum_total) & 0xffffffffn) & 0xffffffff;
+        const bucket = doc._id;
         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
+          bucket,
+          doc.has_clear_op == 1
+            ? ({
+                // full checksum - replaces any previous one
+                bucket,
+                checksum: partialChecksum,
+                count: doc.count
+              } satisfies BucketChecksum)
+            : ({
+                // partial checksum - is added to a previous one
+                bucket,
+                partialCount: doc.count,
+                partialChecksum
+              } satisfies PartialChecksum)
         ];
       })
     );
+
+    return new Map<string, storage.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 = partialChecksums.get(bucket);
+        const merged = addPartialChecksums(bucket, preState?.checksum ?? null, partialChecksum ?? null);
+
+        return [bucket, merged];
+      })
+    );
   }
 
   async terminate(options?: storage.TerminateOptions) {
@@ -575,7 +694,6 @@ export class MongoSyncBucketStorage
             `${this.slot_name} Cleared batch of data in ${lib_mongo.db.MONGO_CLEAR_OPERATION_TIMEOUT_MS}ms, continuing...`
           );
           await timers.setTimeout(lib_mongo.db.MONGO_CLEAR_OPERATION_TIMEOUT_MS / 5);
-          continue;
         } else {
           throw e;
         }
@@ -640,41 +758,6 @@ export class MongoSyncBucketStorage
     );
   }
 
-  async autoActivate(): Promise<void> {
-    await this.db.client.withSession(async (session) => {
-      await session.withTransaction(async () => {
-        const doc = await this.db.sync_rules.findOne({ _id: this.group_id }, { session });
-        if (doc && doc.state == 'PROCESSING') {
-          await this.db.sync_rules.updateOne(
-            {
-              _id: this.group_id
-            },
-            {
-              $set: {
-                state: storage.SyncRuleState.ACTIVE
-              }
-            },
-            { session }
-          );
-
-          await this.db.sync_rules.updateMany(
-            {
-              _id: { $ne: this.group_id },
-              state: { $in: [storage.SyncRuleState.ACTIVE, storage.SyncRuleState.ERRORED] }
-            },
-            {
-              $set: {
-                state: storage.SyncRuleState.STOP
-              }
-            },
-            { session }
-          );
-          await this.db.notifyCheckpoint();
-        }
-      });
-    });
-  }
-
   async reportError(e: any): Promise<void> {
     const message = String(e.message ?? 'Replication failure');
     await this.db.sync_rules.updateOne(
@@ -691,14 +774,29 @@ export class MongoSyncBucketStorage
   }
 
   async compact(options?: storage.CompactOptions) {
-    return new MongoCompactor(this.db, this.group_id, options).compact();
+    let maxOpId = options?.maxOpId;
+    if (maxOpId == null) {
+      const checkpoint = await this.getCheckpointInternal();
+      maxOpId = checkpoint?.checkpoint ?? undefined;
+    }
+    await new MongoCompactor(this.db, this.group_id, { ...options, maxOpId }).compact();
+    if (maxOpId != null && options?.compactParameterData) {
+      await new MongoParameterCompactor(this.db, this.group_id, maxOpId, options).compact();
+    }
   }
 
-  private makeActiveCheckpoint(doc: SyncRuleCheckpointState | null) {
-    return {
-      checkpoint: doc?.last_checkpoint ?? 0n,
-      lsn: doc?.last_checkpoint_lsn ?? null
-    };
+  async populatePersistentChecksumCache(options: Pick<CompactOptions, 'signal' | 'maxOpId'>): Promise<void> {
+    const start = Date.now();
+    // We do a minimal compact, primarily to populate the checksum cache
+    await this.compact({
+      ...options,
+      // Skip parameter data
+      compactParameterData: false,
+      // Don't track updates for MOVE compacting
+      memoryLimitMB: 0
+    });
+    const duration = Date.now() - start;
+    logger.info(`Populated persistent checksum cache in ${(duration / 1000).toFixed(1)}s`);
   }
 
   /**
@@ -720,33 +818,13 @@ export class MongoSyncBucketStorage
         break;
       }
 
-      const doc = await this.db.sync_rules.findOne(
-        {
-          _id: this.group_id,
-          state: { $in: [storage.SyncRuleState.ACTIVE, storage.SyncRuleState.ERRORED] }
-        },
-        {
-          limit: 1,
-          projection: {
-            _id: 1,
-            state: 1,
-            last_checkpoint: 1,
-            last_checkpoint_lsn: 1
-          }
-        }
-      );
-
-      if (doc == null) {
-        // Sync rules not present or not active.
-        // Abort the connections - clients will have to retry later.
-        throw new ServiceError(ErrorCode.PSYNC_S2302, 'No active sync rules available');
-      } else if (doc.state != storage.SyncRuleState.ACTIVE && doc.state != storage.SyncRuleState.ERRORED) {
+      const op = await this.getCheckpointInternal();
+      if (op == null) {
         // Sync rules have changed - abort and restart.
         // We do a soft close of the stream here - no error
         break;
       }
 
-      const op = this.makeActiveCheckpoint(doc);
       // Check for LSN / checkpoint changes - ignore other metadata changes
       if (lastOp == null || op.lsn != lastOp.lsn || op.checkpoint != lastOp.checkpoint) {
         lastOp = op;
@@ -1013,3 +1091,25 @@ interface InternalCheckpointChanges extends CheckpointChanges {
   updatedWriteCheckpoints: Map<string, bigint>;
   invalidateWriteCheckpoints: boolean;
 }
+
+class MongoReplicationCheckpoint implements ReplicationCheckpoint {
+  constructor(
+    private storage: MongoSyncBucketStorage,
+    public readonly checkpoint: InternalOpId,
+    public readonly lsn: string | null,
+    public snapshotTime: mongo.Timestamp
+  ) {}
+
+  async getParameterSets(lookups: ParameterLookup[]): Promise<SqliteJsonRow[]> {
+    return this.storage.getParameterSets(this, lookups);
+  }
+}
+
+class EmptyReplicationCheckpoint implements ReplicationCheckpoint {
+  readonly checkpoint: InternalOpId = 0n;
+  readonly lsn: string | null = null;
+
+  async getParameterSets(lookups: ParameterLookup[]): Promise<SqliteJsonRow[]> {
+    return [];
+  }
+}

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;
 }