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

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

@@ -13,8 +13,10 @@ import {
   UpsertCurrentDataOptions
 } from '../common/PersistedBatch.js';
 import { SourceRecordLookupState } from '../common/SourceRecordStore.js';
+import { MongoWriteBatch } from '../MongoWriteBatch.js';
 import { serializeBucketData } from './bucket-format.js';
 import { chunkBucketData } from './chunking.js';
+import { DEFAULT_MIN_COMPACT_CHUNK_INTERVAL_MS } from './compaction-constants.js';
 import {
   BucketDataDocumentV3,
   BucketStateDocumentV3,
@@ -24,12 +26,21 @@ import {
   taggedBucketParameterDocumentToTagged
 } from './models.js';
 import { ObjectStorageLifecycle } from './object-storage/ObjectStorageLifecycle.js';
+import { ObjectStorageUsage } from './object-storage/ObjectStorageUsage.js';
 import { VersionedPowerSyncMongoV3 } from './VersionedPowerSyncMongoV3.js';
 
+interface SourceRecordWrite {
+  sourceTableId: bson.ObjectId;
+  operation: mongo.AnyBulkWriteOperation<CurrentDataDocumentV3>;
+}
+
 export class PersistedBatchV3 extends PersistedBatch {
-  currentData: { sourceTableId: bson.ObjectId; operation: mongo.AnyBulkWriteOperation<CurrentDataDocumentV3> }[] = [];
+  // Upserts and soft deletes supply the complete source-record state, including
+  // pending_delete. Keep the final state per key, without removing any history.
+  currentData = new Map<string, SourceRecordWrite>();
   sourceTablePendingDeletes = new Map<string, InternalOpId>();
   protected readonly objectStorageLifecycle?: ObjectStorageLifecycle;
+  protected readonly objectStorageUsage?: ObjectStorageUsage;
 
   declare protected readonly db: VersionedPowerSyncMongoV3;
 
@@ -43,6 +54,7 @@ export class PersistedBatchV3 extends PersistedBatch {
     super(db, group_id, mapping, writtenSize, options);
     if (this.objectStorage) {
       this.objectStorageLifecycle = new ObjectStorageLifecycle(this.db, this.group_id, this.objectStorage);
+      this.objectStorageUsage = new ObjectStorageUsage(this.db, this.group_id, this.objectStorageUsageWriterId);
     }
   }
 
@@ -127,7 +139,7 @@ export class PersistedBatchV3 extends PersistedBatch {
   }
 
   hardDeleteCurrentData(sourceTableId: bson.ObjectId, replicaId: storage.ReplicaId) {
-    this.currentData.push({
+    this.currentData.set(this.currentDataKey(sourceTableId, replicaId), {
       sourceTableId,
       operation: {
         deleteOne: {
@@ -143,7 +155,7 @@ export class PersistedBatchV3 extends PersistedBatch {
     replicaId: storage.ReplicaId,
     checkpointGreaterThan: InternalOpId
   ) {
-    this.currentData.push({
+    this.currentData.set(this.currentDataKey(sourceTableId, replicaId), {
       sourceTableId,
       operation: {
         updateOne: {
@@ -191,7 +203,7 @@ export class PersistedBatchV3 extends PersistedBatch {
       };
     });
 
-    this.currentData.push({
+    this.currentData.set(this.currentDataKey(values.sourceTableId, values.replicaId), {
       sourceTableId: values.sourceTableId,
       operation: {
         updateOne: {
@@ -212,12 +224,12 @@ export class PersistedBatchV3 extends PersistedBatch {
   }
 
   protected get currentDataCount() {
-    return this.currentData.length;
+    return this.currentData.size;
   }
 
   // Flush methods
 
-  protected async flushBucketData(session: mongo.ClientSession) {
+  protected async queueBucketData(writes: MongoWriteBatch) {
     const operationsByDefinition = new Map<BucketDefinitionId, BucketDataDoc[]>();
     for (const document of this.bucketData) {
       const existing = operationsByDefinition.get(document.bucketKey.definitionId) ?? [];
@@ -225,15 +237,17 @@ export class PersistedBatchV3 extends PersistedBatch {
       operationsByDefinition.set(document.bucketKey.definitionId, existing);
     }
 
-    let uploadCount = 0;
     const plans = Array.from(operationsByDefinition, ([definitionId, documents]) => {
       const operationsByBucket = Map.groupBy(documents, (document) => document.bucketKey.bucket);
       const lifecycle = this.objectStorageLifecycle;
-      const createInserts: (() => Promise<mongo.AnyBulkWriteOperation<BucketDataDocumentV3>>)[] = [];
+      type BucketDataInsert = { insertOne: { document: BucketDataDocumentV3 } };
+      const createInserts: (() => Promise<BucketDataInsert>)[] = [];
 
       for (const [bucket, ops] of operationsByBucket) {
+        this.resetBucketPersistedBytes(definitionId, bucket);
         for (const chunk of chunkBucketData(ops)) {
           const serialized = serializeBucketData(bucket, chunk);
+          this.incrementBucketPersistedChunk(definitionId, bucket, serialized.size);
           if (lifecycle == null || serialized.size <= this.inlineThresholdBytes) {
             createInserts.push(async () => ({
               insertOne: {
@@ -243,13 +257,12 @@ export class PersistedBatchV3 extends PersistedBatch {
             continue;
           }
 
-          uploadCount += 1;
           createInserts.push(async () => {
             const minOp = chunk[0].o;
             const maxOp = chunk[chunk.length - 1].o;
             const { ops: bucketOps, ...metadata } = serialized;
             const path = lifecycle.allocatePath(definitionId, bucket, minOp, maxOp);
-            const { fileSize } = await lifecycle.bucketData.store(path, bucketOps!);
+            const { fileSize } = await lifecycle.bucketData.store(path, bucketOps!, { signal: this.signal });
             return {
               insertOne: {
                 document: {
@@ -279,19 +292,29 @@ export class PersistedBatchV3 extends PersistedBatch {
     // S3ObjectStorage applies one shared concurrency limit across all callers,
     // so replication can schedule its uploads together without creating a
     // separate limiter here.
-    const writes = await createAllInserts();
-
-    for (const { definitionId, inserts } of writes) {
+    const bucketWrites = await createAllInserts();
+
+    const usageDeltas = this.objectStorageUsage == null ? undefined : new Map<BucketDefinitionId, bigint>();
+    for (const { definitionId, inserts } of bucketWrites) {
+      if (usageDeltas != null) {
+        const delta = inserts.reduce(
+          (sum, operation) => sum + ObjectStorageUsage.bytes(operation.insertOne.document),
+          0n
+        );
+        if (delta !== 0n) {
+          usageDeltas.set(definitionId, delta);
+        }
+      }
       if (inserts.length > 0) {
-        await this.db.bucketData(this.group_id, definitionId).bulkWrite(inserts, {
-          session,
-          ordered: false
-        });
+        writes.bulkWriteUnordered(this.db.bucketData(this.group_id, definitionId), inserts);
       }
     }
+    if (usageDeltas != null) {
+      this.objectStorageUsage!.applyDeltas(usageDeltas, writes);
+    }
   }
 
-  protected async flushBucketParameters(session: mongo.ClientSession) {
+  protected queueBucketParameters(writes: MongoWriteBatch): void {
     const operationsByIndex = new Map<string, typeof this.bucketParameters>();
     for (const document of this.bucketParameters) {
       const existing = operationsByIndex.get(document.index) ?? [];
@@ -300,23 +323,20 @@ export class PersistedBatchV3 extends PersistedBatch {
     }
 
     for (const [indexId, documents] of operationsByIndex.entries()) {
-      await this.db.parameterIndex(this.group_id, indexId).bulkWrite(
+      writes.bulkWriteUnordered(
+        this.db.parameterIndex(this.group_id, indexId),
         documents.map((document) => ({
           insertOne: {
             document: taggedBucketParameterDocumentToTagged(document)
           }
-        })),
-        {
-          session,
-          ordered: false
-        }
+        }))
       );
     }
   }
 
-  protected async flushCurrentData(session: mongo.ClientSession) {
-    const operationsBySourceTable = new Map<string, typeof this.currentData>();
-    for (const operation of this.currentData) {
+  protected queueCurrentData(writes: MongoWriteBatch): void {
+    const operationsBySourceTable = new Map<string, SourceRecordWrite[]>();
+    for (const operation of this.currentData.values()) {
       const sourceTableId = operation.sourceTableId.toHexString();
       const existing = operationsBySourceTable.get(sourceTableId) ?? [];
       existing.push(operation);
@@ -339,33 +359,31 @@ export class PersistedBatchV3 extends PersistedBatch {
     });
 
     if (sourceTableUpdates.length > 0) {
-      await this.db.sourceTables(this.group_id).bulkWrite(sourceTableUpdates, { session, ordered: false });
+      writes.bulkWriteUnordered(this.db.sourceTables(this.group_id), sourceTableUpdates);
     }
 
     for (const operations of operationsBySourceTable.values()) {
       const sourceTableId = operations[0]!.sourceTableId;
-      await this.db.sourceRecords(this.group_id, sourceTableId).bulkWrite(
-        operations.map((entry) => entry.operation),
-        {
-          session,
-          ordered: true
-        }
+      writes.bulkWriteUnordered(
+        this.db.sourceRecords(this.group_id, sourceTableId),
+        operations.map((entry) => entry.operation)
       );
     }
   }
 
-  protected async flushBucketStates(session: mongo.ClientSession) {
-    await this.db.bucketState(this.group_id).bulkWrite(this.getBucketStateUpdates(), {
-      session,
-      ordered: false
-    });
+  protected queueBucketStates(writes: MongoWriteBatch): void {
+    writes.bulkWriteUnordered(this.db.bucketState(this.group_id), this.getBucketStateUpdates());
   }
 
   protected resetCurrentData() {
-    this.currentData = [];
+    this.currentData.clear();
     this.sourceTablePendingDeletes.clear();
   }
 
+  private currentDataKey(sourceTableId: bson.ObjectId, replicaId: storage.ReplicaId): string {
+    return Buffer.from(bson.serialize({ t: sourceTableId, k: replicaId })).toString('base64');
+  }
+
   private getBucketStateUpdates(): mongo.AnyBulkWriteOperation<BucketStateDocumentV3>[] {
     return Array.from(this.bucketStates.values()).map((state: BucketStateUpdate) => {
       if (state.definitionId == null) {
@@ -379,15 +397,42 @@ export class PersistedBatchV3 extends PersistedBatch {
               b: state.bucket
             }
           },
-          update: {
-            $set: {
-              last_op: state.lastOp
-            },
-            $inc: {
-              'estimate_since_compact.count': state.incrementCount,
-              'estimate_since_compact.bytes': state.incrementBytes
+          // A pipeline update makes initialisation and scheduling one atomic
+          // writer operation. In particular, a later write cannot move an
+          // already-due compact check into the future.
+          update: [
+            {
+              $set: {
+                last_op: state.lastOp,
+                bucket_stats: {
+                  count: { $add: [{ $ifNull: ['$bucket_stats.count', 0] }, state.incrementCount] },
+                  bytes: { $add: [{ $ifNull: ['$bucket_stats.bytes', 0n] }, BigInt(state.incrementBytes)] },
+                  chunks: { $add: [{ $ifNull: ['$bucket_stats.chunks', 0] }, state.incrementChunks] }
+                },
+                first_uncompacted_write: { $ifNull: ['$first_uncompacted_write', '$$NOW'] },
+                next_compact_check: {
+                  $let: {
+                    vars: {
+                      requested: {
+                        $dateAdd: {
+                          startDate: '$$NOW',
+                          unit: 'millisecond',
+                          amount: DEFAULT_MIN_COMPACT_CHUNK_INTERVAL_MS
+                        }
+                      }
+                    },
+                    in: {
+                      $cond: [
+                        { $lt: [{ $ifNull: ['$next_compact_check', '$$requested'] }, '$$requested'] },
+                        '$next_compact_check',
+                        '$$requested'
+                      ]
+                    }
+                  }
+                }
+              }
             }
-          },
+          ],
           upsert: true
         }
       } satisfies mongo.AnyBulkWriteOperation<BucketStateDocumentV3>;

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

@@ -1,5 +1,5 @@
 import { mongo } from '@powersync/lib-service-mongodb';
-import { BucketDefinitionId, ParameterIndexId } from '@powersync/service-sync-rules';
+import { BucketDefinitionId, EventDefinitionId, ParameterIndexId } from '@powersync/service-sync-rules';
 import { BaseVersionedPowerSyncMongo } from '../common/VersionedPowerSyncMongoBase.js';
 import { CommonSourceTableDocument } from '../models.js';
 import {
@@ -7,11 +7,19 @@ import {
   BucketParameterDocumentV3,
   BucketStateDocumentV3,
   CurrentDataDocumentV3,
+  CustomCheckpointRequestDocumentV3,
   ObjectStorageDeletionMarker,
+  ObjectStorageUsageDocument,
   SourceTableDocumentV3,
   SyncConfigDefinition
 } from './models.js';
 
+export const OBJECT_STORAGE_USAGE_COLLECTION = 'object_storage_usage';
+export type CustomCheckpointRequestCollectionSpecifier = {
+  replicationStreamId: number;
+  eventId: EventDefinitionId;
+};
+
 export class VersionedPowerSyncMongoV3 extends BaseVersionedPowerSyncMongo {
   constructor(
     upstream: ConstructorParameters<typeof BaseVersionedPowerSyncMongo>[0],
@@ -74,6 +82,42 @@ export class VersionedPowerSyncMongoV3 extends BaseVersionedPowerSyncMongo {
     return this.db.collection<ObjectStorageDeletionMarker>(`pending_object_storage_deletes_${replicationStreamId}`);
   }
 
+  get objectStorageUsage(): mongo.Collection<ObjectStorageUsageDocument> {
+    return this.db.collection<ObjectStorageUsageDocument>(OBJECT_STORAGE_USAGE_COLLECTION);
+  }
+
+  customCheckpointRequests(options: CustomCheckpointRequestCollectionSpecifier) {
+    const collectionName = this.customCheckpointRequestCollectionName(options);
+    return this.db.collection<CustomCheckpointRequestDocumentV3>(collectionName);
+  }
+
+  async initializeCustomCheckpointRequestsCollection(options: CustomCheckpointRequestCollectionSpecifier) {
+    // createIndexes also creates the collection on first use. Callers invoke
+    // this lazily before the transaction which writes the first checkpoint.
+    await this.customCheckpointRequests(options).createIndexes([
+      {
+        key: { user_id: 1 },
+        name: 'user_unique',
+        unique: true
+      },
+      {
+        // Find checkpoint changes over an operation-id range.
+        key: { op_id: 1 },
+        name: 'op_id'
+      },
+      {
+        // Limit retention scans to temporary client-requested checkpoints.
+        key: { checkpoint_requested_at: 1 },
+        name: 'checkpoint_requested_at',
+        partialFilterExpression: { checkpoint_requested_at: { $exists: true } }
+      }
+    ]);
+  }
+
+  listCustomCheckpointRequestCollections(replicationStreamId: number) {
+    return this.upstream.listCustomCheckpointRequestCollections(replicationStreamId);
+  }
+
   listBucketDataCollections(replicationStreamId: number) {
     return this.listCollectionsByPrefix(`bucket_data_${replicationStreamId}_`);
   }
@@ -117,11 +161,18 @@ export class VersionedPowerSyncMongoV3 extends BaseVersionedPowerSyncMongo {
       },
       { name: 'bucket_updates', unique: true }
     );
+    // V3 workers claim only buckets that have reached their scheduled check
+    // time, keeping scheduled scans proportional to pending work.
     await bucketState.createIndex(
+      { next_compact_check: 1 },
       {
-        'estimate_since_compact.count': -1
-      },
-      { name: 'dirty_count' }
+        name: 'next_compact_check',
+        partialFilterExpression: { next_compact_check: { $exists: true } }
+      }
     );
   }
+
+  customCheckpointRequestCollectionName(options: CustomCheckpointRequestCollectionSpecifier) {
+    return this.upstream.customCheckpointRequestCollectionName(options);
+  }
 }

package/src/storage/implementation/v3/bucket-format.ts

@@ -32,6 +32,7 @@ export function serializeBucketData(
       op: op.op,
       source_table: op.source_table,
       source_key: op.source_key,
+      subkey: op.subkey,
       table: op.table,
       row_id: op.row_id,
       checksum: op.checksum,
@@ -78,6 +79,7 @@ export function* loadBucketDataDocument(
       op: op.op,
       source_table: op.source_table,
       source_key: op.source_key,
+      subkey: op.subkey,
       table: op.table,
       row_id: op.row_id,
       checksum: op.checksum,