@powersync/service-module-mongodb 0.15.0 → 0.15.2

package/src/replication/ChangeStream.ts

@@ -38,6 +38,7 @@ import {
 } from './MongoRelation.js';
 import { ChunkedSnapshotQuery } from './MongoSnapshotQuery.js';
 import { CHECKPOINTS_COLLECTION, timestampToDate } from './replication-utils.js';
+import { trackChangeStreamBsonBytes } from './internal-mongodb-utils.js';
 
 export interface ChangeStreamOptions {
   connections: MongoManager;
@@ -264,10 +265,13 @@ export class ChangeStream {
     const LSN_TIMEOUT_SECONDS = 60;
     const LSN_CREATE_INTERVAL_SECONDS = 1;
 
-    await using streamManager = this.openChangeStream({ lsn: null, maxAwaitTimeMs: 0 });
+    // Create a checkpoint, and open a change stream using startAtOperationTime with the checkpoint's operationTime.
+    const firstCheckpointLsn = await createCheckpoint(this.client, this.defaultDb, this.checkpointStreamId);
+    await using streamManager = this.openChangeStream({ lsn: firstCheckpointLsn, maxAwaitTimeMs: 0 });
+
     const { stream } = streamManager;
     const startTime = performance.now();
-    let lastCheckpointCreated = -10_000;
+    let lastCheckpointCreated = performance.now();
     let eventsSeen = 0;
 
     while (performance.now() - startTime < LSN_TIMEOUT_SECONDS * 1000) {
@@ -476,6 +480,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);
@@ -496,11 +504,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);
@@ -525,7 +535,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,
@@ -790,6 +800,7 @@ export class ChangeStream {
     } else {
       // Legacy: We don't persist lsns without resumeTokens anymore, but we do still handle the
       // case if we have an old one.
+      // This is also relevant for getSnapshotLSN().
       streamOptions.startAtOperationTime = startAfter;
     }
 
@@ -815,7 +826,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,
@@ -844,6 +878,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
@@ -860,6 +899,7 @@ export class ChangeStream {
         let changesSinceLastCheckpoint = 0;
 
         let lastEmptyResume = performance.now();
+        let lastTxnKey: string | null = null;
 
         while (true) {
           if (this.abort_signal.aborted) {
@@ -1001,7 +1041,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({
@@ -1040,6 +1093,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
@@ -1052,6 +1106,17 @@ export class ChangeStream {
               if (this.oldestUncommittedChange == null && changeDocument.clusterTime != null) {
                 this.oldestUncommittedChange = 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) {
@@ -1149,3 +1214,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/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/change_stream_utils.ts

@@ -164,9 +164,9 @@ export class ChangeStreamTestContext {
   async markSnapshotConsistent() {
     const checkpoint = await createCheckpoint(this.client, this.db, STANDALONE_CHECKPOINT_ID);
 
-    await this.storage!.startBatch(test_utils.BATCH_OPTIONS, async (batch) => {
-      await batch.keepalive(checkpoint);
-    });
+    await using writer = await this.storage!.createWriter(test_utils.BATCH_OPTIONS);
+    await writer.keepalive(checkpoint);
+    await writer.flush();
   }
 
   startStreaming() {

package/test/src/env.ts

@@ -3,7 +3,7 @@ import { utils } from '@powersync/lib-services-framework';
 export const env = utils.collectEnvironmentVariables({
   MONGO_TEST_URL: utils.type.string.default('mongodb://localhost:27017/powersync_test'),
   MONGO_TEST_DATA_URL: utils.type.string.default('mongodb://localhost:27017/powersync_test_data'),
-  PG_STORAGE_TEST_URL: utils.type.string.default('postgres://postgres:postgres@localhost:5431/powersync_storage_test'),
+  PG_STORAGE_TEST_URL: utils.type.string.default('postgres://postgres:postgres@localhost:5432/powersync_storage_test'),
   CI: utils.type.boolean.default('false'),
   SLOW_TESTS: utils.type.boolean.default('false'),
   TEST_MONGO_STORAGE: utils.type.boolean.default('true'),

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);
+  }
+});