@powersync/service-module-mongodb-storage 0.18.2 → 0.19.0

package/src/storage/implementation/v1/MongoCompactorV1.ts

@@ -1,20 +1,25 @@
 import { mongo } from '@powersync/lib-service-mongodb';
 import { ReplicationAssertionError } from '@powersync/lib-services-framework';
-import { storage } from '@powersync/service-core';
+import { addChecksums, storage, utils } from '@powersync/service-core';
 import { BucketDefinitionId } from '@powersync/service-sync-rules';
-import { SingleBucketStore } from '../common/SingleBucketStore.js';
+import { BucketDataDoc } from '../common/BucketDataDoc.js';
 import { BucketStateDocumentBase, LEGACY_BUCKET_DATA_DEFINITION_ID } from '../models.js';
-import { DirtyBucket, MongoCompactor } from '../MongoCompactor.js';
-import { BucketStateDocumentV1 } from './models.js';
+import { CurrentBucketState, DirtyBucket, MongoCompactor } from '../MongoCompactor.js';
+import { cacheKey } from '../OperationBatch.js';
+import { BucketDataDocumentV1, BucketStateDocumentV1 } from './models.js';
 import type { MongoSyncBucketStorageV1 } from './MongoSyncBucketStorageV1.js';
 import { SingleBucketStoreV1 } from './SingleBucketStoreV1.js';
 import { VersionedPowerSyncMongoV1 } from './VersionedPowerSyncMongoV1.js';
 
+type CompactClearProperties = 'op' | 'checksum' | 'target_op';
+
 export class MongoCompactorV1 extends MongoCompactor {
   // Override types to the more specific ones
   declare protected readonly db: VersionedPowerSyncMongoV1;
   declare protected readonly storage: MongoSyncBucketStorageV1;
 
+  private updates: mongo.AnyBulkWriteOperation<BucketDataDocumentV1>[] = [];
+
   public async *dirtyBucketBatches(options: {
     minBucketChanges: number;
     minChangeRatio: number;
@@ -81,14 +86,293 @@ export class MongoCompactorV1 extends MongoCompactor {
     };
   }
 
-  protected async getBucketDataContext(
-    bucket: string,
-    _definitionId: BucketDefinitionId | null
-  ): Promise<SingleBucketStore | null> {
+  private getBucketDataContext(bucket: string): SingleBucketStoreV1 {
     return new SingleBucketStoreV1(this.db, {
       replicationStreamId: this.group_id,
       definitionId: LEGACY_BUCKET_DATA_DEFINITION_ID,
       bucket
     });
   }
+
+  protected async compactSingleBucket(bucket: string) {
+    // Do not carry queued writes from a failed attempt into the rescan
+    this.updates = [];
+
+    const idLimitBytes = this.idLimitBytes;
+    const bucketContext = this.getBucketDataContext(bucket);
+    const currentState: CurrentBucketState = {
+      bucket,
+      definitionId: bucketContext.key.definitionId,
+      seen: new Map(),
+      trackingSize: 0,
+      lastNotPut: null,
+      opsSincePut: 0,
+      checksum: 0,
+      opCount: 0,
+      opBytes: 0
+    };
+
+    // Constant lower bound.
+    const lowerBound = bucketContext.minId;
+    // Upper bound is adjusted for each batch.
+    let upperBound = bucketContext.maxId;
+
+    // Paginate through bucket data in batches to avoid cursor timeouts.
+    while (true) {
+      this.signal?.throwIfAborted();
+
+      // Query one batch at a time, to avoid cursor timeouts.
+      const pipeline = [
+        {
+          $match: {
+            _id: {
+              $gte: lowerBound,
+              $lt: upperBound
+            }
+          }
+        },
+        { $sort: { _id: -1 } },
+        { $limit: this.moveBatchQueryLimit },
+        {
+          $project: {
+            _id: 1,
+            op: 1,
+            table: 1,
+            row_id: 1,
+            source_table: 1,
+            source_key: 1,
+            checksum: 1,
+            size: { $bsonSize: '$$ROOT' }
+          }
+        }
+      ];
+
+      const cursor = bucketContext.collection.aggregate<BucketDataDocumentV1 & { size: number | bigint }>(pipeline, {
+        // batchSize is 1 more than limit to auto-close the cursor.
+        // See https://github.com/mongodb/node-mongodb-native/pull/4580
+        batchSize: this.moveBatchQueryLimit + 1
+      });
+      // We don't limit to a single batch here, since that often causes MongoDB to scan through more than it returns.
+      // Instead, we load up to the limit.
+      const rawBatch = await cursor.toArray();
+      const batch = rawBatch.map((document) => {
+        const { size, ...rest } = document;
+        return {
+          doc: bucketContext.fromPersistedDocument(rest),
+          size
+        };
+      });
+
+      if (batch.length == 0) {
+        // We've reached the end.
+        break;
+      }
+
+      // Reuse the exact collection _id value from Mongo for the next bound.
+      upperBound = rawBatch[rawBatch.length - 1]._id;
+
+      for (const { doc, size } of batch) {
+        if (doc.o > this.maxOpId) {
+          continue;
+        }
+
+        currentState.checksum = addChecksums(currentState.checksum, Number(doc.checksum));
+        currentState.opCount += 1;
+
+        let isPersistentPut = doc.op == 'PUT';
+
+        currentState.opBytes += Number(size);
+        if (doc.op == 'REMOVE' || doc.op == 'PUT') {
+          const key = `${doc.table}/${doc.row_id}/${cacheKey(doc.source_table!, doc.source_key!)}`;
+          const targetOp = currentState.seen.get(key);
+          if (targetOp) {
+            // Will convert to MOVE, so don't count as PUT.
+            isPersistentPut = false;
+
+            this.updates.push({
+              updateOne: {
+                filter: { _id: bucketContext.docId(doc.o) },
+                update: {
+                  $set: {
+                    op: 'MOVE',
+                    target_op: targetOp
+                  },
+                  $unset: {
+                    source_table: 1,
+                    source_key: 1,
+                    table: 1,
+                    row_id: 1,
+                    data: 1
+                  }
+                } satisfies mongo.UpdateFilter<BucketDataDocumentV1>
+              }
+            });
+
+            // TODO: better estimate for this.
+            currentState.opBytes += 200 - Number(size);
+          } else if (currentState.trackingSize < idLimitBytes) {
+            // flatstr reduces the memory usage by flattening the string.
+            currentState.seen.set(utils.flatstr(key), doc.o);
+            // length + 16 for the string
+            // 24 for the bigint
+            // 50 for map overhead
+            // 50 for additional overhead
+            currentState.trackingSize += key.length + 140;
+          }
+        }
+
+        if (isPersistentPut) {
+          currentState.lastNotPut = null;
+          currentState.opsSincePut = 0;
+        } else if (doc.op != 'CLEAR') {
+          if (currentState.lastNotPut == null) {
+            currentState.lastNotPut = doc.o;
+          }
+          currentState.opsSincePut += 1;
+        }
+
+        if (this.updates.length + this.bucketStateUpdates.length >= this.moveBatchLimit) {
+          await this.flush(bucketContext);
+        }
+      }
+
+      this.logger.info(`Processed batch of length ${batch.length} current bucket: ${bucket}`);
+    }
+
+    // Free memory before clearing the bucket.
+    currentState.seen.clear();
+    if (currentState.lastNotPut != null && currentState.opsSincePut >= 1) {
+      this.logger.info(
+        `Inserting CLEAR at ${this.group_id}:${bucket}:${currentState.lastNotPut} to remove ${currentState.opsSincePut} operations`
+      );
+      // Need flush() before clear().
+      await this.flush(bucketContext);
+      await this.clearBucket(currentState, bucketContext);
+    }
+
+    // Do this after clearBucket so we have accurate counts.
+    this.updateBucketChecksums(currentState, this.maxOpId);
+    // Need another flush after updateBucketChecksums().
+    await this.flush(bucketContext);
+  }
+
+  private async flush(bucketContext: SingleBucketStoreV1) {
+    if (this.updates.length > 0) {
+      this.logger.info(`Compacting ${this.updates.length} ops`);
+      await bucketContext.collection.bulkWrite(this.updates, {
+        // Order is not important. Since checksums are not affected, these operations can happen in any order,
+        // and it's fine if the operations are partially applied. Each individual operation is atomic.
+        ordered: false
+      });
+      this.updates = [];
+    }
+
+    await this.flushBucketStateUpdates();
+  }
+
+  /**
+   * Perform a CLEAR compact for a bucket.
+   *
+   * @param currentState tracks the last non-PUT op, which will be converted to CLEAR.
+   */
+  private async clearBucket(currentState: CurrentBucketState, bucketContext: SingleBucketStoreV1) {
+    const clearOp = currentState.lastNotPut!;
+
+    const opFilter = {
+      _id: {
+        $gte: bucketContext.minId,
+        $lte: bucketContext.docId(clearOp)
+      }
+    };
+
+    const session = this.db.client.startSession();
+    try {
+      let done = false;
+      while (!done) {
+        this.signal?.throwIfAborted();
+        let opCountDiff = 0;
+        // Do the CLEAR operation in batches, with each batch a separate transaction.
+        // The state after each batch is fully consistent.
+        // We need a transaction per batch to make sure checksums stay consistent.
+        await session.withTransaction(
+          async () => {
+            const query = bucketContext.collection.find<Pick<BucketDataDocumentV1, '_id' | CompactClearProperties>>(
+              opFilter,
+              {
+                session,
+                sort: { _id: 1 },
+                projection: {
+                  _id: 1,
+                  op: 1,
+                  checksum: 1,
+                  target_op: 1
+                },
+                limit: this.clearBatchLimit
+              }
+            );
+            let checksum = 0;
+            let lastOp: Pick<BucketDataDoc, 'o' | CompactClearProperties> | null = null;
+            let targetOp: bigint | null = null;
+            let gotAnOp = false;
+            let numberOfOpsToClear = 0;
+            for await (const rawOp of query.stream()) {
+              const op = bucketContext.fromPartialPersistedDocument(rawOp);
+
+              if (op.op == 'MOVE' || op.op == 'REMOVE' || op.op == 'CLEAR') {
+                checksum = utils.addChecksums(checksum, Number(op.checksum));
+                lastOp = op;
+                numberOfOpsToClear += 1;
+                if (op.op != 'CLEAR') {
+                  gotAnOp = true;
+                }
+                if (op.target_op != null && (targetOp == null || op.target_op > targetOp)) {
+                  targetOp = op.target_op;
+                }
+              } else {
+                throw new ReplicationAssertionError(`Unexpected ${op.op} operation at ${this.formatBucketDataKey(op)}`);
+              }
+            }
+            if (!gotAnOp) {
+              done = true;
+              return;
+            }
+
+            this.logger.info(`Flushing CLEAR for ${numberOfOpsToClear} ops at ${lastOp?.o}`);
+            await bucketContext.collection.deleteMany(
+              {
+                _id: {
+                  $gte: bucketContext.minId,
+                  $lte: bucketContext.docId(lastOp!.o)
+                }
+              },
+              { session }
+            );
+
+            const op = bucketContext.toPersistedDocument({
+              o: lastOp!.o,
+              op: 'CLEAR',
+              checksum: BigInt(checksum),
+              data: null,
+              target_op: targetOp
+            });
+            await bucketContext.collection.insertOne(op, { session });
+
+            opCountDiff = -numberOfOpsToClear + 1;
+          },
+          {
+            writeConcern: { w: 'majority' },
+            readConcern: { level: 'snapshot' }
+          }
+        );
+        // Update outside the transaction, since the transaction can be retried multiple times.
+        currentState.opCount += opCountDiff;
+      }
+    } finally {
+      await session.endSession();
+    }
+  }
+
+  private formatBucketDataKey(doc: Pick<BucketDataDoc, 'bucketKey' | 'o'>) {
+    return `${doc.bucketKey.replicationStreamId}:${doc.bucketKey.bucket}:${doc.o}`;
+  }
 }

package/src/storage/implementation/v1/MongoSyncBucketStorageV1.ts

@@ -16,7 +16,15 @@ import {
 import { JSONBig } from '@powersync/service-jsonbig';
 import { ParameterLookupRows, ScopedParameterLookup, SqliteJsonRow } from '@powersync/service-sync-rules';
 import * as bson from 'bson';
-import { idPrefixFilter, mapOpEntry, readSingleBatch, setSessionSnapshotTime } from '../../../utils/util.js';
+import {
+  clearCollectionInIdBatches,
+  clearCollectionInIdRanges,
+  idPrefixFilter,
+  mapOpEntry,
+  readSingleBatch,
+  retryOnMongoMaxTimeMSExpired,
+  setSessionSnapshotTime
+} from '../../../utils/util.js';
 import { MongoBucketStorage } from '../../MongoBucketStorage.js';
 import { MongoSyncBucketStorageCheckpoint } from '../common/MongoSyncBucketStorageCheckpoint.js';
 import { SourceKey } from '../models.js';
@@ -214,64 +222,59 @@ export class MongoSyncBucketStorageV1 extends MongoSyncBucketStorage {
   }
 
   protected async clearBucketData(signal?: AbortSignal): Promise<void> {
-    await this.clearDeleteMany(
+    await clearCollectionInIdRanges(
+      this.logger,
       'bucket data',
-      () =>
-        this.db.bucket_data.deleteMany(
-          {
-            _id: idPrefixFilter<BucketDataKeyV1>({ g: this.replicationStreamId }, ['b', 'o'])
-          },
-          { maxTimeMS: lib_mongo.db.MONGO_CLEAR_OPERATION_TIMEOUT_MS }
-        ),
-      signal
+      this.db.bucket_data,
+      {
+        _id: idPrefixFilter<BucketDataKeyV1>({ g: this.replicationStreamId }, ['b', 'o'])
+      },
+      signal,
+      this.clearBatchThrottleRate
     );
   }
 
   protected async clearParameterIndexes(signal?: AbortSignal): Promise<void> {
-    await this.clearDeleteMany(
+    await clearCollectionInIdBatches(
+      this.logger,
       'parameter index',
-      () =>
-        this.db.parameterIndexV1.deleteMany(
-          {
-            'key.g': this.replicationStreamId
-          },
-          { maxTimeMS: lib_mongo.db.MONGO_CLEAR_OPERATION_TIMEOUT_MS }
-        ),
-      signal
+      this.db.parameterIndexV1,
+      {
+        'key.g': this.replicationStreamId
+      },
+      signal,
+      this.clearBatchThrottleRate
     );
   }
 
   protected async clearSourceRecords(signal?: AbortSignal): Promise<void> {
-    await this.clearDeleteMany(
+    await clearCollectionInIdRanges(
+      this.logger,
       'source records',
-      () =>
-        this.db.sourceRecordsV1.deleteMany(
-          {
-            _id: idPrefixFilter<SourceKey>({ g: this.replicationStreamId }, ['t', 'k'])
-          },
-          { maxTimeMS: lib_mongo.db.MONGO_CLEAR_OPERATION_TIMEOUT_MS }
-        ),
-      signal
+      this.db.sourceRecordsV1,
+      {
+        _id: idPrefixFilter<SourceKey>({ g: this.replicationStreamId }, ['t', 'k'])
+      },
+      signal,
+      this.clearBatchThrottleRate
     );
   }
 
   protected async clearBucketState(signal?: AbortSignal): Promise<void> {
-    await this.clearDeleteMany(
+    await clearCollectionInIdRanges(
+      this.logger,
       'bucket state',
-      () =>
-        this.db.bucketStateV1.deleteMany(
-          {
-            _id: idPrefixFilter<BucketStateDocument['_id']>({ g: this.replicationStreamId }, ['b'])
-          },
-          { maxTimeMS: lib_mongo.db.MONGO_CLEAR_OPERATION_TIMEOUT_MS }
-        ),
-      signal
+      this.db.bucketStateV1,
+      {
+        _id: idPrefixFilter<BucketStateDocument['_id']>({ g: this.replicationStreamId }, ['b'])
+      },
+      signal,
+      this.clearBatchThrottleRate
     );
   }
 
   protected async clearSourceTables(signal?: AbortSignal): Promise<void> {
-    await this.clearDeleteMany(
-      'source tables',
+    await retryOnMongoMaxTimeMSExpired(
       () =>
         this.db.sourceTablesV1(this.replicationStreamId).deleteMany(
           {
@@ -279,7 +282,17 @@ export class MongoSyncBucketStorageV1 extends MongoSyncBucketStorage {
           },
           { maxTimeMS: lib_mongo.db.MONGO_CLEAR_OPERATION_TIMEOUT_MS }
         ),
-      signal
+      {
+        signal,
+        abortMessage: 'Aborted clearing data',
+        // This is a fairly long delay - only expected to hit this when the storage database is under high load.
+        retryDelayMs: lib_mongo.db.MONGO_CLEAR_OPERATION_TIMEOUT_MS,
+        onRetry: () => {
+          this.logger.info(
+            `Clearing batch of source tables timed out after ${lib_mongo.db.MONGO_CLEAR_OPERATION_TIMEOUT_MS}ms, retrying...`
+          );
+        }
+      }
     );
   }
 
@@ -389,6 +402,7 @@ export async function* getBucketDataBatchV1(
 
   if (session != null) {
     session.advanceOperationTime(checkpoint.snapshotTime);
+    session.advanceClusterTime(checkpoint.clusterTime);
   }
 
   let filters: mongo.Filter<BucketDataDocumentV1>[] = [];

package/src/storage/implementation/v1/SingleBucketStoreV1.ts

@@ -1,58 +1,50 @@
 import { mongo } from '@powersync/lib-service-mongodb';
 import { InternalOpId } from '@powersync/service-core';
 import { BucketDataDoc, BucketKey } from '../common/BucketDataDoc.js';
-import {
-  BucketDataDocumentGeneric,
-  BucketDataDocumentGenericId,
-  SingleBucketStore
-} from '../common/SingleBucketStore.js';
 import { BucketDataProperties } from '../models.js';
 import { VersionedPowerSyncMongoV1 } from './VersionedPowerSyncMongoV1.js';
 import { BucketDataDocumentV1, BucketDataKeyV1, serializeBucketDataV1 } from './models.js';
 
-export class SingleBucketStoreV1 implements SingleBucketStore {
-  public readonly collection: mongo.Collection<BucketDataDocumentGeneric>;
+export class SingleBucketStoreV1 {
+  public readonly collection: mongo.Collection<BucketDataDocumentV1>;
 
   constructor(
-    private db: VersionedPowerSyncMongoV1,
+    db: VersionedPowerSyncMongoV1,
     public readonly key: BucketKey
   ) {
-    this.collection = db.bucketDataV1 as unknown as mongo.Collection<BucketDataDocumentGeneric>;
+    this.collection = db.bucketDataV1;
   }
 
-  docId(o: InternalOpId): BucketDataDocumentGenericId {
-    // `satisfies BucketDataKeyV1` checks that we use the correct type for V1 storage
-    // `as anyt` is to allow casting to the interface virtual type
+  docId(o: InternalOpId): BucketDataKeyV1 {
     return {
       g: this.key.replicationStreamId,
       b: this.key.bucket,
       o
-    } satisfies BucketDataKeyV1 as any;
+    };
   }
 
-  get minId(): BucketDataDocumentGenericId {
+  get minId(): BucketDataKeyV1 {
     return {
       g: this.key.replicationStreamId,
       b: this.key.bucket,
-      o: new mongo.MinKey()
-    } as any;
+      o: new mongo.MinKey() as any
+    };
   }
 
-  get maxId(): BucketDataDocumentGenericId {
+  get maxId(): BucketDataKeyV1 {
     return {
       g: this.key.replicationStreamId,
       b: this.key.bucket,
-      o: new mongo.MaxKey()
-    } as any;
+      o: new mongo.MaxKey() as any
+    };
   }
 
-  toPersistedDocument(source: Omit<BucketDataDoc, 'bucketKey'>): BucketDataDocumentGeneric {
-    return serializeBucketDataV1({ bucketKey: this.key, ...source }) as unknown as BucketDataDocumentGeneric;
+  toPersistedDocument(source: Omit<BucketDataDoc, 'bucketKey'>): BucketDataDocumentV1 {
+    return serializeBucketDataV1({ bucketKey: this.key, ...source });
   }
 
-  fromPersistedDocument(doc: BucketDataDocumentGeneric): BucketDataDoc {
-    const document = doc as unknown as BucketDataDocumentV1;
-    const { _id, ...rest } = document;
+  fromPersistedDocument(doc: BucketDataDocumentV1): BucketDataDoc {
+    const { _id, ...rest } = doc;
     return {
       bucketKey: this.key,
       o: _id.o,
@@ -61,10 +53,9 @@ export class SingleBucketStoreV1 implements SingleBucketStore {
   }
 
   fromPartialPersistedDocument<T extends keyof BucketDataProperties>(
-    doc: Pick<BucketDataDocumentGeneric, '_id' | T>
+    doc: Pick<BucketDataDocumentV1, '_id' | T>
   ): Pick<BucketDataDoc, 'bucketKey' | 'o' | T> {
-    const document = doc as Pick<BucketDataDocumentV1, '_id' | T>;
-    const { _id, ...rest } = document;
+    const { _id, ...rest } = doc;
     return {
       bucketKey: this.key,
       o: _id.o,

package/src/storage/implementation/v3/BucketDataContextV3.ts

@@ -0,0 +1,32 @@
+import { mongo } from '@powersync/lib-service-mongodb';
+import { InternalOpId } from '@powersync/service-core';
+import { BucketKey } from '../common/BucketDataDoc.js';
+import { BucketDataKey } from '../models.js';
+import { VersionedPowerSyncMongoV3 } from './VersionedPowerSyncMongoV3.js';
+import { BucketDataDocumentV3 } from './models.js';
+
+export class BucketDataContextV3 {
+  public readonly collection: mongo.Collection<BucketDataDocumentV3>;
+
+  constructor(
+    db: VersionedPowerSyncMongoV3,
+    public readonly key: BucketKey
+  ) {
+    this.collection = db.bucketData(key.replicationStreamId, key.definitionId);
+  }
+
+  docId(o: InternalOpId): BucketDataKey {
+    return {
+      b: this.key.bucket,
+      o
+    };
+  }
+
+  get minId(): BucketDataKey {
+    // MongoDB's MinKey sentinel does not match the bigint type for _id.o.
+    return {
+      b: this.key.bucket,
+      o: new mongo.MinKey()
+    } as unknown as BucketDataKey;
+  }
+}

package/src/storage/implementation/v3/MongoBucketBatchV3.ts

@@ -43,7 +43,9 @@ export class MongoBucketBatchV3 extends MongoBucketBatch {
 
   protected createPersistedBatch(writtenSize: number): PersistedBatch {
     return new PersistedBatchV3(this.db, this.replicationStreamId, this.mapping, writtenSize, {
-      logger: this.logger
+      logger: this.logger,
+      objectStorage: this.options.objectStorage,
+      inlineThresholdBytes: this.options.inlineThresholdBytes
     });
   }