@powersync/service-module-mongodb-storage 0.9.4 → 0.10.0

package/src/storage/implementation/MongoWriteCheckpointAPI.ts

@@ -1,14 +1,7 @@
 import { mongo } from '@powersync/lib-service-mongodb';
 import * as framework from '@powersync/lib-services-framework';
-import {
-  Demultiplexer,
-  DemultiplexerValue,
-  storage,
-  WatchUserWriteCheckpointOptions,
-  WriteCheckpointResult
-} from '@powersync/service-core';
+import { GetCheckpointChangesOptions, InternalOpId, storage } from '@powersync/service-core';
 import { PowerSyncMongo } from './db.js';
-import { CustomWriteCheckpointDocument, WriteCheckpointDocument } from './models.js';
 
 export type MongoCheckpointAPIOptions = {
   db: PowerSyncMongo;
@@ -35,13 +28,9 @@ export class MongoWriteCheckpointAPI implements storage.WriteCheckpointAPI {
     this._mode = mode;
   }
 
-  async batchCreateCustomWriteCheckpoints(checkpoints: storage.CustomWriteCheckpointOptions[]): Promise<void> {
-    return batchCreateCustomWriteCheckpoints(this.db, checkpoints);
-  }
-
   async createManagedWriteCheckpoint(checkpoint: storage.ManagedWriteCheckpointOptions): Promise<bigint> {
     if (this.writeCheckpointMode !== storage.WriteCheckpointMode.MANAGED) {
-      throw new framework.errors.ValidationError(
+      throw new framework.ServiceAssertionError(
         `Attempting to create a managed Write Checkpoint when the current Write Checkpoint mode is set to "${this.writeCheckpointMode}"`
       );
     }
@@ -53,7 +42,8 @@ export class MongoWriteCheckpointAPI implements storage.WriteCheckpointAPI {
       },
       {
         $set: {
-          lsns
+          lsns,
+          processed_at_lsn: null
         },
         $inc: {
           client_id: 1n
@@ -68,12 +58,12 @@ export class MongoWriteCheckpointAPI implements storage.WriteCheckpointAPI {
     switch (this.writeCheckpointMode) {
       case storage.WriteCheckpointMode.CUSTOM:
         if (false == 'sync_rules_id' in filters) {
-          throw new framework.errors.ValidationError(`Sync rules ID is required for custom Write Checkpoint filtering`);
+          throw new framework.ServiceAssertionError(`Sync rules ID is required for custom Write Checkpoint filtering`);
         }
         return this.lastCustomWriteCheckpoint(filters);
       case storage.WriteCheckpointMode.MANAGED:
         if (false == 'heads' in filters) {
-          throw new framework.errors.ValidationError(
+          throw new framework.ServiceAssertionError(
             `Replication HEAD is required for managed Write Checkpoint filtering`
           );
         }
@@ -81,231 +71,15 @@ export class MongoWriteCheckpointAPI implements storage.WriteCheckpointAPI {
     }
   }
 
-  watchUserWriteCheckpoint(options: WatchUserWriteCheckpointOptions): AsyncIterable<storage.WriteCheckpointResult> {
+  async getWriteCheckpointChanges(options: GetCheckpointChangesOptions) {
     switch (this.writeCheckpointMode) {
       case storage.WriteCheckpointMode.CUSTOM:
-        return this.watchCustomWriteCheckpoint(options);
+        return this.getCustomWriteCheckpointChanges(options);
       case storage.WriteCheckpointMode.MANAGED:
-        return this.watchManagedWriteCheckpoint(options);
-      default:
-        throw new Error('Invalid write checkpoint mode');
-    }
-  }
-
-  private sharedManagedIter = new Demultiplexer<WriteCheckpointResult>((signal) => {
-    const clusterTimePromise = this.getClusterTime();
-
-    return {
-      iterator: this.watchAllManagedWriteCheckpoints(clusterTimePromise, signal),
-      getFirstValue: async (user_id: string) => {
-        // Potential race conditions we cater for:
-
-        // Case 1: changestream is behind.
-        // We get a doc now, then the same or older doc again later.
-        // No problem!
-
-        // Case 2: Query is behind. I.e. doc has been created, and emitted on the changestream, but the query doesn't see it yet.
-        // Not possible luckily, but can we make sure?
-
-        // Case 3: changestream delays openeing. A doc is created after our query here, but before the changestream is opened.
-        // Awaiting clusterTimePromise should be sufficient here, but as a sanity check we also confirm that our query
-        // timestamp is > the startClusterTime.
-
-        const changeStreamStart = await clusterTimePromise;
-
-        let doc = null as WriteCheckpointDocument | null;
-        let clusterTime = null as mongo.Timestamp | null;
-
-        await this.db.client.withSession(async (session) => {
-          doc = await this.db.write_checkpoints.findOne(
-            {
-              user_id: user_id
-            },
-            {
-              session
-            }
-          );
-          const time = session.clusterTime?.clusterTime ?? null;
-          clusterTime = time;
-        });
-        if (clusterTime == null) {
-          throw new framework.ServiceAssertionError('Could not get clusterTime for write checkpoint');
-        }
-
-        if (clusterTime.lessThan(changeStreamStart)) {
-          throw new framework.ServiceAssertionError(
-            'clusterTime for write checkpoint is older than changestream start'
-          );
-        }
-
-        if (doc == null) {
-          return {
-            id: null,
-            lsn: null
-          };
-        }
-
-        return {
-          id: doc.client_id,
-          lsn: doc.lsns['1']
-        };
-      }
-    };
-  });
-
-  private async *watchAllManagedWriteCheckpoints(
-    clusterTimePromise: Promise<mongo.BSON.Timestamp>,
-    signal: AbortSignal
-  ): AsyncGenerator<DemultiplexerValue<WriteCheckpointResult>> {
-    const clusterTime = await clusterTimePromise;
-
-    const stream = this.db.write_checkpoints.watch(
-      [{ $match: { operationType: { $in: ['insert', 'update', 'replace'] } } }],
-      {
-        fullDocument: 'updateLookup',
-        startAtOperationTime: clusterTime
-      }
-    );
-
-    signal.onabort = () => {
-      stream.close();
-    };
-
-    if (signal.aborted) {
-      stream.close();
-      return;
-    }
-
-    for await (let event of stream) {
-      if (!('fullDocument' in event) || event.fullDocument == null) {
-        continue;
-      }
-
-      const user_id = event.fullDocument.user_id;
-      yield {
-        key: user_id,
-        value: {
-          id: event.fullDocument.client_id,
-          lsn: event.fullDocument.lsns['1']
-        }
-      };
+        return this.getManagedWriteCheckpointChanges(options);
     }
   }
 
-  watchManagedWriteCheckpoint(options: WatchUserWriteCheckpointOptions): AsyncIterable<storage.WriteCheckpointResult> {
-    const stream = this.sharedManagedIter.subscribe(options.user_id, options.signal);
-    return this.orderedStream(stream);
-  }
-
-  private sharedCustomIter = new Demultiplexer<WriteCheckpointResult>((signal) => {
-    const clusterTimePromise = this.getClusterTime();
-
-    return {
-      iterator: this.watchAllCustomWriteCheckpoints(clusterTimePromise, signal),
-      getFirstValue: async (user_id: string) => {
-        // We cater for the same potential race conditions as for managed write checkpoints.
-
-        const changeStreamStart = await clusterTimePromise;
-
-        let doc = null as CustomWriteCheckpointDocument | null;
-        let clusterTime = null as mongo.Timestamp | null;
-
-        await this.db.client.withSession(async (session) => {
-          doc = await this.db.custom_write_checkpoints.findOne(
-            {
-              user_id: user_id,
-              sync_rules_id: this.sync_rules_id
-            },
-            {
-              session
-            }
-          );
-          const time = session.clusterTime?.clusterTime ?? null;
-          clusterTime = time;
-        });
-        if (clusterTime == null) {
-          throw new framework.ServiceAssertionError('Could not get clusterTime for write checkpoint');
-        }
-
-        if (clusterTime.lessThan(changeStreamStart)) {
-          throw new framework.ServiceAssertionError(
-            'clusterTime for write checkpoint is older than changestream start'
-          );
-        }
-
-        if (doc == null) {
-          // No write checkpoint, but we still need to return a result
-          return {
-            id: null,
-            lsn: null
-          };
-        }
-
-        return {
-          id: doc.checkpoint,
-          // custom write checkpoints are not tied to a LSN
-          lsn: null
-        };
-      }
-    };
-  });
-
-  private async *watchAllCustomWriteCheckpoints(
-    clusterTimePromise: Promise<mongo.BSON.Timestamp>,
-    signal: AbortSignal
-  ): AsyncGenerator<DemultiplexerValue<WriteCheckpointResult>> {
-    const clusterTime = await clusterTimePromise;
-
-    const stream = this.db.custom_write_checkpoints.watch(
-      [
-        {
-          $match: {
-            'fullDocument.sync_rules_id': this.sync_rules_id,
-            operationType: { $in: ['insert', 'update', 'replace'] }
-          }
-        }
-      ],
-      {
-        fullDocument: 'updateLookup',
-        startAtOperationTime: clusterTime
-      }
-    );
-
-    signal.onabort = () => {
-      stream.close();
-    };
-
-    if (signal.aborted) {
-      stream.close();
-      return;
-    }
-
-    for await (let event of stream) {
-      if (!('fullDocument' in event) || event.fullDocument == null) {
-        continue;
-      }
-
-      const user_id = event.fullDocument.user_id;
-      yield {
-        key: user_id,
-        value: {
-          id: event.fullDocument.checkpoint,
-          // Custom write checkpoints are not tied to a specific LSN
-          lsn: null
-        }
-      };
-    }
-  }
-
-  watchCustomWriteCheckpoint(options: WatchUserWriteCheckpointOptions): AsyncIterable<storage.WriteCheckpointResult> {
-    if (options.sync_rules_id != this.sync_rules_id) {
-      throw new framework.ServiceAssertionError('sync_rules_id does not match');
-    }
-
-    const stream = this.sharedCustomIter.subscribe(options.user_id, options.signal);
-    return this.orderedStream(stream);
-  }
-
   protected async lastCustomWriteCheckpoint(filters: storage.CustomWriteCheckpointFilters) {
     const { user_id, sync_rules_id } = filters;
     const lastWriteCheckpoint = await this.db.custom_write_checkpoints.findOne({
@@ -330,34 +104,70 @@ export class MongoWriteCheckpointAPI implements storage.WriteCheckpointAPI {
     return lastWriteCheckpoint?.client_id ?? null;
   }
 
-  private async getClusterTime(): Promise<mongo.Timestamp> {
-    const hello = await this.db.db.command({ hello: 1 });
-    // Note: This is not valid on sharded clusters.
-    const startClusterTime = hello.lastWrite?.majorityOpTime?.ts as mongo.Timestamp;
-    return startClusterTime;
-  }
+  private async getManagedWriteCheckpointChanges(options: GetCheckpointChangesOptions) {
+    const limit = 1000;
+    const changes = await this.db.write_checkpoints
+      .find(
+        {
+          processed_at_lsn: { $gt: options.lastCheckpoint.lsn, $lte: options.nextCheckpoint.lsn }
+        },
+        {
+          limit: limit + 1,
+          batchSize: limit + 1,
+          singleBatch: true
+        }
+      )
+      .toArray();
+    const invalidate = changes.length > limit;
+
+    const updatedWriteCheckpoints = new Map<string, bigint>();
+    if (!invalidate) {
+      for (let c of changes) {
+        updatedWriteCheckpoints.set(c.user_id, c.client_id);
+      }
+    }
 
-  /**
-   * Makes a write checkpoint stream an orderered one - any out-of-order events are discarded.
-   */
-  private async *orderedStream(stream: AsyncIterable<storage.WriteCheckpointResult>) {
-    let lastId = -1n;
+    return {
+      invalidateWriteCheckpoints: invalidate,
+      updatedWriteCheckpoints
+    };
+  }
 
-    for await (let event of stream) {
-      // Guard against out-of-order events
-      if (lastId == -1n || (event.id != null && event.id > lastId)) {
-        yield event;
-        if (event.id != null) {
-          lastId = event.id;
+  private async getCustomWriteCheckpointChanges(options: GetCheckpointChangesOptions) {
+    const limit = 1000;
+    const changes = await this.db.custom_write_checkpoints
+      .find(
+        {
+          op_id: { $gt: options.lastCheckpoint.checkpoint, $lte: options.nextCheckpoint.checkpoint }
+        },
+        {
+          limit: limit + 1,
+          batchSize: limit + 1,
+          singleBatch: true
         }
+      )
+      .toArray();
+    const invalidate = changes.length > limit;
+
+    const updatedWriteCheckpoints = new Map<string, bigint>();
+    if (!invalidate) {
+      for (let c of changes) {
+        updatedWriteCheckpoints.set(c.user_id, c.checkpoint);
       }
     }
+
+    return {
+      invalidateWriteCheckpoints: invalidate,
+      updatedWriteCheckpoints
+    };
   }
 }
 
 export async function batchCreateCustomWriteCheckpoints(
   db: PowerSyncMongo,
-  checkpoints: storage.CustomWriteCheckpointOptions[]
+  session: mongo.ClientSession,
+  checkpoints: storage.CustomWriteCheckpointOptions[],
+  opId: InternalOpId
 ): Promise<void> {
   if (checkpoints.length == 0) {
     return;
@@ -370,12 +180,13 @@ export async function batchCreateCustomWriteCheckpoints(
         update: {
           $set: {
             checkpoint: checkpointOptions.checkpoint,
-            sync_rules_id: checkpointOptions.sync_rules_id
+            sync_rules_id: checkpointOptions.sync_rules_id,
+            op_id: opId
           }
         },
         upsert: true
       }
     })),
-    {}
+    { session }
   );
 }

package/src/storage/implementation/PersistedBatch.ts

@@ -3,7 +3,7 @@ import { JSONBig } from '@powersync/service-jsonbig';
 import { EvaluatedParameters, EvaluatedRow } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
 
-import { logger } from '@powersync/lib-services-framework';
+import { Logger, logger as defaultLogger } from '@powersync/lib-services-framework';
 import { InternalOpId, storage, utils } from '@powersync/service-core';
 import { currentBucketKey, MAX_ROW_SIZE } from './MongoBucketBatch.js';
 import { MongoIdSequence } from './MongoIdSequence.js';
@@ -46,6 +46,7 @@ const MAX_TRANSACTION_DOC_COUNT = 2_000;
  * multiple transactions.
  */
 export class PersistedBatch {
+  logger: Logger;
   bucketData: mongo.AnyBulkWriteOperation<BucketDataDocument>[] = [];
   bucketParameters: mongo.AnyBulkWriteOperation<BucketParameterDocument>[] = [];
   currentData: mongo.AnyBulkWriteOperation<CurrentDataDocument>[] = [];
@@ -63,9 +64,11 @@ export class PersistedBatch {
 
   constructor(
     private group_id: number,
-    writtenSize: number
+    writtenSize: number,
+    options?: { logger?: Logger }
   ) {
     this.currentSize = writtenSize;
+    this.logger = options?.logger ?? defaultLogger;
   }
 
   private incrementBucket(bucket: string, op_id: InternalOpId) {
@@ -107,7 +110,7 @@ export class PersistedBatch {
         // the BSON size is small enough, but the JSON size is too large.
         // In these cases, we can't store the data, so we skip it, or generate a REMOVE operation if the row
         // was synced previously.
-        logger.error(`powersync_${this.group_id} Row ${key} too large: ${recordData.length} bytes. Removing.`);
+        this.logger.error(`Row ${key} too large: ${recordData.length} bytes. Removing.`);
         continue;
       }
 
@@ -206,7 +209,7 @@ export class PersistedBatch {
               k: sourceKey
             },
             lookup: binLookup,
-            bucket_parameters: result.bucket_parameters
+            bucket_parameters: result.bucketParameters
           }
         }
       });
@@ -270,9 +273,11 @@ export class PersistedBatch {
     );
   }
 
-  async flush(db: PowerSyncMongo, session: mongo.ClientSession) {
+  async flush(db: PowerSyncMongo, session: mongo.ClientSession, options?: storage.BucketBatchCommitOptions) {
     const startAt = performance.now();
+    let flushedSomething = false;
     if (this.bucketData.length > 0) {
+      flushedSomething = true;
       await db.bucket_data.bulkWrite(this.bucketData, {
         session,
         // inserts only - order doesn't matter
@@ -280,6 +285,7 @@ export class PersistedBatch {
       });
     }
     if (this.bucketParameters.length > 0) {
+      flushedSomething = true;
       await db.bucket_parameters.bulkWrite(this.bucketParameters, {
         session,
         // inserts only - order doesn't matter
@@ -287,6 +293,7 @@ export class PersistedBatch {
       });
     }
     if (this.currentData.length > 0) {
+      flushedSomething = true;
       await db.current_data.bulkWrite(this.currentData, {
         session,
         // may update and delete data within the same batch - order matters
@@ -295,6 +302,7 @@ export class PersistedBatch {
     }
 
     if (this.bucketStates.size > 0) {
+      flushedSomething = true;
       await db.bucket_state.bulkWrite(this.getBucketStateUpdates(), {
         session,
         // Per-bucket operation - order doesn't matter
@@ -302,12 +310,43 @@ export class PersistedBatch {
       });
     }
 
-    const duration = performance.now() - startAt;
-    logger.info(
-      `powersync_${this.group_id} Flushed ${this.bucketData.length} + ${this.bucketParameters.length} + ${
-        this.currentData.length
-      } updates, ${Math.round(this.currentSize / 1024)}kb in ${duration.toFixed(0)}ms. Last op_id: ${this.debugLastOpId}`
-    );
+    if (flushedSomething) {
+      const duration = Math.round(performance.now() - startAt);
+      if (options?.oldestUncommittedChange != null) {
+        const replicationLag = Math.round((Date.now() - options.oldestUncommittedChange.getTime()) / 1000);
+
+        this.logger.info(
+          `Flushed ${this.bucketData.length} + ${this.bucketParameters.length} + ${
+            this.currentData.length
+          } updates, ${Math.round(this.currentSize / 1024)}kb in ${duration}ms. Last op_id: ${this.debugLastOpId}. Replication lag: ${replicationLag}s`,
+          {
+            flushed: {
+              duration: duration,
+              size: this.currentSize,
+              bucket_data_count: this.bucketData.length,
+              parameter_data_count: this.bucketParameters.length,
+              current_data_count: this.currentData.length,
+              replication_lag_seconds: replicationLag
+            }
+          }
+        );
+      } else {
+        this.logger.info(
+          `Flushed ${this.bucketData.length} + ${this.bucketParameters.length} + ${
+            this.currentData.length
+          } updates, ${Math.round(this.currentSize / 1024)}kb in ${duration}ms. Last op_id: ${this.debugLastOpId}`,
+          {
+            flushed: {
+              duration: duration,
+              size: this.currentSize,
+              bucket_data_count: this.bucketData.length,
+              parameter_data_count: this.bucketParameters.length,
+              current_data_count: this.currentData.length
+            }
+          }
+        );
+      }
+    }
 
     this.bucketData = [];
     this.bucketParameters = [];

package/src/storage/implementation/db.ts

@@ -7,6 +7,7 @@ import {
   BucketDataDocument,
   BucketParameterDocument,
   BucketStateDocument,
+  CheckpointEventDocument,
   CurrentDataDocument,
   CustomWriteCheckpointDocument,
   IdSequenceDocument,
@@ -35,6 +36,7 @@ export class PowerSyncMongo {
   readonly instance: mongo.Collection<InstanceDocument>;
   readonly locks: mongo.Collection<lib_mongo.locks.Lock>;
   readonly bucket_state: mongo.Collection<BucketStateDocument>;
+  readonly checkpoint_events: mongo.Collection<CheckpointEventDocument>;
 
   readonly client: mongo.MongoClient;
   readonly db: mongo.Db;
@@ -58,6 +60,7 @@ export class PowerSyncMongo {
     this.instance = db.collection('instance');
     this.locks = this.db.collection('locks');
     this.bucket_state = this.db.collection('bucket_state');
+    this.checkpoint_events = this.db.collection('checkpoint_events');
   }
 
   /**
@@ -85,6 +88,45 @@ export class PowerSyncMongo {
   async drop() {
     await this.db.dropDatabase();
   }
+
+  /**
+   * Call this after every checkpoint or sync rules status update. Rather call too often than too rarely.
+   *
+   * This is used in a similar way to the Postgres NOTIFY functionality.
+   */
+  async notifyCheckpoint() {
+    await this.checkpoint_events.insertOne({} as any, { forceServerObjectId: true });
+  }
+
+  /**
+   * Only use in migrations and tests.
+   */
+  async createCheckpointEventsCollection() {
+    // We cover the case where the replication process was started before running this migration.
+    const existingCollections = await this.db
+      .listCollections({ name: 'checkpoint_events' }, { nameOnly: false })
+      .toArray();
+    const collection = existingCollections[0];
+    if (collection != null) {
+      if (!collection.options?.capped) {
+        // Collection was auto-created but not capped, so we need to drop it
+        await this.db.dropCollection('checkpoint_events');
+      } else {
+        // Collection previously created somehow - ignore
+        return;
+      }
+    }
+
+    await this.db.createCollection('checkpoint_events', {
+      capped: true,
+      // We want a small size, since opening a tailable cursor scans this entire collection.
+      // On the other hand, if we fill this up faster than a process can read it, it will
+      // invalidate the cursor. We do handle cursor invalidation events, but don't want
+      // that to happen too often.
+      size: 50 * 1024, // size in bytes
+      max: 50 // max number of documents
+    });
+  }
 }
 
 export function createPowerSyncMongo(config: MongoStorageConfig, options?: lib_mongo.MongoConnectionOptions) {

package/src/storage/implementation/models.ts

@@ -1,4 +1,4 @@
-import { storage } from '@powersync/service-core';
+import { InternalOpId, storage } from '@powersync/service-core';
 import { SqliteJsonValue } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
 
@@ -73,6 +73,13 @@ export interface SourceTableDocument {
   replica_id_columns: string[] | null;
   replica_id_columns2: { name: string; type_oid?: number; type?: string }[] | undefined;
   snapshot_done: boolean | undefined;
+  snapshot_status: SourceTableDocumentSnapshotStatus | undefined;
+}
+
+export interface SourceTableDocumentSnapshotStatus {
+  total_estimated_count: number;
+  replicated_count: number;
+  last_key: bson.Binary | null;
 }
 
 /**
@@ -110,6 +117,13 @@ export interface SyncRuleDocument {
    */
   snapshot_done: boolean;
 
+  /**
+   * If snapshot_done = false, this may be the lsn at which we started the snapshot.
+   *
+   * This can be used for resuming the snapshot after a restart.
+   */
+  snapshot_lsn: string | undefined;
+
   /**
    * The last consistent checkpoint.
    *
@@ -159,6 +173,10 @@ export interface SyncRuleDocument {
   content: string;
 }
 
+export interface CheckpointEventDocument {
+  _id: bson.ObjectId;
+}
+
 export type SyncRuleCheckpointState = Pick<
   SyncRuleDocument,
   'last_checkpoint' | 'last_checkpoint_lsn' | '_id' | 'state'
@@ -169,6 +187,13 @@ export interface CustomWriteCheckpointDocument {
   user_id: string;
   checkpoint: bigint;
   sync_rules_id: number;
+  /**
+   * Unlike managed write checkpoints, custom write checkpoints are flushed together with
+   * normal ops. This means we can assign an op_id for ordering / correlating with read checkpoints.
+   *
+   * This is not unique - multiple write checkpoints can have the same op_id.
+   */
+  op_id?: InternalOpId;
 }
 
 export interface WriteCheckpointDocument {
@@ -176,6 +201,12 @@ export interface WriteCheckpointDocument {
   user_id: string;
   lsns: Record<string, string>;
   client_id: bigint;
+  /**
+   * This is set to the checkpoint lsn when the checkpoint lsn >= this lsn.
+   * This is used to make it easier to determine what write checkpoints have been processed
+   * between two checkpoints.
+   */
+  processed_at_lsn: string | null;
 }
 
 export interface InstanceDocument {