[MongoDB] Adaptive batchSize to improve change stream timeout handling (#609)

* Use aggregate batchSize: 1 to reduce timeouts.

* Fix resumeToken handling.

* Handle edge cases.

* Add more low-level change stream tests.

* enableTestCommands=1

* Skip test locally if configureFailPoint is not supported.

* Always flush after a batch.

* Implement adaptive batch sizing.

* Disable test for "resuming with a different source database".

* Changeset.

* Log replication lag on each flush.

* Use constant and update docs.

Author: rkistner

.changeset/tasty-bottles-relate.md

@@ -0,0 +1,5 @@
+---
+'@powersync/service-module-mongodb': patch
+---
+
+Use adaptive batchSize for better recovery in PSYNC_S1345 change stream timeouts.

.github/actions/start-mongodb/action.yml

@@ -1,5 +1,5 @@
 name: Start MongoDB
-description: Pre-pull and start MongoDB with a replica set for tests
+description: Start MongoDB with a replica set for tests
 
 inputs:
   mongodb-version:
@@ -10,14 +10,53 @@ inputs:
 runs:
   using: composite
   steps:
-    # The mongodb-github-action below doesn't use the Docker credentials for the pull.
-    # We pre-pull, so that the image is cached.
-    - name: Pre-pull Mongo image
+    # Previously we used supercharge/mongodb-github-action@1.12.1.
+    # However, we specifically want `--setParameter enableTestCommands=1`, which that doesn't support.
+    # This now just starts MongoDB directly using Docker, and waits until it is ready.
+    - name: Start MongoDB
       shell: bash
-      run: docker pull mongo:${{ inputs.mongodb-version }}
+      run: |
+        set -euo pipefail
 
-    - name: Start MongoDB
-      uses: supercharge/mongodb-github-action@1.12.1
-      with:
-        mongodb-version: ${{ inputs.mongodb-version }}
-        mongodb-replica-set: test-rs
+        container_name="mongodb"
+        image="mongo:${{ inputs.mongodb-version }}"
+
+        docker rm -f "$container_name" >/dev/null 2>&1 || true
+
+        docker run \
+          --name "$container_name" \
+          --publish 27017:27017 \
+          --detach \
+          "$image" \
+          mongod \
+          --port 27017 \
+          --replSet test-rs \
+          --bind_ip_all \
+          --setParameter enableTestCommands=1
+
+        for i in $(seq 1 30); do
+          if docker exec "$container_name" mongosh --quiet --port 27017 --eval "db.adminCommand({ ping: 1 })" >/dev/null; then
+            break
+          fi
+          sleep 1
+          if [ "$i" -eq 30 ]; then
+            echo "MongoDB did not become ready in time"
+            docker logs "$container_name"
+            exit 1
+          fi
+        done
+
+        docker exec "$container_name" mongosh --quiet --port 27017 --eval \
+          'rs.initiate({_id:"test-rs", members:[{_id:0, host:"localhost:27017"}]})'
+
+        for i in $(seq 1 30); do
+          if docker exec "$container_name" mongosh --quiet --port 27017 --eval 'quit(db.hello().isWritablePrimary ? 0 : 1)' >/dev/null; then
+            break
+          fi
+          sleep 1
+          if [ "$i" -eq 30 ]; then
+            echo "MongoDB replica set did not become primary in time"
+            docker logs "$container_name"
+            exit 1
+          fi
+        done

modules/module-mongodb/src/replication/ChangeStream.ts

@@ -1129,6 +1129,17 @@ export class ChangeStream {
               });
             }
           }
+
+          if (splitDocument == null) {
+            // We flush and mark progress on every batch of data we receive.
+            // Batches are generally large (64MB or 6000 events, whichever comes first),
+            // so this is a good natural point to flush and mark progress.
+            // We avoid this when splitDocument is set, since we cannot resume in the middle of a split event.
+            const { comparable: lsn } = MongoLSN.fromResumeToken(resumeToken);
+            await batch.flush({ oldestUncommittedChange: this.replicationLag.oldestUncommittedChange });
+            // TODO: We should consider making this standard behavior of flush().
+            await batch.setResumeLsn(lsn);
+          }
           this.logger.info(`Processed batch of ${events.length} changes in ${Date.now() - batchStart}ms`);
         }
       }

modules/module-mongodb/src/replication/RawChangeStream.ts

@@ -9,14 +9,26 @@ import { ChangeStreamInvalidatedError } from './ChangeStream.js';
 
 export interface RawChangeStreamOptions {
   signal?: AbortSignal;
+
   /**
    * How long to wait for new data per batch (max time for long-polling).
+   *
+   * This is used for maxTimeMS for the getMore command.
    */
   maxAwaitTimeMS: number;
+
   /**
    * Timeout for the initial aggregate command.
    */
   maxTimeMS: number;
+
+  /**
+   * batchSize for the getMore commands.
+   *
+   * The aggregate command always uses a batchSize of 1.
+   *
+   * After a timeout error, the batchSize will be reduced and ramped up again.
+   */
   batchSize: number;
 
   /**
@@ -63,6 +75,8 @@ export async function* rawChangeStream(db: mongo.Db, pipeline: mongo.Document[],
   }
   let lastResumeToken: unknown | null = null;
 
+  const batchSizer = new AdaptiveBatchSize(options.batchSize);
+
   // > If the server supports sessions, the resume attempt MUST use the same session as the previous attempt's command.
   const session = db.client.startSession();
   await using _ = { [Symbol.asyncDispose]: () => session.endSession() };
@@ -79,9 +93,15 @@ export async function* rawChangeStream(db: mongo.Db, pipeline: mongo.Document[],
         changeStreamStage.resumeAfter = lastResumeToken;
         innerPipeline = [{ $changeStream: changeStreamStage }, ...rest];
       }
-      const inner = rawChangeStreamInner(session, db, innerPipeline, {
-        ...options
-      });
+      const inner = rawChangeStreamInner(
+        session,
+        db,
+        innerPipeline,
+        {
+          ...options
+        },
+        batchSizer
+      );
       for await (let batch of inner) {
         yield batch;
         lastResumeToken = batch.resumeToken;
@@ -111,7 +131,8 @@ async function* rawChangeStreamInner(
   session: mongo.ClientSession,
   db: mongo.Db,
   pipeline: mongo.Document[],
-  options: RawChangeStreamOptions
+  options: RawChangeStreamOptions,
+  batchSizer: AdaptiveBatchSize
 ): AsyncGenerator<ChangeStreamBatch> {
   // See specs:
   // https://github.com/mongodb/specifications/blob/master/source/change-streams/change-streams.md
@@ -123,7 +144,6 @@ async function* rawChangeStreamInner(
    */
   let nsCollection: string | null = null;
 
-  const maxTimeMS = options.maxAwaitTimeMS;
   const batchSize = options.batchSize;
   const collection = options.collection ?? 1;
   let abortPromise: Promise<any> | null = null;
@@ -149,7 +169,11 @@ async function* rawChangeStreamInner(
           {
             aggregate: collection,
             pipeline,
-            cursor: { batchSize },
+            cursor: {
+              // Always use batchSize of 1 for the initial aggregate command,
+              // to maximize the chance of success in case of timeouts.
+              batchSize: 1
+            },
             maxTimeMS: options.maxTimeMS
           },
           { session, raw: true }
@@ -183,8 +207,8 @@ async function* rawChangeStreamInner(
           {
             getMore: cursorId,
             collection: nsCollection,
-            batchSize,
-            maxTimeMS
+            batchSize: batchSizer.next(),
+            maxTimeMS: options.maxAwaitTimeMS
           },
           { session, raw: true }
         )
@@ -196,6 +220,9 @@ async function* rawChangeStreamInner(
           }
 
           if (isResumableChangeStreamError(e)) {
+            if (isTimeoutError(e)) {
+              batchSizer.reduceAfterError();
+            }
             throw new ResumableChangeStreamError(e.message, { cause: e });
           }
           throw mapChangeStreamError(e);
@@ -233,6 +260,55 @@ async function* rawChangeStreamInner(
 
 class ResumableChangeStreamError extends Error {}
 
+/**
+ * After a timeout error, we reduce the batch size to this number, then multiply by this for each batch.
+ *
+ * Must be an integer >= 2 with the current implementation.
+ */
+const BATCH_SIZE_MULTIPLIER = 2;
+
+/**
+ * Manage batch sizes after timeout errors.
+ *
+ * This starts with the initial batch size.
+ *
+ * After a timeout error, we reduce the batch size for aggregate command to 1,
+ * then multiply by BATCH_SIZE_MULTIPLIER for each subsequent batch, until we reach the initial batch size again.
+ *
+ * We use this to protect against timeout errors:
+ *   [PSYNC_S1345] Timeout while reading MongoDB ChangeStream
+ *
+ * When we run into that, the stream is restarted automatically. starting with an aggregate command
+ * with batchSize: 1.
+ *
+ * We then ramp up the batchSize for getMore commands.
+ */
+class AdaptiveBatchSize {
+  private nextBatchSize: number;
+
+  constructor(private maxBatchSize: number) {
+    this.nextBatchSize = maxBatchSize;
+  }
+
+  /**
+   * Get the next batchSize for a getMore command.
+   */
+  next() {
+    const current = this.nextBatchSize;
+    this.nextBatchSize = Math.min(this.maxBatchSize, this.nextBatchSize * BATCH_SIZE_MULTIPLIER);
+    return current;
+  }
+
+  /**
+   * After a timeout error, the next aggregate command will start with a batchSize of 1.
+   *
+   * The next getMore will then have a batchSize of BATCH_SIZE_MULTIPLIER.
+   */
+  reduceAfterError() {
+    this.nextBatchSize = BATCH_SIZE_MULTIPLIER;
+  }
+}
+
 type RawBsonValue =
   | string
   | number
@@ -308,31 +384,39 @@ export function parseChangeDocument(buffer: Buffer): ProjectedChangeStreamDocume
   return doc as any;
 }
 
+function isTimeoutError(e: unknown) {
+  return isMongoNetworkTimeoutError(e) || (isMongoServerError(e) && e.codeName == 'MaxTimeMSExpired');
+}
+
 function isResumableChangeStreamError(e: unknown) {
   // See: https://github.com/mongodb/specifications/blob/master/source/change-streams/change-streams.md#resumable-error
   if (!isMongoServerError(e)) {
-    // Any error encountered which is not a server error (e.g. a timeout error or network error)
+    // Any error encountered which is not a server error (e.g. a socket timeout error or network error)
     return true;
   } else if (e.codeName == 'CursorNotFound') {
     // A server error with code 43 (CursorNotFound)
     return true;
   } else if (e.hasErrorLabel('ResumableChangeStreamError')) {
     // For servers with wire version 9 or higher (server version 4.4 or higher), any server error with the ResumableChangeStreamError error label.
     return true;
+  } else if (e.codeName == 'MaxTimeMSExpired') {
+    // Our own exception for MaxTimeMSExpired.
+    // This can help us retry faster, with a smaller batch size (if initialBatchSize is set to 1), which should hopefully avoid the timeout.
+    return true;
   } else {
-    // We ignore servers with wire version less than 9, since we only support MongoDB 6.0+.
+    // Other errors are not retried.
+    // We ignore the spec for servers with wire version less than 9, since we only support MongoDB 6.0+.
     return false;
   }
 }
 
 export function mapChangeStreamError(e: unknown) {
-  if (isMongoNetworkTimeoutError(e)) {
-    // This typically has an unhelpful message like "connection 2 to 159.41.94.47:27017 timed out".
-    // We wrap the error to make it more useful.
-    throw new DatabaseConnectionError(ErrorCode.PSYNC_S1345, `Timeout while reading MongoDB ChangeStream`, e);
-  } else if (isMongoServerError(e) && e.codeName == 'MaxTimeMSExpired') {
-    // maxTimeMS was reached. Example message:
-    // MongoServerError: Executor error during aggregate command on namespace: powersync_test_data.$cmd.aggregate :: caused by :: operation exceeded time limit
+  if (isTimeoutError(e)) {
+    // For isMongoNetworkTimeoutError():
+    //   This typically has an unhelpful message like "connection 2 to 159.41.94.47:27017 timed out".
+    //   We wrap the error to make it more useful.
+    // Example for MaxTimeMSExpired:
+    //   MongoServerError: Executor error during aggregate command on namespace: powersync_test_data.$cmd.aggregate :: caused by :: operation exceeded time limit
     throw new DatabaseConnectionError(ErrorCode.PSYNC_S1345, `Timeout while reading MongoDB ChangeStream`, e);
   } else if (
     isMongoServerError(e) &&