@powersync/service-module-mongodb-storage 0.22.0 → 0.23.0

package/src/storage/implementation/MongoBucketBatch.ts

@@ -14,6 +14,7 @@ import {
   ErrorCode,
   errors,
   Logger,
+  ReplicationAbortedError,
   ReplicationAssertionError,
   ServiceError
 } from '@powersync/lib-services-framework';
@@ -33,22 +34,20 @@ import { mongoTableId } from '../../utils/util.js';
 import { PersistedBatch } from './common/PersistedBatch.js';
 import { LoadedSourceRecord, SourceRecordStore } from './common/SourceRecordStore.js';
 import type { VersionedPowerSyncMongo } from './db.js';
+import { SyncRuleDocumentBase } from './models.js';
 import { MAX_ROW_SIZE } from './MongoBucketBatchShared.js';
-import { MongoIdSequence } from './MongoIdSequence.js';
+import { MongoIdSequence, OpIdRangeExhausted } from './MongoIdSequence.js';
+import { MongoOpIdAllocator } from './MongoOpIdAllocator.js';
 import { MongoParsedSyncConfigSet } from './MongoParsedSyncConfigSet.js';
+import { MongoSyncRulesLock } from './MongoSyncRulesLock.js';
 import { MongoWriteBatch } from './MongoWriteBatch.js';
 import { OperationBatch, RecordOperation } from './OperationBatch.js';
 import { ObjectStorage } from './v3/object-storage/ObjectStorage.js';
 import { createObjectStorageUsageWriterId } from './v3/object-storage/ObjectStorageUsage.js';
 
-// Currently, we can only have a single flush() at a time, since it locks the op_id sequence.
-// While the MongoDB transaction retry mechanism handles this okay, using an in-process Mutex
-// makes it more fair and has less overhead.
-//
-// In the future, we can investigate allowing multiple replication streams operating independently.
-const replicationMutex = new utils.Mutex();
-
 export interface MongoBucketBatchOptions {
+  replicationLock: MongoSyncRulesLock;
+  opIdAllocator: MongoOpIdAllocator;
   db: VersionedPowerSyncMongo;
   /**
    * The parsed sync config set for this batch.
@@ -145,7 +144,13 @@ export abstract class MongoBucketBatch
   constructor(options: MongoBucketBatchOptions) {
     super();
     this.logger = options.logger;
-    this.options = options;
+    this.options = {
+      ...options,
+      signal:
+        options.signal == null
+          ? options.replicationLock.signal
+          : AbortSignal.any([options.signal, options.replicationLock.signal])
+    };
     this.client = options.db.client;
     this.db = options.db;
     this.replicationStreamId = options.replicationStreamId;
@@ -215,7 +220,34 @@ export abstract class MongoBucketBatch
     return result;
   }
 
-  private async flushInner(options?: storage.BatchBucketFlushOptions): Promise<storage.FlushedResult | null> {
+  /** Publish the checkpoint in the final flush transaction, or use the empty-commit path. */
+  protected async flushAndCommit<T>(
+    checkpoint: (stream: SyncRuleDocumentBase, lastOp: InternalOpId) => Promise<T>,
+    commitWithoutFlush: () => Promise<T>,
+    options?: storage.BatchBucketFlushOptions,
+    projection?: mongo.Document
+  ): Promise<T> {
+    if (this.batch == null && this.write_checkpoint_batch.length === 0) {
+      return commitWithoutFlush();
+    }
+    let result!: T;
+    while (this.batch != null || this.write_checkpoint_batch.length > 0) {
+      await this.flushInner(
+        options,
+        async (stream, lastOp) => {
+          result = await checkpoint(stream, lastOp);
+        },
+        projection
+      );
+    }
+    return result;
+  }
+
+  private async flushInner(
+    options?: storage.BatchBucketFlushOptions,
+    checkpoint?: (stream: SyncRuleDocumentBase, lastOp: InternalOpId) => Promise<void>,
+    projection?: mongo.Document
+  ): Promise<storage.FlushedResult | null> {
     const batch = this.batch;
     let last_op: InternalOpId | null = null;
     let resumeBatch: OperationBatch | null = null;
@@ -229,21 +261,33 @@ export abstract class MongoBucketBatch
       await this.prepareCustomWriteCheckpoints();
     }
 
-    await this.withReplicationTransaction(`Flushing ${batch?.length ?? 0} ops`, async (session, opSeq) => {
-      clearedError = false;
-      if (batch != null) {
-        const result = await this.replicateBatch(session, batch, opSeq, options);
-        resumeBatch = result.resumeBatch;
-        clearedError ||= result.clearedError;
-      }
+    await this.withReplicationTransaction(
+      `Flushing ${batch?.length ?? 0} ops`,
+      async (session, opSeq) => {
+        clearedError = false;
+        if (batch != null) {
+          const result = await this.replicateBatch(session, batch, opSeq, options);
+          resumeBatch = result.resumeBatch;
+          clearedError ||= result.clearedError;
+        }
 
-      if (this.write_checkpoint_batch.length > 0) {
-        this.logger.info(`Writing ${this.write_checkpoint_batch.length} custom write checkpoints`);
-        await this.batchCreateCustomWriteCheckpoints(session, opSeq.next());
-      }
+        if (this.write_checkpoint_batch.length > 0) {
+          this.logger.info(`Writing ${this.write_checkpoint_batch.length} custom write checkpoints`);
+          await this.batchCreateCustomWriteCheckpoints(session, opSeq.next());
+        }
 
-      last_op = opSeq.last();
-    });
+        last_op = opSeq.last();
+      },
+      async (stream, lastOp) => {
+        if (checkpoint != null && resumeBatch == null) {
+          // The checkpoint must also persist the head, including when visibility is blocked.
+          await checkpoint(stream, lastOp);
+          return true;
+        }
+        return false;
+      },
+      projection
+    );
 
     // Keep checkpoints available if the transaction retries after writing them.
     this.write_checkpoint_batch = [];
@@ -643,17 +687,46 @@ export abstract class MongoBucketBatch
     return result;
   }
 
-  protected async withTransaction(cb: () => Promise<void>) {
-    using lockSpan = this.tracer.span('storage', 'internal_lock');
-    await replicationMutex.exclusiveLock(async () => {
-      lockSpan.end();
-      await this.session.withTransaction(
+  protected async fence(session = this.session, projection?: mongo.Document) {
+    // Graceful cancellation belongs to the source connector's page/batch boundary.
+    // It must still be able to persist progress for rows already flushed. Lease
+    // loss is different: the fence rejects every subsequent writer transaction.
+    return MongoSyncRulesLock.fence(
+      this.db,
+      this.replicationStreamId,
+      this.options.replicationLock,
+      session,
+      projection
+    );
+  }
+
+  protected async withTransaction<T>(
+    cb: (stream: SyncRuleDocumentBase) => Promise<T>,
+    session = this.session
+  ): Promise<T> {
+    return this.withFencedTransaction(() => this.fence(session), cb, session);
+  }
+
+  /** The first operation must modify the stream document conditional on lease ownership. */
+  protected async withFencedTransaction<S, T>(
+    acquireFence: () => Promise<S>,
+    cb: (stream: S) => Promise<T>,
+    session = this.session
+  ): Promise<T> {
+    return this.withWriterLock(() =>
+      session.withTransaction(
         async () => {
+          const stream = await acquireFence();
           try {
-            await cb();
+            const result = await cb(stream);
+            this.options.replicationLock.throwIfAborted();
+            return result;
           } catch (e: unknown) {
+            if (e instanceof OpIdRangeExhausted) {
+              throw e;
+            }
             if (e instanceof mongo.MongoError && e.hasErrorLabel('TransientTransactionError')) {
-              // Likely write conflict caused by concurrent write stream replicating
+              // Likely write conflict caused by concurrent writes to this replication stream.
             } else {
               this.logger.warn('Transaction error', e as Error);
             }
@@ -663,102 +736,118 @@ export abstract class MongoBucketBatch
             throw e;
           }
         },
-        { maxCommitTimeMS: 10000 }
-      );
+        { maxCommitTimeMS: 10000, writeConcern: { w: 'majority' } }
+      )
+    );
+  }
+
+  private async withWriterLock<T>(callback: () => Promise<T>): Promise<T> {
+    using lockSpan = this.tracer.span('storage', 'internal_lock');
+    return this.options.replicationLock.writerMutex.exclusiveLock(async () => {
+      lockSpan.end();
+      this.options.replicationLock.throwIfAborted();
+      return callback();
     });
   }
 
-  private async withReplicationTransaction(
-    description: string,
-    callback: (session: mongo.ClientSession, opSeq: MongoIdSequence) => Promise<void>
+  /** Single-document updates enforce ownership atomically, without a separate transaction. */
+  protected async updateStreamMetadata(
+    set: mongo.Document,
+    filter: mongo.Document = {},
+    writeConcern: mongo.WriteConcernSettings = { w: 'majority' }
   ): Promise<void> {
-    let flushTry = 0;
-
-    const start = Date.now();
-    const lastTry = start + 90000;
-
-    const session = this.session;
-
-    await this.withTransaction(async () => {
-      flushTry += 1;
-      if (flushTry % 10 == 0) {
-        this.logger.info(`${description} - try ${flushTry}`);
-      }
-      if (flushTry > 20 && Date.now() > lastTry) {
-        throw new ServiceError(ErrorCode.PSYNC_S1402, 'Max transaction tries exceeded');
-      }
-
-      const next_op_id_doc = await this.db.op_id_sequence.findOneAndUpdate(
-        {
-          _id: 'main'
-        },
-        {
-          $setOnInsert: { op_id: 0n },
-          $set: {
-            // Force update to ensure we get a mongo lock
-            ts: Date.now()
-          }
-        },
-        {
-          upsert: true,
-          returnDocument: 'after',
-          session
-        }
-      );
-      const opSeq = new MongoIdSequence(next_op_id_doc?.op_id ?? 0n);
+    await this.withWriterLock(() => this.writeStreamMetadata(set, filter, writeConcern));
+  }
 
-      await callback(session, opSeq);
+  /** Caller must already be inside a fenced transaction holding the writer mutex. */
+  protected async updateStreamMetadataInTransaction(set: mongo.Document, filter: mongo.Document = {}): Promise<void> {
+    await this.writeStreamMetadata(set, filter);
+  }
 
-      const writes = this.db.createWriteBatch(session, { ordered: false });
-      writes.updateOne(
-        this.db.op_id_sequence,
-        {
-          _id: 'main'
-        },
-        {
-          $set: {
-            op_id: opSeq.last()
-          }
-        }
-      );
+  private async writeStreamMetadata(
+    set: mongo.Document,
+    filter: mongo.Document,
+    writeConcern?: mongo.WriteConcernSettings
+  ): Promise<void> {
+    const result = await this.db.sync_rules.updateOne(
+      { ...filter, ...MongoSyncRulesLock.ownerFilter(this.replicationStreamId, this.options.replicationLock) },
+      [{ $set: { ...set, ...MongoSyncRulesLock.heartbeatUpdate() } }],
+      { session: this.session, ...(writeConcern == null ? {} : { writeConcern }) }
+    );
+    if (result.matchedCount === 0) {
+      throw new ReplicationAbortedError('Replication writer no longer owns the stream');
+    }
+  }
 
-      writes.updateOne(
-        this.db.sync_rules,
-        {
-          _id: this.replicationStreamId
-        },
-        {
-          $set: {
-            last_keepalive_ts: new Date()
+  private async withReplicationTransaction(
+    description: string,
+    callback: (session: mongo.ClientSession, opSeq: MongoIdSequence) => Promise<void>,
+    finish?: (stream: SyncRuleDocumentBase, lastOp: InternalOpId) => Promise<boolean>,
+    projection?: mongo.Document
+  ): Promise<void> {
+    let flushTry = 0;
+    const lastTry = Date.now() + 90000;
+    const allocator = this.options.opIdAllocator;
+    let lastOp = 0n;
+    // Refill outside the publication transaction, before evaluating any rows.
+    // Exhaustion remains a fallback for large batches or a newer stream head.
+    this.options.replicationLock.throwIfAborted();
+    await allocator.ensureCapacity();
+    for (;;) {
+      try {
+        await this.withFencedTransaction(
+          () => this.fence(this.session, projection),
+          async (stream) => {
+            flushTry += 1;
+            if (flushTry % 10 == 0) {
+              this.logger.info(`${description} - try ${flushTry}`);
+            }
+            if (flushTry > 20 && Date.now() > lastTry) {
+              throw new ServiceError(ErrorCode.PSYNC_S1402, 'Max transaction tries exceeded');
+            }
+            // The fence already holds this stream document for the transaction.
+            // A different writer may have used higher IDs since our last flush.
+            const opSeq = allocator.sequence(this.persistedOpHead(stream));
+            await callback(this.session, opSeq);
+            lastOp = opSeq.last();
+
+            if (!(await finish?.(stream, lastOp))) {
+              const writes = this.db.createWriteBatch(this.session, { ordered: false });
+              // Allow subclasses to persist additional flush-time state in the same transaction.
+              this.onReplicationTransactionFlush(writes, lastOp);
+              await writes.execute();
+            }
+            // Notifications and cleanup must wait until the transaction commits.
           }
+        );
+        allocator.committed(lastOp);
+        return;
+      } catch (error) {
+        if (!(error instanceof OpIdRangeExhausted)) {
+          throw error;
         }
-      );
-
-      // Allow subclasses to persist additional flush-time state in the same transaction
-      // (e.g. v3 advances the stream-level last_persisted_op).
-      this.onReplicationTransactionFlush(writes, opSeq.last());
-      await writes.execute();
-      // We don't notify checkpoint here - we don't make any checkpoint updates directly
-    });
+        // withTransaction has aborted every write. Reserve outside that transaction,
+        // then replay evaluation from the same committed stream head. Existing
+        // ranges can be reused on retry because no operation from this attempt committed.
+        this.options.replicationLock.throwIfAborted();
+        await allocator.reserve();
+      }
+    }
   }
 
-  /**
-   * Hook called inside the replication flush transaction, after ops have been persisted.
-   *
-   * The base implementation does nothing; v3 storage overrides this to `$max` the stream-level
-   * `last_persisted_op` durably within the same transaction.
-   */
-  protected onReplicationTransactionFlush(_writes: MongoWriteBatch, _lastOp: InternalOpId): void {
-    // No-op by default (v1 behaviour unchanged).
-  }
+  /** Highest durably persisted operation, including operations not yet checkpointed. */
+  protected abstract persistedOpHead(stream: SyncRuleDocumentBase): InternalOpId;
+
+  /** Advance the persisted head in the same transaction as the operations. */
+  protected abstract onReplicationTransactionFlush(writes: MongoWriteBatch, lastOp: InternalOpId): void;
 
   /**
    * Called after a replication transaction has successfully committed, with the last persisted op id.
    *
    * v1 storage tracks this in memory to fold into the next checkpoint. v3 storage does not need it:
    * the stream-level `last_persisted_op` is already advanced durably by
-   * {@link onReplicationTransactionFlush} within the same transaction, and checkpoints read it
-   * from the document.
+   * {@link onReplicationTransactionFlush} or the combined checkpoint update within the same
+   * transaction, and empty commits read it from the document.
    *
    * Keep calls to this adjacent to {@link withReplicationTransaction} usage - both must observe
    * every path that persists ops.
@@ -837,7 +926,9 @@ export abstract class MongoBucketBatch
 
     await this.withTransaction(async () => {
       for (let table of sourceTables) {
-        await this.db.commonSourceTables(this.replicationStreamId).deleteOne({ _id: mongoTableId(table.id) });
+        await this.db
+          .commonSourceTables(this.replicationStreamId)
+          .deleteOne({ _id: mongoTableId(table.id) }, { session: this.session });
       }
     });
 

package/src/storage/implementation/MongoCompactor.ts

@@ -16,6 +16,8 @@ export interface MongoCompactOptions extends storage.CompactOptions {
    * the lightweight pass.
    */
   compactChunksOnly?: boolean;
+  /** Internal/testing use: bypass the new-chunk threshold during chunk-only compaction. */
+  forceChunkCompaction?: boolean;
 }
 
 const DEFAULT_CLEAR_BATCH_LIMIT = 5000;

package/src/storage/implementation/MongoIdSequence.ts

@@ -9,7 +9,10 @@ import { ReplicationAssertionError } from '@powersync/lib-services-framework';
 export class MongoIdSequence {
   private _last: bigint;
 
-  constructor(last: bigint) {
+  constructor(
+    last: bigint,
+    private readonly ranges?: readonly OpIdRange[]
+  ) {
     if (typeof last != 'bigint') {
       throw new ReplicationAssertionError(`BigInt required, got ${last} ${typeof last}`);
     }
@@ -17,6 +20,14 @@ export class MongoIdSequence {
   }
 
   next() {
+    if (this.ranges != null) {
+      const next = this.ranges.find((range) => range.end > this._last);
+      if (next == null) {
+        throw new OpIdRangeExhausted();
+      }
+      this._last = this._last + 1n < next.start ? next.start : this._last + 1n;
+      return this._last;
+    }
     return ++this._last;
   }
 
@@ -24,3 +35,12 @@ export class MongoIdSequence {
     return this._last;
   }
 }
+
+/** Inclusive, durably reserved operation IDs. Gaps between ranges are valid. */
+export interface OpIdRange {
+  start: bigint;
+  end: bigint;
+}
+
+/** Abort the transaction before reserving more IDs and retrying its evaluation. */
+export class OpIdRangeExhausted extends Error {}

package/src/storage/implementation/MongoOpIdAllocator.ts

@@ -0,0 +1,105 @@
+import { mongo } from '@powersync/lib-service-mongodb';
+import { ReplicationAssertionError } from '@powersync/lib-services-framework';
+import { VersionedPowerSyncMongo } from './db.js';
+import { MongoIdSequence, OpIdRange } from './MongoIdSequence.js';
+
+const RESERVATION_SIZE = 65_536n;
+const LOW_WATER_MARK = 16_384n;
+const MAX_OP_ID = (1n << 63n) - 1n;
+
+/**
+ * Shared by writers of one stream and lease owner in this process. Reservations
+ * use short majority writes, never the publication session. A crashed owner
+ * abandons its unused ranges; the global sequence is an allocation watermark.
+ */
+export class MongoOpIdAllocator {
+  private ranges: OpIdRange[] = [];
+  private reserving?: Promise<void>;
+
+  constructor(private readonly db: VersionedPowerSyncMongo) {}
+
+  discard() {
+    this.ranges = [];
+  }
+
+  sequence(persistedOp: bigint): MongoIdSequence {
+    // The stream is fenced before reading this head. Another writer may have
+    // committed since this allocator was last used, including in another process.
+    // Never reuse reserved IDs below that durable head.
+    return new MongoIdSequence(persistedOp, this.ranges.slice());
+  }
+
+  committed(lastOp: bigint) {
+    this.ranges = this.ranges
+      .filter((range) => range.end > lastOp)
+      .map((range) => ({ start: range.start > lastOp ? range.start : lastOp + 1n, end: range.end }));
+  }
+
+  /** Top up between transactions, retaining every unused ID in existing ranges. */
+  async ensureCapacity(): Promise<void> {
+    const remaining = this.ranges.reduce((total, range) => total + range.end - range.start + 1n, 0n);
+    if (remaining < LOW_WATER_MARK) {
+      await this.reserve();
+    }
+  }
+
+  async reserve(): Promise<void> {
+    if (this.reserving != null) {
+      return this.reserving;
+    }
+    this.reserving = this.reserveRange();
+    try {
+      await this.reserving;
+    } finally {
+      this.reserving = undefined;
+    }
+  }
+
+  /**
+   * Atomically reserve a range of operation IDs in the global sequence.
+   *
+   * We continue using the range across multiple batches, until the range is exhausted,
+   * so we don't "waste" large segements in long-running processes with small batches.
+   * However, we can't return unused portions back to the database if we don't use them.
+   */
+  private async reserveRange(): Promise<void> {
+    try {
+      // Return the previous watermark so an unchanged value at the limit can be
+      // distinguished from a successful reservation ending near the limit.
+      const previous = await this.db.op_id_sequence.findOneAndUpdate(
+        { _id: 'main' },
+        [
+          {
+            $set: {
+              op_id: {
+                $let: {
+                  vars: { current: { $ifNull: ['$op_id', 0n] } },
+                  in: {
+                    $cond: [
+                      { $lte: ['$$current', MAX_OP_ID - RESERVATION_SIZE] },
+                      { $add: ['$$current', RESERVATION_SIZE] },
+                      '$$current'
+                    ]
+                  }
+                }
+              }
+            }
+          }
+        ],
+        { upsert: true, returnDocument: 'before', writeConcern: { w: 'majority' } }
+      );
+      const last = previous?.op_id ?? 0n;
+      if (last > MAX_OP_ID - RESERVATION_SIZE) {
+        // Not expected to ever happen, unless the sequence was manually modified
+        throw new ReplicationAssertionError('Operation ID sequence exhausted');
+      }
+      this.ranges.push({ start: last + 1n, end: last + RESERVATION_SIZE });
+    } catch (error) {
+      // Concurrent first reservations may race to insert the sequence document - retry in that case.
+      if (error instanceof mongo.MongoServerError && error.code === 11000) {
+        return this.reserveRange();
+      }
+      throw error;
+    }
+  }
+}

package/src/storage/implementation/MongoSyncBucketStorage.ts

@@ -43,6 +43,7 @@ import { MongoCompactOptions, MongoCompactor } from './MongoCompactor.js';
 import { MongoParameterCompactor } from './MongoParameterCompactor.js';
 import { MongoParsedSyncConfigSet } from './MongoParsedSyncConfigSet.js';
 import { MongoPersistedReplicationStream } from './MongoPersistedReplicationStream.js';
+import { MongoSyncRulesLock } from './MongoSyncRulesLock.js';
 import { MongoCheckpointAPIOptions, MongoWriteCheckpointAPI } from './MongoWriteCheckpointAPI.js';
 import { ObjectStorage } from './v3/object-storage/ObjectStorage.js';
 
@@ -109,6 +110,10 @@ export abstract class MongoSyncBucketStorage
   public readonly readPreference: mongo.ReadPreference | undefined;
   public readonly clearBatchThrottleRate: number;
   #storageInitialized = false;
+  /**
+   * Required for writing, optional for reading.
+   */
+  private readonly replicationLock: MongoSyncRulesLock | null;
 
   constructor(
     public readonly factory: MongoBucketStorage,
@@ -119,6 +124,7 @@ export abstract class MongoSyncBucketStorage
     options: MongoSyncBucketStorageOptions
   ) {
     super();
+    this.replicationLock = replicationStream.current_lock;
     this.storageConfig = options.storageConfig;
     this.objectStorage = options.objectStorage;
     // Keep small chunks inline in MongoDB rather than offloading them to S3.
@@ -269,12 +275,15 @@ export abstract class MongoSyncBucketStorage
    * The version-independent part of the batch options.
    */
   protected writerBatchOptions(options: storage.CreateWriterOptions): Omit<MongoBucketBatchOptions, 'resumeFromLsn'> {
+    const replicationLock = this.requireReplicationLock();
     return {
       logger: options.logger ?? this.logger,
       db: this.db,
       parsedSyncConfig: this.getParsedSyncConfigSet(options),
       replicationStreamId: this.replicationStreamId,
       replicationStreamName: this.replicationStreamName,
+      replicationLock,
+      opIdAllocator: this.factory.getOpIdAllocator(this.replicationStream, replicationLock),
       storeCurrentData: options.storeCurrentData,
       skipExistingRows: options.skipExistingRows ?? false,
       markRecordUnavailable: options.markRecordUnavailable,
@@ -286,7 +295,16 @@ export abstract class MongoSyncBucketStorage
     };
   }
 
+  private requireReplicationLock(): MongoSyncRulesLock {
+    if (this.replicationLock == null) {
+      throw new ServiceAssertionError('A replication lease is required to create a writer');
+    }
+    this.replicationLock.throwIfAborted();
+    return this.replicationLock;
+  }
+
   async createWriter(options: storage.CreateWriterOptions): Promise<storage.BucketStorageBatch> {
+    this.requireReplicationLock();
     await this.initializeStorage();
 
     const writer = await this.createWriterImpl(options);
@@ -385,6 +403,7 @@ export abstract class MongoSyncBucketStorage
       throw new ReplicationAbortedError('Aborted clearing data', signal.reason);
     }
 
+    this.factory.discardOpIdAllocator(this.replicationStreamId);
     await this.clearSyncRuleState();
 
     await this.clearBucketData(signal);