@powersync/service-module-mongodb-storage 0.7.2 → 0.8.1

package/src/storage/implementation/MongoSyncBucketStorage.ts

@@ -11,25 +11,29 @@ import {
   BroadcastIterable,
   CHECKPOINT_INVALIDATE_ALL,
   CheckpointChanges,
+  deserializeParameterLookup,
   GetCheckpointChangesOptions,
   InternalOpId,
   internalToExternalOpId,
+  mergeAsyncIterables,
   ProtocolOpId,
   ReplicationCheckpoint,
-  SourceTable,
   storage,
   utils,
-  WatchWriteCheckpointOptions
+  WatchWriteCheckpointOptions,
+  WriteCheckpointResult
 } from '@powersync/service-core';
-import { SqliteJsonRow, SqliteJsonValue, SqlSyncRules } from '@powersync/service-sync-rules';
+import { JSONBig } from '@powersync/service-jsonbig';
+import { ParameterLookup, SqliteJsonRow, 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,
@@ -65,7 +69,8 @@ export class MongoSyncBucketStorage
     this.db = factory.db;
     this.writeCheckpointAPI = new MongoWriteCheckpointAPI({
       db: this.db,
-      mode: writeCheckpointMode
+      mode: writeCheckpointMode,
+      sync_rules_id: group_id
     });
   }
 
@@ -83,13 +88,6 @@ export class MongoSyncBucketStorage
     );
   }
 
-  createCustomWriteCheckpoint(checkpoint: storage.BatchedCustomWriteCheckpointOptions): Promise<bigint> {
-    return this.writeCheckpointAPI.createCustomWriteCheckpoint({
-      ...checkpoint,
-      sync_rules_id: this.group_id
-    });
-  }
-
   createManagedWriteCheckpoint(checkpoint: storage.ManagedWriteCheckpointOptions): Promise<bigint> {
     return this.writeCheckpointAPI.createManagedWriteCheckpoint(checkpoint);
   }
@@ -154,7 +152,7 @@ export class MongoSyncBucketStorage
 
     await callback(batch);
     await batch.flush();
-    if (batch.last_flushed_op) {
+    if (batch.last_flushed_op != null) {
       return { flushed_op: batch.last_flushed_op };
     } else {
       return null;
@@ -252,7 +250,7 @@ export class MongoSyncBucketStorage
     return result!;
   }
 
-  async getParameterSets(checkpoint: utils.InternalOpId, lookups: SqliteJsonValue[][]): Promise<SqliteJsonRow[]> {
+  async getParameterSets(checkpoint: utils.InternalOpId, lookups: ParameterLookup[]): Promise<SqliteJsonRow[]> {
     const lookupFilter = lookups.map((lookup) => {
       return storage.serializeLookup(lookup);
     });
@@ -585,6 +583,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
@@ -694,8 +699,7 @@ export class MongoSyncBucketStorage
     if (doc == null) {
       // Sync rules not present or not active.
       // Abort the connections - clients will have to retry later.
-      // Should this error instead?
-      return;
+      throw new ServiceError(ErrorCode.PSYNC_S2302, 'No active sync rules available');
     }
 
     yield this.makeActiveCheckpoint(doc);
@@ -739,7 +743,7 @@ export class MongoSyncBucketStorage
       }
       if (doc.state != storage.SyncRuleState.ACTIVE && doc.state != storage.SyncRuleState.ERRORED) {
         // Sync rules have changed - abort and restart.
-        // Should this error instead?
+        // We do a soft close of the stream here - no error
         break;
       }
 
@@ -762,28 +766,60 @@ export class MongoSyncBucketStorage
   /**
    * User-specific watch on the latest checkpoint and/or write checkpoint.
    */
-  async *watchWriteCheckpoint(options: WatchWriteCheckpointOptions): AsyncIterable<storage.StorageCheckpointUpdate> {
-    const { user_id, signal } = options;
+  async *watchCheckpointChanges(options: WatchWriteCheckpointOptions): AsyncIterable<storage.StorageCheckpointUpdate> {
+    const { signal } = options;
     let lastCheckpoint: utils.InternalOpId | null = null;
     let lastWriteCheckpoint: bigint | null = null;
+    let lastWriteCheckpointDoc: WriteCheckpointResult | null = null;
+    let nextWriteCheckpoint: bigint | null = null;
+    let lastCheckpointEvent: ReplicationCheckpoint | null = null;
+    let receivedWriteCheckpoint = false;
+
+    const writeCheckpointIter = this.writeCheckpointAPI.watchUserWriteCheckpoint({
+      user_id: options.user_id,
+      signal,
+      sync_rules_id: this.group_id
+    });
+    const iter = mergeAsyncIterables<ReplicationCheckpoint | storage.WriteCheckpointResult>(
+      [this.sharedIter, writeCheckpointIter],
+      signal
+    );
 
-    const iter = wrapWithAbort(this.sharedIter, signal);
     for await (const event of iter) {
-      const { checkpoint, lsn } = event;
+      if ('checkpoint' in event) {
+        lastCheckpointEvent = event;
+      } else {
+        lastWriteCheckpointDoc = event;
+        receivedWriteCheckpoint = true;
+      }
+
+      if (lastCheckpointEvent == null || !receivedWriteCheckpoint) {
+        // We need to wait until we received at least on checkpoint, and one write checkpoint.
+        continue;
+      }
 
       // lsn changes are not important by itself.
       // What is important is:
       // 1. checkpoint (op_id) changes.
       // 2. write checkpoint changes for the specific user
 
-      const lsnFilters: Record<string, string> = lsn ? { 1: lsn } : {};
+      const lsn = lastCheckpointEvent?.lsn;
 
-      const currentWriteCheckpoint = await this.lastWriteCheckpoint({
-        user_id,
-        heads: {
-          ...lsnFilters
+      if (
+        lastWriteCheckpointDoc != null &&
+        (lastWriteCheckpointDoc.lsn == null || (lsn != null && lsn >= lastWriteCheckpointDoc.lsn))
+      ) {
+        const writeCheckpoint = lastWriteCheckpointDoc.id;
+        if (nextWriteCheckpoint == null || (writeCheckpoint != null && writeCheckpoint > nextWriteCheckpoint)) {
+          nextWriteCheckpoint = writeCheckpoint;
         }
-      });
+        // We used the doc - clear it
+        lastWriteCheckpointDoc = null;
+      }
+
+      const { checkpoint } = lastCheckpointEvent;
+
+      const currentWriteCheckpoint = nextWriteCheckpoint;
 
       if (currentWriteCheckpoint == lastWriteCheckpoint && checkpoint == lastCheckpoint) {
         // No change - wait for next one
@@ -795,12 +831,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
@@ -810,7 +841,7 @@ export class MongoSyncBucketStorage
       lastCheckpoint = checkpoint;
 
       yield {
-        base: event,
+        base: lastCheckpointEvent,
         writeCheckpoint: currentWriteCheckpoint,
         update: updates
       };
@@ -869,7 +900,105 @@ export class MongoSyncBucketStorage
     return pipeline;
   }
 
+  private async getDataBucketChanges(
+    options: GetCheckpointChangesOptions
+  ): Promise<Pick<CheckpointChanges, 'updatedDataBuckets' | 'invalidateDataBuckets'>> {
+    const limit = 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: limit + 1,
+          batchSize: limit + 1,
+          singleBatch: true
+        }
+      )
+      .toArray();
+
+    const buckets = bucketStateUpdates.map((doc) => doc._id.b);
+    const invalidateDataBuckets = buckets.length > limit;
+
+    return {
+      invalidateDataBuckets: invalidateDataBuckets,
+      updatedDataBuckets: invalidateDataBuckets ? new Set<string>() : new Set(buckets)
+    };
+  }
+
+  private async getParameterBucketChanges(
+    options: GetCheckpointChangesOptions
+  ): Promise<Pick<CheckpointChanges, 'updatedParameterLookups' | 'invalidateParameterBuckets'>> {
+    const limit = 1000;
+    const parameterUpdates = await this.db.bucket_parameters
+      .find(
+        {
+          _id: { $gt: BigInt(options.lastCheckpoint), $lte: BigInt(options.nextCheckpoint) },
+          'key.g': this.group_id
+        },
+        {
+          projection: {
+            lookup: 1
+          },
+          limit: limit + 1,
+          batchSize: limit + 1,
+          singleBatch: true
+        }
+      )
+      .toArray();
+    const invalidateParameterUpdates = parameterUpdates.length > limit;
+
+    return {
+      invalidateParameterBuckets: invalidateParameterUpdates,
+      updatedParameterLookups: invalidateParameterUpdates
+        ? new Set<string>()
+        : new Set<string>(parameterUpdates.map((p) => JSONBig.stringify(deserializeParameterLookup(p.lookup))))
+    };
+  }
+
+  // If we processed all connections together for each checkpoint, we could do a single lookup for all connections.
+  // In practice, specific connections may fall behind. So instead, we just cache the results of each specific lookup.
+  // TODO (later):
+  // 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 }>({
+    // Limit to 50 cache entries, or 10MB, whichever comes first.
+    // Some rough calculations:
+    // If we process 10 checkpoints per second, and a connection may be 2 seconds behind, we could have
+    // up to 20 relevant checkpoints. That gives us 20*20 = 400 potentially-relevant cache entries.
+    // That is a worst-case scenario, so we don't actually store that many. In real life, the cache keys
+    // would likely be clustered around a few values, rather than spread over all 400 potential values.
+    max: 50,
+    maxSize: 10 * 1024 * 1024,
+    sizeCalculation: (value: CheckpointChanges) => {
+      // Estimate of memory usage
+      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);
+    }
+  });
+
   async getCheckpointChanges(options: GetCheckpointChangesOptions): Promise<CheckpointChanges> {
-    return CHECKPOINT_INVALIDATE_ALL;
+    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/MongoWriteCheckpointAPI.ts

@@ -1,19 +1,30 @@
+import { mongo } from '@powersync/lib-service-mongodb';
 import * as framework from '@powersync/lib-services-framework';
-import { storage } from '@powersync/service-core';
+import {
+  Demultiplexer,
+  DemultiplexerValue,
+  storage,
+  WatchUserWriteCheckpointOptions,
+  WriteCheckpointResult
+} from '@powersync/service-core';
 import { PowerSyncMongo } from './db.js';
+import { CustomWriteCheckpointDocument, WriteCheckpointDocument } from './models.js';
 
 export type MongoCheckpointAPIOptions = {
   db: PowerSyncMongo;
   mode: storage.WriteCheckpointMode;
+  sync_rules_id: number;
 };
 
 export class MongoWriteCheckpointAPI implements storage.WriteCheckpointAPI {
   readonly db: PowerSyncMongo;
   private _mode: storage.WriteCheckpointMode;
+  private sync_rules_id: number;
 
   constructor(options: MongoCheckpointAPIOptions) {
     this.db = options.db;
     this._mode = options.mode;
+    this.sync_rules_id = options.sync_rules_id;
   }
 
   get writeCheckpointMode() {
@@ -28,29 +39,6 @@ export class MongoWriteCheckpointAPI implements storage.WriteCheckpointAPI {
     return batchCreateCustomWriteCheckpoints(this.db, checkpoints);
   }
 
-  async createCustomWriteCheckpoint(options: storage.CustomWriteCheckpointOptions): Promise<bigint> {
-    if (this.writeCheckpointMode !== storage.WriteCheckpointMode.CUSTOM) {
-      throw new framework.errors.ValidationError(
-        `Creating a custom Write Checkpoint when the current Write Checkpoint mode is set to "${this.writeCheckpointMode}"`
-      );
-    }
-
-    const { checkpoint, user_id, sync_rules_id } = options;
-    const doc = await this.db.custom_write_checkpoints.findOneAndUpdate(
-      {
-        user_id: user_id,
-        sync_rules_id
-      },
-      {
-        $set: {
-          checkpoint
-        }
-      },
-      { upsert: true, returnDocument: 'after' }
-    );
-    return doc!.checkpoint;
-  }
-
   async createManagedWriteCheckpoint(checkpoint: storage.ManagedWriteCheckpointOptions): Promise<bigint> {
     if (this.writeCheckpointMode !== storage.WriteCheckpointMode.MANAGED) {
       throw new framework.errors.ValidationError(
@@ -93,6 +81,231 @@ export class MongoWriteCheckpointAPI implements storage.WriteCheckpointAPI {
     }
   }
 
+  watchUserWriteCheckpoint(options: WatchUserWriteCheckpointOptions): AsyncIterable<storage.WriteCheckpointResult> {
+    switch (this.writeCheckpointMode) {
+      case storage.WriteCheckpointMode.CUSTOM:
+        return this.watchCustomWriteCheckpoint(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']
+        }
+      };
+    }
+  }
+
+  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({
@@ -116,13 +329,37 @@ 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;
+  }
+
+  /**
+   * 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;
+
+    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;
+        }
+      }
+    }
+  }
 }
 
 export async function batchCreateCustomWriteCheckpoints(
   db: PowerSyncMongo,
   checkpoints: storage.CustomWriteCheckpointOptions[]
 ): Promise<void> {
-  if (!checkpoints.length) {
+  if (checkpoints.length == 0) {
     return;
   }