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

package/src/storage/implementation/MongoSyncBucketStorage.ts

@@ -26,8 +26,9 @@ import { HydratedSyncConfig, ParameterLookupRows, ScopedParameterLookup } from '
 import * as bson from 'bson';
 import { LRUCache } from 'lru-cache';
 import * as timers from 'timers/promises';
-import { retryOnMongoMaxTimeMSExpired } from '../../utils/util.js';
+import { DEFAULT_CLEAR_BATCH_THROTTLE_RATE } from '../../types/types.js';
 import { MongoBucketStorage } from '../MongoBucketStorage.js';
+import { DEFAULT_INLINE_THRESHOLD_BYTES } from './common/PersistedBatch.js';
 import type { VersionedPowerSyncMongo } from './db.js';
 import { StorageConfig } from './models.js';
 import { MongoBucketBatchOptions } from './MongoBucketBatch.js';
@@ -37,12 +38,16 @@ import { MongoParameterCompactor } from './MongoParameterCompactor.js';
 import { MongoParsedSyncConfigSet } from './MongoParsedSyncConfigSet.js';
 import { MongoPersistedReplicationStream } from './MongoPersistedReplicationStream.js';
 import { MongoWriteCheckpointAPI } from './MongoWriteCheckpointAPI.js';
+import { ObjectStorage } from './v3/object-storage/ObjectStorage.js';
 
 export interface MongoSyncBucketStorageOptions {
   checksumOptions?: Omit<MongoChecksumOptions, 'storageConfig'>;
   readPreference?: mongo.ReadPreference;
   checksumCacheTtlMs?: number;
+  clearBatchThrottleRate?: number;
   storageConfig: StorageConfig;
+  objectStorage?: ObjectStorage;
+  inlineThresholdBytes?: number;
 }
 
 interface InternalCheckpointChanges extends CheckpointChanges {
@@ -71,6 +76,9 @@ export abstract class MongoSyncBucketStorage
 
   readonly checksums: MongoChecksums;
 
+  readonly objectStorage?: ObjectStorage;
+  readonly inlineThresholdBytes: number;
+
   /**
    * Canonical parsed sync config sets, keyed by defaultSchema.
    *
@@ -83,6 +91,7 @@ export abstract class MongoSyncBucketStorage
   public readonly logger: Logger;
   public readonly storageConfig: StorageConfig;
   public readonly readPreference: mongo.ReadPreference | undefined;
+  public readonly clearBatchThrottleRate: number;
   #storageInitialized = false;
 
   constructor(
@@ -95,7 +104,12 @@ export abstract class MongoSyncBucketStorage
   ) {
     super();
     this.storageConfig = options.storageConfig;
+    this.objectStorage = options.objectStorage;
+    // Keep small chunks inline in MongoDB rather than offloading them to S3.
+    // Configurable via object_storage.inline_threshold_bytes.
+    this.inlineThresholdBytes = options.inlineThresholdBytes ?? DEFAULT_INLINE_THRESHOLD_BYTES;
     this.readPreference = options.readPreference;
+    this.clearBatchThrottleRate = options.clearBatchThrottleRate ?? DEFAULT_CLEAR_BATCH_THROTTLE_RATE;
     this.db = factory.db.versioned(this.storageConfig);
     this.checksums = this.createMongoChecksums(options);
     this.writeCheckpointAPI = new MongoWriteCheckpointAPI({
@@ -134,7 +148,9 @@ export abstract class MongoSyncBucketStorage
     this.writeCheckpointAPI.setWriteCheckpointMode(mode);
   }
 
-  createManagedWriteCheckpoints(checkpoints: storage.ManagedWriteCheckpointOptions[]): Promise<Map<string, bigint>> {
+  createManagedWriteCheckpoints(
+    checkpoints: storage.ManagedWriteCheckpointOptions[]
+  ): Promise<storage.CreateManagedWriteCheckpointsResult> {
     return this.writeCheckpointAPI.createManagedWriteCheckpoints(checkpoints);
   }
 
@@ -218,7 +234,9 @@ export abstract class MongoSyncBucketStorage
       skipExistingRows: options.skipExistingRows ?? false,
       markRecordUnavailable: options.markRecordUnavailable,
       hooks: options.hooks,
-      tracer: options.tracer
+      tracer: options.tracer,
+      objectStorage: this.objectStorage,
+      inlineThresholdBytes: this.inlineThresholdBytes
     };
   }
 
@@ -258,13 +276,13 @@ export abstract class MongoSyncBucketStorage
     checkpoint: MongoReplicationCheckpoint,
     dataBuckets: storage.BucketDataRequest[],
     options?: storage.BucketDataBatchOptions
-  ): AsyncIterable<storage.SyncBucketDataChunk>;
+  ): AsyncIterable<storage.SyncBucketDataChunk | storage.SyncBucketDataBatchEnd>;
 
   async *getBucketDataBatch(
     checkpoint: storage.ReplicationCheckpoint,
     dataBuckets: storage.BucketDataRequest[],
     options?: storage.BucketDataBatchOptions
-  ): AsyncIterable<storage.SyncBucketDataChunk> {
+  ): AsyncIterable<storage.SyncBucketDataChunk | storage.SyncBucketDataBatchEnd> {
     yield* this.getBucketDataBatchImpl(checkpoint as MongoReplicationCheckpoint, dataBuckets, options);
   }
 
@@ -331,23 +349,6 @@ export abstract class MongoSyncBucketStorage
     this.#storageInitialized = false;
   }
 
-  protected async clearDeleteMany(
-    label: string,
-    operation: () => Promise<mongo.DeleteResult>,
-    signal?: AbortSignal
-  ): Promise<void> {
-    await retryOnMongoMaxTimeMSExpired(operation, {
-      signal,
-      abortMessage: 'Aborted clearing data',
-      retryDelayMs: lib_mongo.db.MONGO_CLEAR_OPERATION_TIMEOUT_MS / 5,
-      onRetry: () => {
-        this.logger.info(
-          `Cleared batch of ${label} in ${lib_mongo.db.MONGO_CLEAR_OPERATION_TIMEOUT_MS}ms, continuing...`
-        );
-      }
-    });
-  }
-
   async reportError(e: any): Promise<void> {
     const message = String(e.message ?? 'Replication failure');
     await this.db.sync_rules.updateOne(

package/src/storage/implementation/MongoWriteCheckpointAPI.ts

@@ -28,21 +28,26 @@ export class MongoWriteCheckpointAPI implements storage.WriteCheckpointAPI {
 
   async createManagedWriteCheckpoints(
     checkpoints: storage.ManagedWriteCheckpointOptions[]
-  ): Promise<Map<string, bigint>> {
+  ): Promise<storage.CreateManagedWriteCheckpointsResult> {
     if (this.writeCheckpointMode !== storage.WriteCheckpointMode.MANAGED) {
       throw new framework.ServiceAssertionError(
         `Attempting to create a managed Write Checkpoint when the current Write Checkpoint mode is set to "${this.writeCheckpointMode}"`
       );
     }
 
-    const uniqueCheckpoints = [...new Map(checkpoints.map((checkpoint) => [checkpoint.user_id, checkpoint])).values()];
+    const uniqueCheckpoints = storage.uniqueManagedWriteCheckpoints(checkpoints);
     if (uniqueCheckpoints.length == 0) {
-      return new Map();
+      return { writeCheckpoints: new Map(), shouldAdvance: false };
     }
 
-    if (uniqueCheckpoints.length == 1) {
+    let shouldAdvance = false;
+    const writeCheckpoints = new Map<string, bigint>();
+    const generatedCheckpoints = uniqueCheckpoints.filter((checkpoint) => checkpoint.checkpoint_request_id == null);
+    const suppliedCheckpoints = uniqueCheckpoints.filter((checkpoint) => checkpoint.checkpoint_request_id != null);
+
+    if (generatedCheckpoints.length == 1) {
       // For the common case of a single checkpoint, we can do this in a single request.
-      const { user_id, heads: lsns } = uniqueCheckpoints[0];
+      const { user_id, heads: lsns } = generatedCheckpoints[0];
       const doc = await this.db.write_checkpoints.findOneAndUpdate(
         {
           user_id
@@ -52,6 +57,11 @@ export class MongoWriteCheckpointAPI implements storage.WriteCheckpointAPI {
             lsns,
             processed_at_lsn: null
           },
+          // Unset rather than set to null, to keep the document out of the
+          // partial checkpoint_requested_at index used for retention cleanup.
+          $unset: {
+            checkpoint_requested_at: 1
+          },
           $inc: {
             client_id: 1n
           }
@@ -59,29 +69,123 @@ export class MongoWriteCheckpointAPI implements storage.WriteCheckpointAPI {
         { upsert: true, returnDocument: 'after' }
       );
 
-      return new Map([[doc!.user_id, doc!.client_id]]);
+      // Generated checkpoints always reset processed_at_lsn, so they are pending
+      // and the source marker must be forced.
+      shouldAdvance = true;
+      writeCheckpoints.set(doc!.user_id, doc!.client_id);
+    } else if (generatedCheckpoints.length > 1) {
+      // For more than one generated checkpoint, this gives a constant 2 requests.
+      await this.db.write_checkpoints.bulkWrite(
+        generatedCheckpoints.map(({ user_id, heads: lsns }) => ({
+          updateOne: {
+            filter: { user_id },
+            update: {
+              $set: {
+                lsns,
+                processed_at_lsn: null
+              },
+              $unset: {
+                checkpoint_requested_at: 1
+              },
+              $inc: {
+                client_id: 1n
+              }
+            },
+            upsert: true
+          }
+        }))
+      );
+
+      const userIds = generatedCheckpoints.map((checkpoint) => checkpoint.user_id);
+      const docs = await this.db.write_checkpoints
+        .find(
+          {
+            user_id: { $in: userIds }
+          },
+          {
+            projection: {
+              user_id: 1,
+              client_id: 1
+            }
+          }
+        )
+        .toArray();
+
+      shouldAdvance = true;
+      for (const doc of docs) {
+        writeCheckpoints.set(doc.user_id, doc.client_id);
+      }
     }
 
-    // For more than one checkpoint, this gives a constant 2 requests
+    if (suppliedCheckpoints.length > 0) {
+      const suppliedResult = await this.createSuppliedManagedWriteCheckpoints(suppliedCheckpoints);
+      shouldAdvance ||= suppliedResult.shouldAdvance;
+      for (const [userId, writeCheckpoint] of suppliedResult.writeCheckpoints) {
+        writeCheckpoints.set(userId, writeCheckpoint);
+      }
+    }
+
+    return { writeCheckpoints, shouldAdvance };
+  }
+
+  private async createSuppliedManagedWriteCheckpoints(checkpoints: storage.ManagedWriteCheckpointOptions[]) {
+    const checkpointRequestedAt = new Date();
+
     await this.db.write_checkpoints.bulkWrite(
-      uniqueCheckpoints.map(({ user_id, heads: lsns }) => ({
-        updateOne: {
-          filter: { user_id },
-          update: {
-            $set: {
-              lsns,
-              processed_at_lsn: null
-            },
-            $inc: {
-              client_id: 1n
-            }
-          },
-          upsert: true
-        }
-      }))
+      checkpoints.map((checkpoint) => {
+        const { user_id, heads: lsns } = checkpoint;
+        const checkpointRequestId = checkpoint.checkpoint_request_id!;
+
+        // Supplied request ids are monotonic. updateMany preserves the existing
+        // upsert behavior for first checkpoint requests, while also advancing
+        // every duplicate row for the user_id that is lower than this request id.
+        const shouldApplyRequestId = {
+          $or: [
+            { $eq: [{ $ifNull: ['$client_id', null] }, null] },
+            { $lt: ['$client_id', { $literal: checkpointRequestId }] }
+          ]
+        };
+        // Equal retries keep their original source head but refresh retention.
+        const shouldRefreshRequest = {
+          $or: [
+            { $eq: [{ $ifNull: ['$client_id', null] }, null] },
+            { $lte: ['$client_id', { $literal: checkpointRequestId }] }
+          ]
+        };
+
+        return {
+          updateMany: {
+            filter: { user_id },
+            update: [
+              {
+                $set: {
+                  user_id,
+                  client_id: {
+                    $cond: [shouldApplyRequestId, { $literal: checkpointRequestId }, '$client_id']
+                  },
+                  lsns: {
+                    $cond: [shouldApplyRequestId, { $literal: lsns }, '$lsns']
+                  },
+                  processed_at_lsn: {
+                    $cond: [shouldApplyRequestId, null, '$processed_at_lsn']
+                  },
+                  checkpoint_requested_at: {
+                    $cond: [shouldRefreshRequest, { $literal: checkpointRequestedAt }, '$checkpoint_requested_at']
+                  }
+                }
+              }
+            ],
+            upsert: true
+          }
+        };
+      })
     );
 
-    const userIds = uniqueCheckpoints.map((checkpoint) => checkpoint.user_id);
+    // Fetch the final ids separately so stale requests can still return the
+    // checkpoint currently stored. If duplicate rows already exist for a
+    // user_id, use the highest stored id instead of whichever document MongoDB
+    // happens to return first.
+    const userIds = checkpoints.map((checkpoint) => checkpoint.user_id);
     const docs = await this.db.write_checkpoints
       .find(
         {
@@ -90,13 +194,28 @@ export class MongoWriteCheckpointAPI implements storage.WriteCheckpointAPI {
         {
           projection: {
             user_id: 1,
-            client_id: 1
+            client_id: 1,
+            processed_at_lsn: 1
           }
         }
       )
       .toArray();
+    const maxDocsByUser = new Map<string, (typeof docs)[number]>();
+    for (const doc of docs) {
+      const existing = maxDocsByUser.get(doc.user_id);
+      if (
+        existing == null ||
+        doc.client_id > existing.client_id ||
+        (doc.client_id == existing.client_id && existing.processed_at_lsn != null && doc.processed_at_lsn == null)
+      ) {
+        maxDocsByUser.set(doc.user_id, doc);
+      }
+    }
 
-    return new Map(docs.map((doc) => [doc.user_id, doc.client_id]));
+    return {
+      writeCheckpoints: new Map([...maxDocsByUser].map(([userId, doc]) => [userId, doc.client_id])),
+      shouldAdvance: [...maxDocsByUser.values()].some((doc) => doc.processed_at_lsn == null)
+    };
   }
 
   async lastWriteCheckpoint(filters: storage.LastWriteCheckpointFilters): Promise<bigint | null> {
@@ -225,19 +344,27 @@ export async function batchCreateCustomWriteCheckpoints(
   }
 
   await db.custom_write_checkpoints.bulkWrite(
-    checkpoints.map((checkpointOptions) => ({
-      updateOne: {
-        filter: { user_id: checkpointOptions.user_id, sync_rules_id: checkpointOptions.sync_rules_id },
-        update: {
-          $set: {
-            checkpoint: checkpointOptions.checkpoint,
-            sync_rules_id: checkpointOptions.sync_rules_id,
-            op_id: opId
-          }
-        },
-        upsert: true
+    checkpoints.map((checkpointOptions) => {
+      const set: Record<string, unknown> = {
+        checkpoint: checkpointOptions.checkpoint,
+        sync_rules_id: checkpointOptions.sync_rules_id,
+        op_id: opId
+      };
+      if (checkpointOptions.checkpoint_requested_at != null) {
+        set.checkpoint_requested_at = checkpointOptions.checkpoint_requested_at;
       }
-    })),
+
+      return {
+        updateOne: {
+          filter: { user_id: checkpointOptions.user_id, sync_rules_id: checkpointOptions.sync_rules_id },
+          update: {
+            $set: set,
+            ...(checkpointOptions.checkpoint_requested_at == null ? { $unset: { checkpoint_requested_at: 1 } } : {})
+          },
+          upsert: true
+        }
+      };
+    }),
     { session }
   );
 }

package/src/storage/implementation/common/PersistedBatch.ts

@@ -10,6 +10,7 @@ import { currentBucketKey, MAX_ROW_SIZE } from '../MongoBucketBatchShared.js';
 import { MongoIdSequence } from '../MongoIdSequence.js';
 import type { VersionedPowerSyncMongo } from '../db.js';
 import { TaggedBucketParameterDocument } from '../models.js';
+import { ObjectStorage } from '../v3/object-storage/ObjectStorage.js';
 import { BucketDataDoc, BucketKey } from './BucketDataDoc.js';
 import { SourceRecordBucketState, SourceRecordLookupState } from './SourceRecordStore.js';
 
@@ -34,6 +35,8 @@ const MAX_TRANSACTION_BATCH_SIZE = 30_000_000;
  */
 const MAX_TRANSACTION_DOC_COUNT = 2_000;
 
+export const DEFAULT_INLINE_THRESHOLD_BYTES = 16 * 1024;
+
 export interface SaveBucketDataOptions {
   op_seq: MongoIdSequence;
   sourceKey: storage.ReplicaId;
@@ -60,6 +63,8 @@ export interface UpsertCurrentDataOptions {
 
 export interface PersistedBatchOptions {
   logger?: Logger;
+  objectStorage?: ObjectStorage;
+  inlineThresholdBytes?: number;
 }
 
 /**
@@ -74,6 +79,9 @@ export abstract class PersistedBatch {
   bucketParameters: TaggedBucketParameterDocument[] = [];
   bucketStates: Map<string, BucketStateUpdate> = new Map();
 
+  protected readonly objectStorage?: ObjectStorage;
+  protected readonly inlineThresholdBytes: number = DEFAULT_INLINE_THRESHOLD_BYTES;
+
   /**
    * For debug logging only.
    */
@@ -93,6 +101,10 @@ export abstract class PersistedBatch {
   ) {
     this.currentSize = writtenSize;
     this.logger = options?.logger ?? defaultLogger;
+    this.objectStorage = options?.objectStorage;
+    if (options?.inlineThresholdBytes != null) {
+      this.inlineThresholdBytes = options.inlineThresholdBytes;
+    }
   }
 
   saveBucketData(options: SaveBucketDataOptions) {
@@ -309,32 +321,32 @@ export abstract class PersistedBatch {
         const replicationLag = Math.round((Date.now() - options.oldestUncommittedChange.getTime()) / 1000);
 
         this.logger.info(
-          `Flushed ${this.bucketDataCount} + ${this.bucketParameters.length} + ${
+          `Flushed: ${this.bucketDataCount} ops, ${this.bucketParameters.length} index entries, ${
             this.currentDataCount
-          } updates, ${Math.round(this.currentSize / 1024)}kb in ${duration}ms. Last op_id: ${this.debugLastOpId}. Replication lag: ${replicationLag}s`,
+          } records. ${Math.round(this.currentSize / 1024)}kb in ${duration}ms. Last op_id: ${this.debugLastOpId}. Replication lag: ${replicationLag}s`,
           {
             flushed: {
               duration: duration,
               size: this.currentSize,
-              bucket_data_count: this.bucketDataCount,
-              parameter_data_count: this.bucketParameters.length,
-              current_data_count: this.currentDataCount,
+              bucket_ops_count: this.bucketDataCount,
+              parameter_indexes_count: this.bucketParameters.length,
+              source_records_count: this.currentDataCount,
               replication_lag_seconds: replicationLag
             }
           }
         );
       } else {
         this.logger.info(
-          `Flushed ${this.bucketDataCount} + ${this.bucketParameters.length} + ${
+          `Flushed: ${this.bucketDataCount} ops, ${this.bucketParameters.length} index entries, ${
             this.currentDataCount
-          } updates, ${Math.round(this.currentSize / 1024)}kb in ${duration}ms. Last op_id: ${this.debugLastOpId}`,
+          } records. ${Math.round(this.currentSize / 1024)}kb in ${duration}ms. Last op_id: ${this.debugLastOpId}`,
           {
             flushed: {
               duration: duration,
               size: this.currentSize,
-              bucket_data_count: this.bucketDataCount,
-              parameter_data_count: this.bucketParameters.length,
-              current_data_count: this.currentDataCount
+              bucket_ops_count: this.bucketDataCount,
+              parameter_indexes_count: this.bucketParameters.length,
+              source_records_count: this.currentDataCount
             }
           }
         );

package/src/storage/implementation/models.ts

@@ -78,6 +78,11 @@ export interface BucketDataProperties {
   row_id?: string;
   checksum: bigint;
   data: string | null;
+  /**
+   * V1-only.
+   *
+   * V3 stores this on the BucketDataDocumentV3 instead of the individual ops.
+   */
   target_op?: bigint | null;
 }
 
@@ -105,6 +110,10 @@ export interface SourceTableDocument {
   replica_id_columns2: { name: string; type_oid?: number; type?: string }[] | undefined;
   snapshot_done: boolean | undefined;
   snapshot_status: SourceTableDocumentSnapshotStatus | undefined;
+  /**
+   * Source-specific metadata. Absent for legacy records.
+   */
+  source_metadata?: storage.JsonValue;
 }
 
 export interface SourceTableDocumentSnapshotStatus {
@@ -273,6 +282,11 @@ export interface CustomWriteCheckpointDocument {
    * This is not unique - multiple write checkpoints can have the same op_id.
    */
   op_id?: InternalOpId;
+  /**
+   * Set when this checkpoint was created from a client-supplied checkpoint
+   * request rather than a persistent custom write checkpoint.
+   */
+  checkpoint_requested_at?: Date | null;
 }
 
 export interface WriteCheckpointDocument {
@@ -286,6 +300,12 @@ export interface WriteCheckpointDocument {
    * between two checkpoints.
    */
   processed_at_lsn: string | null;
+  /**
+   * Set when this checkpoint was created from a client-supplied checkpoint request
+   * rather than a server-generated id. A missing/null value means this is not a
+   * checkpoint request.
+   */
+  checkpoint_requested_at?: Date | null;
 }
 
 export interface InstanceDocument {