@powersync/service-module-mongodb 0.15.1 → 0.15.3

package/src/replication/ChangeStream.ts

@@ -12,6 +12,7 @@ import {
 import {
   MetricsEngine,
   RelationCache,
+  ReplicationLagTracker,
   SaveOperationTag,
   SourceEntityDescriptor,
   SourceTable,
@@ -19,15 +20,16 @@ import {
 } from '@powersync/service-core';
 import {
   DatabaseInputRow,
+  HydratedSyncRules,
   SqliteInputRow,
   SqliteRow,
-  HydratedSyncRules,
   TablePattern
 } from '@powersync/service-sync-rules';
 import { ReplicationMetric } from '@powersync/service-types';
 import { MongoLSN } from '../common/MongoLSN.js';
 import { PostImagesOption } from '../types/types.js';
 import { escapeRegExp } from '../utils.js';
+import { trackChangeStreamBsonBytes } from './internal-mongodb-utils.js';
 import { MongoManager } from './MongoManager.js';
 import {
   constructAfterRecord,
@@ -98,16 +100,7 @@ export class ChangeStream {
 
   private relationCache = new RelationCache(getCacheIdentifier);
 
-  /**
-   * Time of the oldest uncommitted change, according to the source db.
-   * This is used to determine the replication lag.
-   */
-  private oldestUncommittedChange: Date | null = null;
-  /**
-   * Keep track of whether we have done a commit or keepalive yet.
-   * We can only compute replication lag if isStartingReplication == false, or oldestUncommittedChange is present.
-   */
-  private isStartingReplication = true;
+  private replicationLag = new ReplicationLagTracker();
 
   private checkpointStreamId = new mongo.ObjectId();
 
@@ -479,6 +472,10 @@ export class ChangeStream {
   }
 
   private async snapshotTable(batch: storage.BucketStorageBatch, table: storage.SourceTable) {
+    const rowsReplicatedMetric = this.metrics.getCounter(ReplicationMetric.ROWS_REPLICATED);
+    const bytesReplicatedMetric = this.metrics.getCounter(ReplicationMetric.DATA_REPLICATED_BYTES);
+    const chunksReplicatedMetric = this.metrics.getCounter(ReplicationMetric.CHUNKS_REPLICATED);
+
     const totalEstimatedCount = await this.estimatedCountNumber(table);
     let at = table.snapshotStatus?.replicatedCount ?? 0;
     const db = this.client.db(table.schema);
@@ -499,11 +496,13 @@ export class ChangeStream {
     let lastBatch = performance.now();
     let nextChunkPromise = query.nextChunk();
     while (true) {
-      const { docs: docBatch, lastKey } = await nextChunkPromise;
+      const { docs: docBatch, lastKey, bytes: chunkBytes } = await nextChunkPromise;
       if (docBatch.length == 0) {
         // No more data - stop iterating
         break;
       }
+      bytesReplicatedMetric.add(chunkBytes);
+      chunksReplicatedMetric.add(1);
 
       if (this.abort_signal.aborted) {
         throw new ReplicationAbortedError(`Aborted initial replication`, this.abort_signal.reason);
@@ -528,7 +527,7 @@ export class ChangeStream {
       // Important: flush before marking progress
       await batch.flush();
       at += docBatch.length;
-      this.metrics.getCounter(ReplicationMetric.ROWS_REPLICATED).add(docBatch.length);
+      rowsReplicatedMetric.add(docBatch.length);
 
       table = await batch.updateTableProgress(table, {
         lastKey,
@@ -819,7 +818,30 @@ export class ChangeStream {
     };
   }
 
+  private getBufferedChangeCount(stream: mongo.ChangeStream<mongo.Document>): number {
+    // The driver keeps fetched change stream documents on the underlying cursor, but does
+    // not expose that through the public ChangeStream API. We use this to detect backlog
+    // building up before we have processed the corresponding source changes locally.
+    // If the driver API changes, we'll have a hard error here.
+    // We specifically want to avoid a silent performance regression if the driver behavior changes.
+    const cursor = (
+      stream as mongo.ChangeStream<mongo.Document> & {
+        cursor: mongo.AbstractCursor<mongo.ChangeStreamDocument<mongo.Document>>;
+      }
+    ).cursor;
+    if (cursor == null || typeof cursor.bufferedCount != 'function') {
+      throw new ReplicationAssertionError(
+        'MongoDB ChangeStream no longer exposes an internal cursor with bufferedCount'
+      );
+    }
+    return cursor.bufferedCount();
+  }
+
   async streamChangesInternal() {
+    const transactionsReplicatedMetric = this.metrics.getCounter(ReplicationMetric.TRANSACTIONS_REPLICATED);
+    const bytesReplicatedMetric = this.metrics.getCounter(ReplicationMetric.DATA_REPLICATED_BYTES);
+    const chunksReplicatedMetric = this.metrics.getCounter(ReplicationMetric.CHUNKS_REPLICATED);
+
     await this.storage.startBatch(
       {
         logger: this.logger,
@@ -848,6 +870,11 @@ export class ChangeStream {
           await stream.close();
           return;
         }
+        trackChangeStreamBsonBytes(stream, (bytes) => {
+          bytesReplicatedMetric.add(bytes);
+          // Each of these represent a single response message from MongoDB.
+          chunksReplicatedMetric.add(1);
+        });
 
         // Always start with a checkpoint.
         // This helps us to clear errors when restarting, even if there is
@@ -864,6 +891,7 @@ export class ChangeStream {
         let changesSinceLastCheckpoint = 0;
 
         let lastEmptyResume = performance.now();
+        let lastTxnKey: string | null = null;
 
         while (true) {
           if (this.abort_signal.aborted) {
@@ -903,7 +931,7 @@ export class ChangeStream {
               this.logger.info(
                 `Idle change stream. Persisted resumeToken for ${timestampToDate(timestamp).toISOString()}`
               );
-              this.isStartingReplication = false;
+              this.replicationLag.markStarted();
             }
             continue;
           }
@@ -1005,7 +1033,20 @@ export class ChangeStream {
             // However, these typically have a much lower rate than batch checkpoints, so we don't do that for now.
 
             const checkpointId = changeDocument.documentKey._id as string | mongo.ObjectId;
-            if (!(checkpointId == STANDALONE_CHECKPOINT_ID || this.checkpointStreamId.equals(checkpointId))) {
+
+            if (checkpointId == STANDALONE_CHECKPOINT_ID) {
+              // Standalone / write checkpoint received.
+              // When we are caught up, commit immediately to keep write checkpoint latency low.
+              // Once there is already a batch checkpoint pending, or the driver has buffered more
+              // change stream events, collapse standalone checkpoints into the normal batch
+              // checkpoint flow to avoid commit churn under sustained load.
+              if (waitForCheckpointLsn != null || this.getBufferedChangeCount(stream) > 0) {
+                if (waitForCheckpointLsn == null) {
+                  waitForCheckpointLsn = await createCheckpoint(this.client, this.defaultDb, this.checkpointStreamId);
+                }
+                continue;
+              }
+            } else if (!this.checkpointStreamId.equals(checkpointId)) {
               continue;
             }
             const { comparable: lsn } = new MongoLSN({
@@ -1027,12 +1068,11 @@ export class ChangeStream {
               waitForCheckpointLsn = null;
             }
             const { checkpointBlocked } = await batch.commit(lsn, {
-              oldestUncommittedChange: this.oldestUncommittedChange
+              oldestUncommittedChange: this.replicationLag.oldestUncommittedChange
             });
 
             if (!checkpointBlocked) {
-              this.oldestUncommittedChange = null;
-              this.isStartingReplication = false;
+              this.replicationLag.markCommitted();
               changesSinceLastCheckpoint = 0;
             }
           } else if (
@@ -1044,6 +1084,7 @@ export class ChangeStream {
             if (waitForCheckpointLsn == null) {
               waitForCheckpointLsn = await createCheckpoint(this.client, this.defaultDb, this.checkpointStreamId);
             }
+
             const rel = getMongoRelation(changeDocument.ns);
             const table = await this.getRelation(batch, rel, {
               // In most cases, we should not need to snapshot this. But if this is the first time we see the collection
@@ -1053,9 +1094,20 @@ export class ChangeStream {
               snapshot: true
             });
             if (table.syncAny) {
-              if (this.oldestUncommittedChange == null && changeDocument.clusterTime != null) {
-                this.oldestUncommittedChange = timestampToDate(changeDocument.clusterTime);
+              this.replicationLag.trackUncommittedChange(
+                changeDocument.clusterTime == null ? null : timestampToDate(changeDocument.clusterTime)
+              );
+
+              const transactionKeyValue = transactionKey(changeDocument);
+
+              if (transactionKeyValue == null || lastTxnKey != transactionKeyValue) {
+                // Very crude metric for counting transactions replicated.
+                // We ignore operations other than basic CRUD, and ignore changes to _powersync_checkpoints.
+                // Individual writes may not have a txnNumber, in which case we count them as separate transactions.
+                lastTxnKey = transactionKeyValue;
+                transactionsReplicatedMetric.add(1);
               }
+
               const flushResult = await this.writeChange(batch, table, changeDocument);
               changesSinceLastCheckpoint += 1;
               if (flushResult != null && changesSinceLastCheckpoint >= 20_000) {
@@ -1106,17 +1158,8 @@ export class ChangeStream {
     );
   }
 
-  async getReplicationLagMillis(): Promise<number | undefined> {
-    if (this.oldestUncommittedChange == null) {
-      if (this.isStartingReplication) {
-        // We don't have anything to compute replication lag with yet.
-        return undefined;
-      } else {
-        // We don't have any uncommitted changes, so replication is up-to-date.
-        return 0;
-      }
-    }
-    return Date.now() - this.oldestUncommittedChange.getTime();
+  getReplicationLagMillis(): number | undefined {
+    return this.replicationLag.getLagMillis();
   }
 
   private lastTouchedAt = performance.now();
@@ -1153,3 +1196,13 @@ function mapChangeStreamError(e: any) {
     throw new DatabaseConnectionError(ErrorCode.PSYNC_S1346, `Error reading MongoDB ChangeStream`, e);
   }
 }
+
+/**
+ * Transaction key for a change stream event, used to detect transaction boundaries. Returns null if the event is not part of a transaction.
+ */
+function transactionKey(doc: mongo.ChangeStreamDocument): string | null {
+  if (doc.txnNumber == null || doc.lsid == null) {
+    return null;
+  }
+  return `${doc.lsid.id.toString('hex')}:${doc.txnNumber}`;
+}

package/src/replication/ChangeStreamReplicationJob.ts

@@ -80,7 +80,7 @@ export class ChangeStreamReplicationJob extends replication.AbstractReplicationJ
     }
   }
 
-  async getReplicationLagMillis(): Promise<number | undefined> {
+  getReplicationLagMillis(): number | undefined {
     return this.lastStream?.getReplicationLagMillis();
   }
 }

package/src/replication/ChangeStreamReplicator.ts

@@ -41,25 +41,4 @@ export class ChangeStreamReplicator extends replication.AbstractReplicator<Chang
   async testConnection() {
     return await MongoModule.testConnection(this.connectionFactory.dbConnectionConfig);
   }
-
-  async getReplicationLagMillis(): Promise<number | undefined> {
-    const lag = await super.getReplicationLagMillis();
-    if (lag != null) {
-      return lag;
-    }
-
-    // Booting or in an error loop. Check last active replication status.
-    // This includes sync rules in an ERROR state.
-    const content = await this.storage.getActiveSyncRulesContent();
-    if (content == null) {
-      return undefined;
-    }
-    // Measure the lag from the last resume token's time
-    const lsn = content.last_checkpoint_lsn;
-    if (lsn == null) {
-      return undefined;
-    }
-    const { timestamp } = MongoLSN.fromSerialized(lsn);
-    return Date.now() - timestampToDate(timestamp).getTime();
-  }
 }

package/src/replication/MongoSnapshotQuery.ts

@@ -1,6 +1,7 @@
 import { mongo } from '@powersync/lib-service-mongodb';
 import { ReplicationAssertionError } from '@powersync/lib-services-framework';
 import { bson } from '@powersync/service-core';
+import { getCursorBatchBytes } from './internal-mongodb-utils.js';
 
 /**
  * Performs a collection snapshot query, chunking by ranges of _id.
@@ -21,7 +22,9 @@ export class ChunkedSnapshotQuery implements AsyncDisposable {
     this.batchSize = options.batchSize;
   }
 
-  async nextChunk(): Promise<{ docs: mongo.Document[]; lastKey: Uint8Array } | { docs: []; lastKey: null }> {
+  async nextChunk(): Promise<
+    { docs: mongo.Document[]; lastKey: Uint8Array; bytes: number } | { docs: []; lastKey: null; bytes: 0 }
+  > {
     let cursor = this.lastCursor;
     let newCursor = false;
     if (cursor == null || cursor.closed) {
@@ -52,12 +55,13 @@ export class ChunkedSnapshotQuery implements AsyncDisposable {
       this.lastCursor = null;
       if (newCursor) {
         // We just created a new cursor and it has no results - we have finished the end of the query.
-        return { docs: [], lastKey: null };
+        return { docs: [], lastKey: null, bytes: 0 };
       } else {
         // The cursor may have hit the batch limit - retry
         return this.nextChunk();
       }
     }
+    const bytes = getCursorBatchBytes(cursor);
     const docBatch = cursor.readBufferedDocuments();
     this.lastCursor = cursor;
     if (docBatch.length == 0) {
@@ -65,7 +69,7 @@ export class ChunkedSnapshotQuery implements AsyncDisposable {
     }
     const lastKey = docBatch[docBatch.length - 1]._id;
     this.lastKey = lastKey;
-    return { docs: docBatch, lastKey: bson.serialize({ _id: lastKey }) };
+    return { docs: docBatch, lastKey: bson.serialize({ _id: lastKey }), bytes };
   }
 
   async [Symbol.asyncDispose](): Promise<void> {

package/src/replication/internal-mongodb-utils.ts

@@ -0,0 +1,91 @@
+import { mongo } from '@powersync/lib-service-mongodb';
+import { get } from 'http';
+
+/**
+ * Track bytes read on a change stream.
+ *
+ * This is after decompression, and without TLS overhead.
+ *
+ * This excludes some protocol overhead, but does include per-batch overhead.
+ *
+ * This is built on internal APIs, and may stop working in future driver versions.
+ *
+ * @param add Called once for each batch of data.
+ */
+export function trackChangeStreamBsonBytes(changeStream: mongo.ChangeStream, add: (bytes: number) => void) {
+  let internalChangeStream = changeStream as ChangeStreamWithCursor;
+  let current = internalChangeStream.cursor;
+  let degisterCursor = trackCursor(current, add);
+
+  const refresh = () => {
+    // The cursor may be replaced closed and re-opened (replaced) in various scenarios, such as
+    // after a primary fail-over event.
+    // There is no direct even to track that, but the `resumeTokenChanged` event is a good proxy.
+    // It may be called more often than the cursor is replaced, so we just check whether the cursor changed.
+    // This might miss the init batch, so we may under-count slightly in that case. It is a rare event
+    // and typically a small number of bytes, so it's fine to ignore.
+    const next = internalChangeStream.cursor;
+    if (next !== current) {
+      degisterCursor();
+      current = next;
+      degisterCursor = trackCursor(current, add);
+    }
+  };
+
+  changeStream.on('resumeTokenChanged', refresh);
+
+  // We return this to allow de-registration of the event listeners.
+  // However, these are garbage collected automatically when the stream is closed, so it's not strictly necessary to call this.
+  return () => {
+    changeStream.off('resumeTokenChanged', refresh);
+  };
+}
+
+/**
+ * Get the byte size of the current batch on a cursor.
+ *
+ * Call after hasNext(), before or after readBufferedDocuments().
+ *
+ * This is built on internal APIs, and may stop working in future driver versions.
+ */
+export function getCursorBatchBytes(cursor: mongo.AbstractCursor): number {
+  const documents = (cursor as any).documents as CursorResponse | undefined;
+  return getResponseBytes(documents);
+}
+
+// Define the internal types from the driver.
+// Here we're using them defensively, assuming it may be undefined at any point.
+
+interface CursorResponse {
+  toBytes?(): Uint8Array;
+}
+
+interface ChangeStreamWithCursor extends mongo.ChangeStream {
+  cursor?: mongo.AbstractCursor;
+}
+
+function trackCursor(cursor: mongo.AbstractCursor | undefined, add: (bytes: number) => void) {
+  if (cursor == null) {
+    return () => {};
+  }
+  const countBatch = (response: CursorResponse | undefined) => {
+    const bytes = getResponseBytes(response);
+    if (bytes > 0) {
+      add(bytes);
+    }
+  };
+
+  // The `init` event is emitted for the first batch, and the `more` event is emitted for subsequent batches.
+  cursor.on('init', countBatch);
+  cursor.on('more', countBatch);
+
+  return () => {
+    cursor.off('init', countBatch);
+    cursor.off('more', countBatch);
+  };
+}
+
+function getResponseBytes(response: CursorResponse | undefined): number {
+  const buffer = response?.toBytes?.();
+  return buffer?.byteLength ?? 0;
+}

package/src/replication/replication-index.ts

@@ -2,3 +2,4 @@ export * from './MongoRelation.js';
 export * from './ChangeStream.js';
 export * from './ChangeStreamReplicator.js';
 export * from './ChangeStreamReplicationJob.js';
+export * from './internal-mongodb-utils.js';

package/test/src/change_stream.test.ts

@@ -3,11 +3,14 @@ import { setTimeout } from 'node:timers/promises';
 import { describe, expect, test, vi } from 'vitest';
 
 import { mongo } from '@powersync/lib-service-mongodb';
+import { createWriteCheckpoint } from '@powersync/service-core';
 import { test_utils } from '@powersync/service-core-tests';
 
+import { MongoRouteAPIAdapter } from '@module/api/MongoRouteAPIAdapter.js';
 import { PostImagesOption } from '@module/types/types.js';
 import { ChangeStreamTestContext } from './change_stream_utils.js';
-import { describeWithStorage, StorageVersionTestContext } from './util.js';
+import { describeWithStorage, StorageVersionTestContext, TEST_CONNECTION_OPTIONS } from './util.js';
+import { createCheckpoint, STANDALONE_CHECKPOINT_ID } from '@module/replication/MongoRelation.js';
 
 const BASIC_SYNC_RULES = `
 bucket_definitions:
@@ -368,6 +371,59 @@ bucket_definitions:
     expect(data).toMatchObject([test_utils.putOp('test_data', { id: test_id, description: 'test1' })]);
   });
 
+  test('coalesces standalone checkpoints when backlog is buffered', async () => {
+    await using context = await openContext();
+    await context.updateSyncRules(BASIC_SYNC_RULES);
+    await context.replicateSnapshot();
+    await context.markSnapshotConsistent();
+    await using api = new MongoRouteAPIAdapter({
+      type: 'mongodb',
+      ...TEST_CONNECTION_OPTIONS
+    });
+
+    let commitCount = 0;
+    // This relies on internals to count how often checkpoints are committed
+    context.storage!.registerListener({
+      batchStarted: (batch) => {
+        const originalCommit = batch.commit.bind(batch);
+        batch.commit = async (...args) => {
+          commitCount += 1;
+          return await originalCommit(...args);
+        };
+      }
+    });
+
+    context.startStreaming();
+
+    // Wait until the stream is active and caught up, then start counting from zero.
+    await context.getCheckpoint();
+    commitCount = 0;
+
+    // Create a large number of write checkpoints together.
+    // We could alternatively use createCheckpoint() directly, but this gives more of
+    // an end-to-end test of the checkpointing behavior under load.
+    const checkpointCount = 30;
+    await Promise.all(
+      Array.from({ length: checkpointCount }, (i) =>
+        createWriteCheckpoint({
+          userId: 'test_user',
+          clientId: 'test_client' + i,
+          api,
+          storage: context.factory
+        })
+      )
+    );
+
+    // Wait for the checkpoints to be processed.
+    await context.getCheckpoint();
+
+    // We need at least 1 commit.
+    expect(commitCount).toBeGreaterThan(0);
+    // The previous implementation greated 1 commit per checkpoint, which is bad for performance.
+    // We expect a small number here - typically 2-10, but allow for anything less than the total number of checkpoints.
+    expect(commitCount).toBeLessThan(checkpointCount + 1);
+  });
+
   test('large record', async () => {
     // Test a large update.
 

package/test/src/internal_mongodb_utils.test.ts

@@ -0,0 +1,103 @@
+import { describe, expect, test } from 'vitest';
+
+import { getCursorBatchBytes, trackChangeStreamBsonBytes } from '@module/replication/replication-index.js';
+import { mongo } from '@powersync/lib-service-mongodb';
+import { clearTestDb, connectMongoData } from './util.js';
+
+describe('internal mongodb utils', () => {
+  // The implementation relies on internal APIs, so we verify this works as expected for various types of change streams.
+  test('collection change stream size tracking', async () => {
+    await testChangeStreamBsonBytes('collection');
+  });
+
+  test('db change stream size tracking', async () => {
+    await testChangeStreamBsonBytes('db');
+  });
+
+  test('cluster change stream size tracking', async () => {
+    await testChangeStreamBsonBytes('cluster');
+  });
+
+  test('cursor batch size tracking', async () => {
+    const { db, client } = await connectMongoData();
+    await using _ = { [Symbol.asyncDispose]: async () => await client.close() };
+    await clearTestDb(db);
+    const collection = db.collection('test_data');
+    await collection.insertMany([{ test: 1 }, { test: 2 }, { test: 3 }, { test: 4 }, { test: 5 }]);
+
+    const cursor = collection.find({}, { batchSize: 2 });
+    let batchBytes: number[] = [];
+    let totalBytes = 0;
+    // We use this in the same way as ChunkedSnapshotQuery
+    while (await cursor.hasNext()) {
+      batchBytes.push(getCursorBatchBytes(cursor));
+      totalBytes += batchBytes[batchBytes.length - 1];
+      cursor.readBufferedDocuments();
+    }
+
+    // 3 batches: [2, 2, 1] documents. Should not change
+    expect(batchBytes.length).toEqual(3);
+    // Current tests show 839, but this may change depending on the MongoDB version and other conditions.
+    expect(totalBytes).toBeGreaterThan(400);
+    expect(totalBytes).toBeLessThan(1200);
+  });
+
+  async function testChangeStreamBsonBytes(type: 'db' | 'collection' | 'cluster') {
+    // With MongoDB, replication uses the exact same document format
+    // as normal queries. We test it anyway.
+    const { db, client } = await connectMongoData();
+    await using _ = { [Symbol.asyncDispose]: async () => await client.close() };
+    await clearTestDb(db);
+    const collection = db.collection('test_data');
+
+    let stream: mongo.ChangeStream;
+    if (type === 'collection') {
+      stream = collection.watch([], {
+        maxAwaitTimeMS: 5,
+        fullDocument: 'updateLookup'
+      });
+    } else if (type === 'db') {
+      stream = db.watch([], {
+        maxAwaitTimeMS: 5,
+        fullDocument: 'updateLookup'
+      });
+    } else {
+      stream = client.watch([], {
+        maxAwaitTimeMS: 5,
+        fullDocument: 'updateLookup'
+      });
+    }
+
+    let batchBytes: number[] = [];
+    let totalBytes = 0;
+    trackChangeStreamBsonBytes(stream, (bytes) => {
+      batchBytes.push(bytes);
+      totalBytes += bytes;
+    });
+
+    const readAll = async () => {
+      while ((await stream.tryNext()) != null) {}
+    };
+
+    await readAll();
+
+    await collection.insertOne({ test: 1 });
+    await readAll();
+    await collection.insertOne({ test: 2 });
+    await readAll();
+    await collection.insertOne({ test: 3 });
+    await readAll();
+
+    await stream.close();
+
+    // The exact length by vary based on exact batching logic, but we do want to know when it changes.
+    // Note: If this causes unstable tests, we can relax this check.
+    expect(batchBytes.length).toEqual(8);
+
+    // Current tests show 4464-4576 bytes for the size, depending on the type of change stream.
+    // This can easily vary based on the mongodb version and general conditions, so we just check the general range.
+    // For the most part, if any bytes are reported, the tracking is working.
+    expect(totalBytes).toBeGreaterThan(2000);
+    expect(totalBytes).toBeLessThan(8000);
+  }
+});