@powersync/service-core 0.8.8 → 0.9.0

package/src/runner/teardown.ts

@@ -1,108 +1,70 @@
 // Script to tear down the data when deleting an instance.
-// This deletes:
-// 1. The replication slots on the source postgres instance (if available).
-// 2. The mongo database.
+// This should:
+// 1. Attempt to clean up any remote configuration of data sources that was set up.
+// 2. Delete the storage
 
-import * as timers from 'timers/promises';
-
-import * as db from '../db/db-index.js';
+import { container, logger } from '@powersync/lib-services-framework';
+import timers from 'timers/promises';
+import * as modules from '../modules/modules-index.js';
 import * as storage from '../storage/storage-index.js';
+import * as system from '../system/system-index.js';
 import * as utils from '../util/util-index.js';
-import * as replication from '../replication/replication-index.js';
-import { logger } from '@powersync/lib-services-framework';
 
-/**
- * Attempt to terminate a single sync rules instance.
- *
- * This may fail with a lock error.
- */
-async function terminateReplicator(
-  storageFactory: storage.BucketStorageFactory,
-  connection: utils.ResolvedConnection,
-  syncRules: storage.PersistedSyncRulesContent
-) {
-  // The lock may still be active if the current replication instance
-  // hasn't stopped yet.
-  const lock = await syncRules.lock();
+export async function teardown(runnerConfig: utils.RunnerConfig) {
   try {
-    const parsed = syncRules.parsed();
-    const storage = storageFactory.getInstance(parsed);
-    const stream = new replication.WalStreamRunner({
-      factory: storageFactory,
-      storage: storage,
-      source_db: connection,
-      lock
-    });
+    logger.info(`Tearing down PowerSync instance...`);
+    const config = await utils.loadConfig(runnerConfig);
+    const serviceContext = new system.ServiceContextContainer(config);
+    const moduleManager = container.getImplementation(modules.ModuleManager);
+    await moduleManager.initialize(serviceContext);
+    // This is mostly done to ensure that the storage is ready
+    await serviceContext.lifeCycleEngine.start();
 
-    logger.info(`Terminating replication slot ${stream.slot_name}`);
-    await stream.terminate();
-    logger.info(`Terminated replication slot ${stream.slot_name}`);
-  } finally {
-    await lock.release();
+    await terminateSyncRules(serviceContext.storageEngine.activeBucketStorage, moduleManager);
+    await serviceContext.storageEngine.activeStorage.tearDown();
+    logger.info(`Teardown complete.`);
+    process.exit(0);
+  } catch (e) {
+    logger.error(`Teardown failure`, e);
+    process.exit(1);
   }
 }
 
-/**
- * Terminate all replicating sync rules, deleting the replication slots.
- *
- * Retries lock and other errors for up to two minutes.
- *
- * This is a best-effot attempt. In some cases it may not be possible to delete the replication
- * slot, such as when the postgres instance is unreachable.
- */
-async function terminateReplicators(
-  storageFactory: storage.BucketStorageFactory,
-  connection: utils.ResolvedConnection
-) {
+async function terminateSyncRules(storageFactory: storage.BucketStorageFactory, moduleManager: modules.ModuleManager) {
+  logger.info(`Terminating sync rules...`);
   const start = Date.now();
-  while (Date.now() - start < 12_000) {
+  const locks: storage.ReplicationLock[] = [];
+  while (Date.now() - start < 120_000) {
     let retry = false;
-    const replicationRules = await storageFactory.getReplicatingSyncRules();
-    for (let syncRules of replicationRules) {
-      try {
-        await terminateReplicator(storageFactory, connection, syncRules);
-      } catch (e) {
-        retry = true;
-        console.error(e);
-        logger.warn(`Failed to terminate ${syncRules.slot_name}`, e);
+    const replicatingSyncRules = await storageFactory.getReplicatingSyncRules();
+    // Lock all the replicating sync rules
+    for (const replicatingSyncRule of replicatingSyncRules) {
+      const lock = await replicatingSyncRule.lock();
+      locks.push(lock);
+    }
+
+    const stoppedSyncRules = await storageFactory.getStoppedSyncRules();
+    const combinedSyncRules = [...replicatingSyncRules, ...stoppedSyncRules];
+    try {
+      // Clean up any module specific configuration for the sync rules
+      await moduleManager.tearDown({ syncRules: combinedSyncRules });
+
+      // Mark the sync rules as terminated
+      for (let syncRules of combinedSyncRules) {
+        using syncRulesStorage = storageFactory.getInstance(syncRules);
+        // The storage will be dropped at the end of the teardown, so we don't need to clear it here
+        await syncRulesStorage.terminate({ clearStorage: false });
+      }
+    } catch (e) {
+      retry = true;
+      for (const lock of locks) {
+        await lock.release();
       }
     }
+
     if (!retry) {
       break;
     }
     await timers.setTimeout(5_000);
   }
 }
-
-export async function teardown(runnerConfig: utils.RunnerConfig) {
-  const config = await utils.loadConfig(runnerConfig);
-  const mongoDB = storage.createPowerSyncMongo(config.storage);
-  try {
-    logger.info(`Waiting for auth`);
-    await db.mongo.waitForAuth(mongoDB.db);
-
-    const bucketStorage = new storage.MongoBucketStorage(mongoDB, { slot_name_prefix: config.slot_name_prefix });
-    const connection = config.connection;
-
-    logger.info(`Terminating replication slots`);
-
-    if (connection) {
-      await terminateReplicators(bucketStorage, connection);
-    }
-
-    const database = mongoDB.db;
-    logger.info(`Dropping database ${database.namespace}`);
-    await database.dropDatabase();
-    logger.info(`Done`);
-    await mongoDB.client.close();
-
-    // If there was an error connecting to postgress, the process may stay open indefinitely.
-    // This forces an exit.
-    // We do not consider those errors a teardown failure.
-    process.exit(0);
-  } catch (e) {
-    logger.error(`Teardown failure`, e);
-    await mongoDB.client.close();
-    process.exit(1);
-  }
-}

package/src/storage/BucketStorage.ts

@@ -1,3 +1,4 @@
+import { DisposableListener, DisposableObserverClient } from '@powersync/lib-services-framework';
 import {
   EvaluatedParameters,
   EvaluatedRow,
@@ -7,12 +8,19 @@ import {
   SqliteRow,
   ToastableSqliteRow
 } from '@powersync/service-sync-rules';
-
-import * as replication from '../replication/replication-index.js';
 import * as util from '../util/util-index.js';
+import { ReplicationEventPayload } from './ReplicationEventPayload.js';
+import { SourceEntityDescriptor } from './SourceEntity.js';
 import { SourceTable } from './SourceTable.js';
+import { BatchedCustomWriteCheckpointOptions, ReplicaId } from './storage-index.js';
+import { SyncStorageWriteCheckpointAPI } from './WriteCheckpointAPI.js';
+
+export interface BucketStorageFactoryListener extends DisposableListener {
+  syncStorageCreated: (storage: SyncRulesBucketStorage) => void;
+  replicationEvent: (event: ReplicationEventPayload) => void;
+}
 
-export interface BucketStorageFactory {
+export interface BucketStorageFactory extends DisposableObserverClient<BucketStorageFactoryListener> {
   /**
    * Update sync rules from configuration, if changed.
    */
@@ -24,7 +32,7 @@ export interface BucketStorageFactory {
   /**
    * Get a storage instance to query sync data for specific sync rules.
    */
-  getInstance(options: PersistedSyncRules): SyncRulesBucketStorage;
+  getInstance(options: PersistedSyncRulesContent): SyncRulesBucketStorage;
 
   /**
    * Deploy new sync rules.
@@ -48,7 +56,7 @@ export interface BucketStorageFactory {
   /**
    * Get the sync rules used for querying.
    */
-  getActiveSyncRules(): Promise<PersistedSyncRules | null>;
+  getActiveSyncRules(options: ParseSyncRulesOptions): Promise<PersistedSyncRules | null>;
 
   /**
    * Get the sync rules used for querying.
@@ -58,7 +66,7 @@ export interface BucketStorageFactory {
   /**
    * Get the sync rules that will be active next once done with initial replicatino.
    */
-  getNextSyncRules(): Promise<PersistedSyncRules | null>;
+  getNextSyncRules(options: ParseSyncRulesOptions): Promise<PersistedSyncRules | null>;
 
   /**
    * Get the sync rules that will be active next once done with initial replicatino.
@@ -81,10 +89,9 @@ export interface BucketStorageFactory {
    */
   getActiveCheckpoint(): Promise<ActiveCheckpoint>;
 
-  createWriteCheckpoint(user_id: string, lsns: Record<string, string>): Promise<bigint>;
-
-  lastWriteCheckpoint(user_id: string, lsn: string): Promise<bigint | null>;
-
+  /**
+   * Yields the latest user write checkpoint whenever the sync checkpoint updates.
+   */
   watchWriteCheckpoint(user_id: string, signal: AbortSignal): AsyncIterable<WriteCheckpoint>;
 
   /**
@@ -98,20 +105,22 @@ export interface BucketStorageFactory {
   getPowerSyncInstanceId(): Promise<string>;
 }
 
-export interface WriteCheckpoint {
-  base: ActiveCheckpoint;
-  writeCheckpoint: bigint | null;
-}
-
-export interface ActiveCheckpoint {
+export interface ReplicationCheckpoint {
   readonly checkpoint: util.OpId;
-  readonly lsn: string;
+  readonly lsn: string | null;
+}
 
+export interface ActiveCheckpoint extends ReplicationCheckpoint {
   hasSyncRules(): boolean;
 
   getBucketStorage(): Promise<SyncRulesBucketStorage | null>;
 }
 
+export interface WriteCheckpoint {
+  base: ActiveCheckpoint;
+  writeCheckpoint: bigint | null;
+}
+
 export interface StorageMetrics {
   /**
    * Size of operations (bucket_data)
@@ -131,6 +140,10 @@ export interface StorageMetrics {
   replication_size_bytes: number;
 }
 
+export interface ParseSyncRulesOptions {
+  defaultSchema: string;
+}
+
 export interface PersistedSyncRulesContent {
   readonly id: number;
   readonly sync_rules_content: string;
@@ -140,7 +153,7 @@ export interface PersistedSyncRulesContent {
   readonly last_keepalive_ts?: Date | null;
   readonly last_checkpoint_ts?: Date | null;
 
-  parsed(): PersistedSyncRules;
+  parsed(options: ParseSyncRulesOptions): PersistedSyncRules;
 
   lock(): Promise<ReplicationLock>;
 }
@@ -157,18 +170,6 @@ export interface PersistedSyncRules {
   readonly slot_name: string;
 }
 
-export class DefaultPersistedSyncRules implements PersistedSyncRules {
-  public readonly checkpoint_lsn: string | null;
-
-  constructor(public readonly id: number, public readonly sync_rules: SqlSyncRules, checkpoint_lsn: string | null) {
-    this.checkpoint_lsn = checkpoint_lsn;
-  }
-
-  get slot_name(): string {
-    return `powersync_${this.id}`;
-  }
-}
-
 export interface UpdateSyncRulesOptions {
   content: string;
   lock?: boolean;
@@ -198,8 +199,27 @@ export interface BucketDataBatchOptions {
   chunkLimitBytes?: number;
 }
 
-export interface SyncRulesBucketStorage {
-  readonly sync_rules: SqlSyncRules;
+export interface StartBatchOptions extends ParseSyncRulesOptions {
+  zeroLSN: string;
+  /**
+   * Whether or not to store a copy of the current data.
+   *
+   * This is needed if we need to apply partial updates, for example
+   * when we get TOAST values from Postgres.
+   *
+   * This is not needed when we get the full document from the source
+   * database, for example from MongoDB.
+   */
+  storeCurrentData: boolean;
+}
+
+export interface SyncRulesBucketStorageListener extends DisposableListener {
+  batchStarted: (batch: BucketStorageBatch) => void;
+}
+
+export interface SyncRulesBucketStorage
+  extends DisposableObserverClient<SyncRulesBucketStorageListener>,
+    SyncStorageWriteCheckpointAPI {
   readonly group_id: number;
   readonly slot_name: string;
 
@@ -207,9 +227,14 @@ export interface SyncRulesBucketStorage {
 
   resolveTable(options: ResolveTableOptions): Promise<ResolveTableResult>;
 
-  startBatch(options: {}, callback: (batch: BucketStorageBatch) => Promise<void>): Promise<FlushedResult | null>;
+  startBatch(
+    options: StartBatchOptions,
+    callback: (batch: BucketStorageBatch) => Promise<void>
+  ): Promise<FlushedResult | null>;
+
+  getCheckpoint(): Promise<ReplicationCheckpoint>;
 
-  getCheckpoint(): Promise<{ checkpoint: util.OpId; lsn: string }>;
+  getParsedSyncRules(options: ParseSyncRulesOptions): SqlSyncRules;
 
   getParameterSets(checkpoint: util.OpId, lookups: SqliteJsonValue[][]): Promise<SqliteJsonRow[]>;
 
@@ -244,7 +269,7 @@ export interface SyncRulesBucketStorage {
    *
    * Must only be called on stopped sync rules.
    */
-  terminate(): Promise<void>;
+  terminate(options?: TerminateOptions): Promise<void>;
 
   getStatus(): Promise<SyncRuleStatus>;
 
@@ -277,7 +302,7 @@ export interface ResolveTableOptions {
   group_id: number;
   connection_id: number;
   connection_tag: string;
-  relation: replication.PgRelation;
+  entity_descriptor: SourceEntityDescriptor;
 
   sync_rules: SqlSyncRules;
 }
@@ -291,7 +316,11 @@ export interface FlushedResult {
   flushed_op: string;
 }
 
-export interface BucketStorageBatch {
+export interface BucketBatchStorageListener extends DisposableListener {
+  replicationEvent: (payload: ReplicationEventPayload) => void;
+}
+
+export interface BucketStorageBatch extends DisposableObserverClient<BucketBatchStorageListener> {
   /**
    * Save an op, and potentially flush.
    *
@@ -337,7 +366,17 @@ export interface BucketStorageBatch {
    */
   keepalive(lsn: string): Promise<boolean>;
 
+  /**
+   * Get the last checkpoint LSN, from either commit or keepalive.
+   */
+  lastCheckpointLsn: string | null;
+
   markSnapshotDone(tables: SourceTable[], no_checkpoint_before_lsn: string): Promise<SourceTable[]>;
+
+  /**
+   * Queues the creation of a custom Write Checkpoint. This will be persisted after operations are flushed.
+   */
+  addCustomWriteCheckpoint(checkpoint: BatchedCustomWriteCheckpointOptions): void;
 }
 
 export interface SaveParameterData {
@@ -355,23 +394,34 @@ export interface SaveBucketData {
   evaluated: EvaluatedRow[];
 }
 
+export type SaveOp = 'insert' | 'update' | 'delete';
+
 export type SaveOptions = SaveInsert | SaveUpdate | SaveDelete;
 
+export enum SaveOperationTag {
+  INSERT = 'insert',
+  UPDATE = 'update',
+  DELETE = 'delete'
+}
+
 export interface SaveInsert {
-  tag: 'insert';
+  tag: SaveOperationTag.INSERT;
   sourceTable: SourceTable;
   before?: undefined;
+  beforeReplicaId?: undefined;
   after: SqliteRow;
+  afterReplicaId: ReplicaId;
 }
 
 export interface SaveUpdate {
-  tag: 'update';
+  tag: SaveOperationTag.UPDATE;
   sourceTable: SourceTable;
 
   /**
    * This is only present when the id has changed, and will only contain replica identity columns.
    */
   before?: SqliteRow;
+  beforeReplicaId?: ReplicaId;
 
   /**
    * A null value means null column.
@@ -379,13 +429,16 @@ export interface SaveUpdate {
    * An undefined value means it's a TOAST value - must be copied from another record.
    */
   after: ToastableSqliteRow;
+  afterReplicaId: ReplicaId;
 }
 
 export interface SaveDelete {
-  tag: 'delete';
+  tag: SaveOperationTag.DELETE;
   sourceTable: SourceTable;
-  before: SqliteRow;
+  before?: SqliteRow;
+  beforeReplicaId: ReplicaId;
   after?: undefined;
+  afterReplicaId?: undefined;
 }
 
 export interface SyncBucketDataBatch {
@@ -429,6 +482,15 @@ export interface CompactOptions {
    * If specified, compact only the specific buckets.
    *
    * If not specified, compacts all buckets.
+   *
+   * These can be individual bucket names, or bucket definition names.
    */
   compactBuckets?: string[];
 }
+
+export interface TerminateOptions {
+  /**
+   * If true, also clear the storage before terminating.
+   */
+  clearStorage: boolean;
+}

package/src/storage/MongoBucketStorage.ts

@@ -1,36 +1,39 @@
-import * as mongo from 'mongodb';
-import * as timers from 'timers/promises';
-import { LRUCache } from 'lru-cache/min';
 import { SqlSyncRules } from '@powersync/service-sync-rules';
 import { wrapWithAbort } from 'ix/asynciterable/operators/withabort.js';
+import { LRUCache } from 'lru-cache/min';
+import * as mongo from 'mongodb';
+import * as timers from 'timers/promises';
 
-import * as replication from '../replication/replication-index.js';
+import * as locks from '../locks/locks-index.js';
 import * as sync from '../sync/sync-index.js';
 import * as util from '../util/util-index.js';
-import * as locks from '../locks/locks-index.js';
 
+import { DisposableObserver, logger } from '@powersync/lib-services-framework';
+import { v4 as uuid } from 'uuid';
 import {
   ActiveCheckpoint,
   BucketStorageFactory,
+  BucketStorageFactoryListener,
+  ParseSyncRulesOptions,
   PersistedSyncRules,
   PersistedSyncRulesContent,
   StorageMetrics,
   UpdateSyncRulesOptions,
   WriteCheckpoint
 } from './BucketStorage.js';
+import { PowerSyncMongo } from './mongo/db.js';
+import { SyncRuleDocument, SyncRuleState } from './mongo/models.js';
 import { MongoPersistedSyncRulesContent } from './mongo/MongoPersistedSyncRulesContent.js';
 import { MongoSyncBucketStorage } from './mongo/MongoSyncBucketStorage.js';
-import { PowerSyncMongo, PowerSyncMongoOptions } from './mongo/db.js';
-import { SyncRuleDocument, SyncRuleState } from './mongo/models.js';
 import { generateSlotName } from './mongo/util.js';
-import { v4 as uuid } from 'uuid';
-import { logger } from '@powersync/lib-services-framework';
-
-export interface MongoBucketStorageOptions extends PowerSyncMongoOptions {}
 
-export class MongoBucketStorage implements BucketStorageFactory {
+export class MongoBucketStorage
+  extends DisposableObserver<BucketStorageFactoryListener>
+  implements BucketStorageFactory
+{
   private readonly client: mongo.MongoClient;
   private readonly session: mongo.ClientSession;
+  // TODO: This is still Postgres specific and needs to be reworked
   public readonly slot_name_prefix: string;
 
   private readonly storageCache = new LRUCache<number, MongoSyncBucketStorage>({
@@ -47,26 +50,44 @@ export class MongoBucketStorage implements BucketStorageFactory {
         return undefined;
       }
       const rules = new MongoPersistedSyncRulesContent(this.db, doc2);
-      const storage = this.getInstance(rules.parsed());
-      return storage;
+      return this.getInstance(rules);
+    },
+    dispose: (storage) => {
+      storage[Symbol.dispose]();
     }
   });
 
   public readonly db: PowerSyncMongo;
 
-  constructor(db: PowerSyncMongo, options: { slot_name_prefix: string }) {
+  constructor(
+    db: PowerSyncMongo,
+    options: {
+      slot_name_prefix: string;
+    }
+  ) {
+    super();
     this.client = db.client;
     this.db = db;
     this.session = this.client.startSession();
     this.slot_name_prefix = options.slot_name_prefix;
   }
 
-  getInstance(options: PersistedSyncRules): MongoSyncBucketStorage {
-    let { id, sync_rules, slot_name } = options;
+  getInstance(options: PersistedSyncRulesContent): MongoSyncBucketStorage {
+    let { id, slot_name } = options;
     if ((typeof id as any) == 'bigint') {
       id = Number(id);
     }
-    return new MongoSyncBucketStorage(this, id, sync_rules, slot_name);
+    const storage = new MongoSyncBucketStorage(this, id, options, slot_name);
+    this.iterateListeners((cb) => cb.syncStorageCreated?.(storage));
+    storage.registerListener({
+      batchStarted: (batch) => {
+        // This nested listener will be automatically disposed when the storage is disposed
+        batch.registerManagedListener(storage, {
+          replicationEvent: (payload) => this.iterateListeners((cb) => cb.replicationEvent?.(payload))
+        });
+      }
+    });
+    return storage;
   }
 
   async configureSyncRules(sync_rules: string, options?: { lock?: boolean }) {
@@ -136,7 +157,12 @@ export class MongoBucketStorage implements BucketStorageFactory {
 
   async updateSyncRules(options: UpdateSyncRulesOptions): Promise<MongoPersistedSyncRulesContent> {
     // Parse and validate before applying any changes
-    const parsed = SqlSyncRules.fromYaml(options.content);
+    const parsed = SqlSyncRules.fromYaml(options.content, {
+      // No schema-based validation at this point
+      schema: undefined,
+      defaultSchema: 'not_applicable', // Not needed for validation
+      throwOnError: true
+    });
 
     let rules: MongoPersistedSyncRulesContent | undefined = undefined;
 
@@ -204,9 +230,9 @@ export class MongoBucketStorage implements BucketStorageFactory {
     return new MongoPersistedSyncRulesContent(this.db, doc);
   }
 
-  async getActiveSyncRules(): Promise<PersistedSyncRules | null> {
+  async getActiveSyncRules(options: ParseSyncRulesOptions): Promise<PersistedSyncRules | null> {
     const content = await this.getActiveSyncRulesContent();
-    return content?.parsed() ?? null;
+    return content?.parsed(options) ?? null;
   }
 
   async getNextSyncRulesContent(): Promise<MongoPersistedSyncRulesContent | null> {
@@ -223,9 +249,9 @@ export class MongoBucketStorage implements BucketStorageFactory {
     return new MongoPersistedSyncRulesContent(this.db, doc);
   }
 
-  async getNextSyncRules(): Promise<PersistedSyncRules | null> {
+  async getNextSyncRules(options: ParseSyncRulesOptions): Promise<PersistedSyncRules | null> {
     const content = await this.getNextSyncRulesContent();
-    return content?.parsed() ?? null;
+    return content?.parsed(options) ?? null;
   }
 
   async getReplicatingSyncRules(): Promise<PersistedSyncRulesContent[]> {
@@ -252,32 +278,6 @@ export class MongoBucketStorage implements BucketStorageFactory {
     });
   }
 
-  async createWriteCheckpoint(user_id: string, lsns: Record<string, string>): Promise<bigint> {
-    const doc = await this.db.write_checkpoints.findOneAndUpdate(
-      {
-        user_id: user_id
-      },
-      {
-        $set: {
-          lsns: lsns
-        },
-        $inc: {
-          client_id: 1n
-        }
-      },
-      { upsert: true, returnDocument: 'after' }
-    );
-    return doc!.client_id;
-  }
-
-  async lastWriteCheckpoint(user_id: string, lsn: string): Promise<bigint | null> {
-    const lastWriteCheckpoint = await this.db.write_checkpoints.findOne({
-      user_id: user_id,
-      'lsns.1': { $lte: lsn }
-    });
-    return lastWriteCheckpoint?.client_id ?? null;
-  }
-
   async getActiveCheckpoint(): Promise<ActiveCheckpoint> {
     const doc = await this.db.sync_rules.findOne(
       {
@@ -303,7 +303,7 @@ export class MongoBucketStorage implements BucketStorageFactory {
       }
     };
 
-    const active_sync_rules = await this.getActiveSyncRules();
+    const active_sync_rules = await this.getActiveSyncRules({ defaultSchema: 'public' });
     if (active_sync_rules == null) {
       return {
         operations_size_bytes: 0,
@@ -379,7 +379,7 @@ export class MongoBucketStorage implements BucketStorageFactory {
   private makeActiveCheckpoint(doc: SyncRuleDocument | null) {
     return {
       checkpoint: util.timestampToOpId(doc?.last_checkpoint ?? 0n),
-      lsn: doc?.last_checkpoint_lsn ?? replication.ZERO_LSN,
+      lsn: doc?.last_checkpoint_lsn ?? null,
       hasSyncRules() {
         return doc != null;
       },
@@ -389,7 +389,7 @@ export class MongoBucketStorage implements BucketStorageFactory {
         }
         return (await this.storageCache.fetch(doc._id)) ?? null;
       }
-    };
+    } satisfies ActiveCheckpoint;
   }
 
   /**
@@ -479,6 +479,7 @@ export class MongoBucketStorage implements BucketStorageFactory {
       if (doc == null) {
         continue;
       }
+
       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) {
@@ -508,8 +509,19 @@ export class MongoBucketStorage implements BucketStorageFactory {
       // What is important is:
       // 1. checkpoint (op_id) changes.
       // 2. write checkpoint changes for the specific user
+      const bucketStorage = await cp.getBucketStorage();
+      if (!bucketStorage) {
+        continue;
+      }
+
+      const lsnFilters: Record<string, string> = lsn ? { 1: lsn } : {};
 
-      const currentWriteCheckpoint = await this.lastWriteCheckpoint(user_id, lsn ?? '');
+      const currentWriteCheckpoint = await bucketStorage.lastWriteCheckpoint({
+        user_id,
+        heads: {
+          ...lsnFilters
+        }
+      });
 
       if (currentWriteCheckpoint == lastWriteCheckpoint && checkpoint == lastCheckpoint) {
         // No change - wait for next one