@powersync/service-module-mongodb 0.15.4 → 0.17.0

package/src/replication/ChangeStream.ts

@@ -1,8 +1,7 @@
-import { isMongoNetworkTimeoutError, isMongoServerError, mongo } from '@powersync/lib-service-mongodb';
+import { mongo } from '@powersync/lib-service-mongodb';
 import {
   container,
   DatabaseConnectionError,
-  logger as defaultLogger,
   ErrorCode,
   Logger,
   ReplicationAbortedError,
@@ -11,6 +10,7 @@ import {
 } from '@powersync/lib-services-framework';
 import {
   MetricsEngine,
+  PerformanceTracer,
   RelationCache,
   ReplicationLagTracker,
   SaveOperationTag,
@@ -18,29 +18,23 @@ import {
   SourceTable,
   storage
 } from '@powersync/service-core';
-import {
-  DatabaseInputRow,
-  HydratedSyncRules,
-  SqliteInputRow,
-  SqliteRow,
-  TablePattern
-} from '@powersync/service-sync-rules';
+import { HydratedSyncConfig } from '@powersync/service-sync-rules';
 import { ReplicationMetric } from '@powersync/service-types';
+import { performance } from 'node:perf_hooks';
 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 { createCheckpoint, getCacheIdentifier, getMongoRelation, STANDALONE_CHECKPOINT_ID } from './MongoRelation.js';
+import { MongoSnapshotter, MongoSnapshotterHooks } from './MongoSnapshotter.js';
 import {
-  constructAfterRecord,
-  createCheckpoint,
-  getCacheIdentifier,
-  getMongoRelation,
-  STANDALONE_CHECKPOINT_ID
-} from './MongoRelation.js';
-import { ChunkedSnapshotQuery } from './MongoSnapshotQuery.js';
+  ChangeStreamBatch,
+  parseChangeDocument,
+  ProjectedChangeStreamDocument,
+  rawChangeStream
+} from './RawChangeStream.js';
 import { CHECKPOINTS_COLLECTION, timestampToDate } from './replication-utils.js';
-
+import { DirectSourceRowConverter, SourceRowConverter } from './SourceRowConverter.js';
 export interface ChangeStreamOptions {
   connections: MongoManager;
   storage: storage.SyncRulesBucketStorage;
@@ -59,12 +53,10 @@ export interface ChangeStreamOptions {
    */
   snapshotChunkLength?: number;
 
-  logger?: Logger;
-}
+  storageHooks?: storage.StorageHooks;
+  snapshotHooks?: MongoSnapshotterHooks;
 
-interface InitResult {
-  needsInitialSync: boolean;
-  snapshotLsn: string | null;
+  logger?: Logger;
 }
 
 /**
@@ -82,7 +74,7 @@ export class ChangeStreamInvalidatedError extends DatabaseConnectionError {
 }
 
 export class ChangeStream {
-  sync_rules: HydratedSyncRules;
+  sync_rules: HydratedSyncConfig;
   group_id: number;
 
   connection_id = 1;
@@ -96,8 +88,15 @@ export class ChangeStream {
 
   private readonly maxAwaitTimeMS: number;
 
-  private abort_signal: AbortSignal;
+  private abortController = new AbortController();
+  private abortSignal: AbortSignal = this.abortController.signal;
+
+  private initPromise: Promise<void> | null = null;
+  private snapshotter: MongoSnapshotter;
 
+  /**
+   * We use the relationCache _only_ for caching static SourceTable info, not for snapshot status.
+   */
   private relationCache = new RelationCache(getCacheIdentifier);
 
   private replicationLag = new ReplicationLagTracker();
@@ -110,6 +109,10 @@ export class ChangeStream {
 
   private changeStreamTimeout: number;
 
+  private storageHooks: storage.StorageHooks | undefined;
+
+  private readonly sourceRowConverter: SourceRowConverter;
+
   constructor(options: ChangeStreamOptions) {
     this.storage = options.storage;
     this.metrics = options.metrics;
@@ -117,29 +120,40 @@ export class ChangeStream {
     this.connections = options.connections;
     this.maxAwaitTimeMS = options.maxAwaitTimeMS ?? 10_000;
     this.snapshotChunkLength = options.snapshotChunkLength ?? 6_000;
+    this.storageHooks = options.storageHooks;
     this.client = this.connections.client;
     this.defaultDb = this.connections.db;
     this.sync_rules = options.storage.getParsedSyncRules({
       defaultSchema: this.defaultDb.databaseName
     });
+    this.sourceRowConverter = new DirectSourceRowConverter(this.sync_rules.compatibility);
+
     // The change stream aggregation command should timeout before the socket times out,
     // so we use 90% of the socket timeout value.
     this.changeStreamTimeout = Math.ceil(this.client.options.socketTimeoutMS * 0.9);
 
-    this.abort_signal = options.abort_signal;
-    this.abort_signal.addEventListener(
+    this.logger = options.logger ?? this.storage.logger;
+    this.snapshotter = new MongoSnapshotter({
+      ...options,
+      abortSignal: this.abortSignal,
+      logger: this.logger,
+      checkpointStreamId: this.checkpointStreamId
+    });
+
+    options.abort_signal.addEventListener(
       'abort',
       () => {
-        // TODO: Fast abort?
+        this.abortController.abort(options.abort_signal.reason);
       },
       { once: true }
     );
-
-    this.logger = options.logger ?? defaultLogger;
+    if (options.abort_signal.aborted) {
+      this.abortController.abort(options.abort_signal.reason);
+    }
   }
 
   get stopped() {
-    return this.abort_signal.aborted;
+    return this.abortSignal.aborted;
   }
 
   private get usePostImages() {
@@ -150,270 +164,6 @@ export class ChangeStream {
     return this.connections.options.postImages == PostImagesOption.AUTO_CONFIGURE;
   }
 
-  /**
-   * This resolves a pattern, persists the related metadata, and returns
-   * the resulting SourceTables.
-   *
-   * This implicitly checks the collection postImage configuration.
-   */
-  async resolveQualifiedTableNames(
-    batch: storage.BucketStorageBatch,
-    tablePattern: TablePattern
-  ): Promise<storage.SourceTable[]> {
-    const schema = tablePattern.schema;
-    if (tablePattern.connectionTag != this.connections.connectionTag) {
-      return [];
-    }
-
-    let nameFilter: RegExp | string;
-    if (tablePattern.isWildcard) {
-      nameFilter = new RegExp('^' + escapeRegExp(tablePattern.tablePrefix));
-    } else {
-      nameFilter = tablePattern.name;
-    }
-    let result: storage.SourceTable[] = [];
-
-    // Check if the collection exists
-    const collections = await this.client
-      .db(schema)
-      .listCollections(
-        {
-          name: nameFilter
-        },
-        { nameOnly: false }
-      )
-      .toArray();
-
-    if (!tablePattern.isWildcard && collections.length == 0) {
-      this.logger.warn(`Collection ${schema}.${tablePattern.name} not found`);
-    }
-
-    for (let collection of collections) {
-      const table = await this.handleRelation(
-        batch,
-        getMongoRelation({ db: schema, coll: collection.name }),
-        // This is done as part of the initial setup - snapshot is handled elsewhere
-        { snapshot: false, collectionInfo: collection }
-      );
-
-      result.push(table);
-    }
-
-    return result;
-  }
-
-  async initSlot(): Promise<InitResult> {
-    const status = await this.storage.getStatus();
-    if (status.snapshot_done && status.checkpoint_lsn) {
-      this.logger.info(`Initial replication already done`);
-      return { needsInitialSync: false, snapshotLsn: null };
-    }
-
-    return { needsInitialSync: true, snapshotLsn: status.snapshot_lsn };
-  }
-
-  async estimatedCount(table: storage.SourceTable): Promise<string> {
-    const count = await this.estimatedCountNumber(table);
-    return `~${count}`;
-  }
-
-  async estimatedCountNumber(table: storage.SourceTable): Promise<number> {
-    const db = this.client.db(table.schema);
-    return await db.collection(table.name).estimatedDocumentCount();
-  }
-
-  /**
-   * This gets a LSN before starting a snapshot, which we can resume streaming from after the snapshot.
-   *
-   * This LSN can survive initial replication restarts.
-   */
-  private async getSnapshotLsn(): Promise<string> {
-    const hello = await this.defaultDb.command({ hello: 1 });
-    // Basic sanity check
-    if (hello.msg == 'isdbgrid') {
-      throw new ServiceError(
-        ErrorCode.PSYNC_S1341,
-        'Sharded MongoDB Clusters are not supported yet (including MongoDB Serverless instances).'
-      );
-    } else if (hello.setName == null) {
-      throw new ServiceError(
-        ErrorCode.PSYNC_S1342,
-        'Standalone MongoDB instances are not supported - use a replicaset.'
-      );
-    }
-
-    // Open a change stream just to get a resume token for later use.
-    // We could use clusterTime from the hello command, but that won't tell us if the
-    // snapshot isn't valid anymore.
-    // If we just use the first resumeToken from the stream, we get two potential issues:
-    // 1. The resumeToken may just be a wrapped clusterTime, which does not detect changes
-    //    in source db or other stream issues.
-    // 2. The first actual change we get may have the same clusterTime, causing us to incorrect
-    //    skip that event.
-    // Instead, we create a new checkpoint document, and wait until we get that document back in the stream.
-    // To avoid potential race conditions with the checkpoint creation, we create a new checkpoint document
-    // periodically until the timeout is reached.
-
-    const LSN_TIMEOUT_SECONDS = 60;
-    const LSN_CREATE_INTERVAL_SECONDS = 1;
-
-    // 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 = performance.now();
-    let eventsSeen = 0;
-
-    while (performance.now() - startTime < LSN_TIMEOUT_SECONDS * 1000) {
-      if (performance.now() - lastCheckpointCreated >= LSN_CREATE_INTERVAL_SECONDS * 1000) {
-        await createCheckpoint(this.client, this.defaultDb, this.checkpointStreamId);
-        lastCheckpointCreated = performance.now();
-      }
-
-      // tryNext() doesn't block, while next() / hasNext() does block until there is data on the stream
-      const changeDocument = await stream.tryNext().catch((e) => {
-        throw mapChangeStreamError(e);
-      });
-      if (changeDocument == null) {
-        continue;
-      }
-
-      const ns = 'ns' in changeDocument && 'coll' in changeDocument.ns ? changeDocument.ns : undefined;
-
-      if (ns?.coll == CHECKPOINTS_COLLECTION && 'documentKey' in changeDocument) {
-        const checkpointId = changeDocument.documentKey._id as string | mongo.ObjectId;
-        if (!this.checkpointStreamId.equals(checkpointId)) {
-          continue;
-        }
-        const { comparable: lsn } = new MongoLSN({
-          timestamp: changeDocument.clusterTime!,
-          resume_token: changeDocument._id
-        });
-        return lsn;
-      }
-
-      eventsSeen += 1;
-    }
-
-    // Could happen if there is a very large replication lag?
-    throw new ServiceError(
-      ErrorCode.PSYNC_S1301,
-      `Timeout after while waiting for checkpoint document for ${LSN_TIMEOUT_SECONDS}s. Streamed events = ${eventsSeen}`
-    );
-  }
-
-  /**
-   * Given a snapshot LSN, validate that we can read from it, by opening a change stream.
-   */
-  private async validateSnapshotLsn(lsn: string) {
-    await using streamManager = this.openChangeStream({ lsn: lsn, maxAwaitTimeMs: 0 });
-    const { stream } = streamManager;
-    try {
-      // tryNext() doesn't block, while next() / hasNext() does block until there is data on the stream
-      await stream.tryNext();
-    } catch (e) {
-      // Note: A timeout here is not handled as a ChangeStreamInvalidatedError, even though
-      // we possibly cannot recover from it.
-      throw mapChangeStreamError(e);
-    }
-  }
-
-  async initialReplication(snapshotLsn: string | null) {
-    const sourceTables = this.sync_rules.getSourceTables();
-    await this.client.connect();
-
-    const flushResult = await this.storage.startBatch(
-      {
-        logger: this.logger,
-        zeroLSN: MongoLSN.ZERO.comparable,
-        defaultSchema: this.defaultDb.databaseName,
-        storeCurrentData: false,
-        skipExistingRows: true
-      },
-      async (batch) => {
-        if (snapshotLsn == null) {
-          // First replication attempt - get a snapshot and store the timestamp
-          snapshotLsn = await this.getSnapshotLsn();
-          await batch.setResumeLsn(snapshotLsn);
-          this.logger.info(`Marking snapshot at ${snapshotLsn}`);
-        } else {
-          this.logger.info(`Resuming snapshot at ${snapshotLsn}`);
-          // Check that the snapshot is still valid.
-          await this.validateSnapshotLsn(snapshotLsn);
-        }
-
-        // Start by resolving all tables.
-        // This checks postImage configuration, and that should fail as
-        // early as possible.
-        let allSourceTables: SourceTable[] = [];
-        for (let tablePattern of sourceTables) {
-          const tables = await this.resolveQualifiedTableNames(batch, tablePattern);
-          allSourceTables.push(...tables);
-        }
-
-        let tablesWithStatus: SourceTable[] = [];
-        for (let table of allSourceTables) {
-          if (table.snapshotComplete) {
-            this.logger.info(`Skipping ${table.qualifiedName} - snapshot already done`);
-            continue;
-          }
-          let count = await this.estimatedCountNumber(table);
-          const updated = await batch.updateTableProgress(table, {
-            totalEstimatedCount: count
-          });
-          tablesWithStatus.push(updated);
-          this.relationCache.update(updated);
-          this.logger.info(
-            `To replicate: ${table.qualifiedName}: ${updated.snapshotStatus?.replicatedCount}/~${updated.snapshotStatus?.totalEstimatedCount}`
-          );
-        }
-
-        for (let table of tablesWithStatus) {
-          await this.snapshotTable(batch, table);
-          await batch.markTableSnapshotDone([table]);
-
-          this.touch();
-        }
-
-        // The checkpoint here is a marker - we need to replicate up to at least this
-        // point before the data can be considered consistent.
-        // We could do this for each individual table, but may as well just do it once for the entire snapshot.
-        const checkpoint = await createCheckpoint(this.client, this.defaultDb, STANDALONE_CHECKPOINT_ID);
-        await batch.markAllSnapshotDone(checkpoint);
-
-        // This will not create a consistent checkpoint yet, but will persist the op.
-        // Actual checkpoint will be created when streaming replication caught up.
-        await batch.commit(snapshotLsn);
-
-        this.logger.info(`Snapshot done. Need to replicate from ${snapshotLsn} to ${checkpoint} to be consistent`);
-      }
-    );
-    return { lastOpId: flushResult?.flushed_op };
-  }
-
-  private async setupCheckpointsCollection() {
-    const collection = await this.getCollectionInfo(this.defaultDb.databaseName, CHECKPOINTS_COLLECTION);
-    if (collection == null) {
-      await this.defaultDb.createCollection(CHECKPOINTS_COLLECTION, {
-        changeStreamPreAndPostImages: { enabled: true }
-      });
-    } else if (this.usePostImages && collection.options?.changeStreamPreAndPostImages?.enabled != true) {
-      // Drop + create requires less permissions than collMod,
-      // and we don't care about the data in this collection.
-      await this.defaultDb.dropCollection(CHECKPOINTS_COLLECTION);
-      await this.defaultDb.createCollection(CHECKPOINTS_COLLECTION, {
-        changeStreamPreAndPostImages: { enabled: true }
-      });
-    } else {
-      // Clear the collection on startup, to keep it clean
-      // We never query this collection directly, and don't want to keep the data around.
-      // We only use this to get data into the oplog/changestream.
-      await this.defaultDb.collection(CHECKPOINTS_COLLECTION).deleteMany({});
-    }
-  }
-
   private getSourceNamespaceFilters(): { $match: any; multipleDatabases: boolean } {
     const sourceTables = this.sync_rules.getSourceTables();
 
@@ -465,95 +215,14 @@ export class ChangeStream {
     return { $match: nsFilter, multipleDatabases };
   }
 
-  static *getQueryData(results: Iterable<DatabaseInputRow>): Generator<SqliteInputRow> {
-    for (let row of results) {
-      yield constructAfterRecord(row);
-    }
-  }
-
-  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);
-    const collection = db.collection(table.name);
-    await using query = new ChunkedSnapshotQuery({
-      collection,
-      key: table.snapshotStatus?.lastKey,
-      batchSize: this.snapshotChunkLength
-    });
-    if (query.lastKey != null) {
-      this.logger.info(
-        `Replicating ${table.qualifiedName} ${table.formatSnapshotProgress()} - resuming at _id > ${query.lastKey}`
-      );
-    } else {
-      this.logger.info(`Replicating ${table.qualifiedName} ${table.formatSnapshotProgress()}`);
-    }
-
-    let lastBatch = performance.now();
-    let nextChunkPromise = query.nextChunk();
-    while (true) {
-      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);
-      }
-
-      // Pre-fetch next batch, so that we can read and write concurrently
-      nextChunkPromise = query.nextChunk();
-      for (let document of docBatch) {
-        const record = this.constructAfterRecord(document);
-
-        // This auto-flushes when the batch reaches its size limit
-        await batch.save({
-          tag: SaveOperationTag.INSERT,
-          sourceTable: table,
-          before: undefined,
-          beforeReplicaId: undefined,
-          after: record,
-          afterReplicaId: document._id
-        });
-      }
-
-      // Important: flush before marking progress
-      await batch.flush();
-      at += docBatch.length;
-      rowsReplicatedMetric.add(docBatch.length);
-
-      table = await batch.updateTableProgress(table, {
-        lastKey,
-        replicatedCount: at,
-        totalEstimatedCount: totalEstimatedCount
-      });
-      this.relationCache.update(table);
-
-      const duration = performance.now() - lastBatch;
-      lastBatch = performance.now();
-      this.logger.info(
-        `Replicating ${table.qualifiedName} ${table.formatSnapshotProgress()} in ${duration.toFixed(0)}ms`
-      );
-      this.touch();
-    }
-    // In case the loop was interrupted, make sure we await the last promise.
-    await nextChunkPromise;
-  }
-
-  private async getRelation(
+  private async getRelations(
     batch: storage.BucketStorageBatch,
     descriptor: SourceEntityDescriptor,
     options: { snapshot: boolean }
-  ): Promise<SourceTable> {
-    const existing = this.relationCache.get(descriptor);
+  ): Promise<SourceTable[]> {
+    const existing = this.relationCache.getAll(descriptor);
     if (existing != null) {
+      // We do this even when it's an empty result: Empty means nothing to sync, and we don't need to re-resolve.
       return existing;
     }
 
@@ -611,14 +280,11 @@ export class ChangeStream {
     }
 
     const snapshot = options.snapshot;
-    const result = await this.storage.resolveTable({
-      group_id: this.group_id,
+    const result = await batch.resolveTables({
       connection_id: this.connection_id,
-      connection_tag: this.connections.connectionTag,
-      entity_descriptor: descriptor,
-      sync_rules: this.sync_rules
+      source: descriptor
     });
-    this.relationCache.update(result.table);
+    this.relationCache.updateAll(descriptor, result.tables);
 
     // Drop conflicting collections.
     // This is generally not expected for MongoDB source dbs, so we log an error.
@@ -632,47 +298,38 @@ export class ChangeStream {
     // Snapshot if:
     // 1. Snapshot is requested (false for initial snapshot, since that process handles it elsewhere)
     // 2. Snapshot is not already done, AND:
-    // 3. The table is used in sync rules.
-    const shouldSnapshot = snapshot && !result.table.snapshotComplete && result.table.syncAny;
-    if (shouldSnapshot) {
+    // 3. The table is used in sync config.
+    const snapshotCandidates = result.tables.filter((table) => snapshot && !table.snapshotComplete && table.syncAny);
+    if (snapshotCandidates.length > 0) {
       this.logger.info(`New collection: ${descriptor.schema}.${descriptor.name}`);
-      // Truncate this table, in case a previous snapshot was interrupted.
-      await batch.truncate([result.table]);
-
-      await this.snapshotTable(batch, result.table);
-      const no_checkpoint_before_lsn = await createCheckpoint(this.client, this.defaultDb, STANDALONE_CHECKPOINT_ID);
-
-      const [table] = await batch.markTableSnapshotDone([result.table], no_checkpoint_before_lsn);
-      return table;
+      await this.snapshotter.snapshotTables(batch, snapshotCandidates);
     }
 
-    return result.table;
-  }
-
-  private constructAfterRecord(document: mongo.Document): SqliteRow {
-    const inputRow = constructAfterRecord(document);
-    return this.sync_rules.applyRowContext<never>(inputRow);
+    return result.tables;
   }
 
   async writeChange(
     batch: storage.BucketStorageBatch,
     table: storage.SourceTable,
-    change: mongo.ChangeStreamDocument
+    change: ProjectedChangeStreamDocument
   ): Promise<storage.FlushedResult | null> {
     if (!table.syncAny) {
-      this.logger.debug(`Collection ${table.qualifiedName} not used in sync rules - skipping`);
+      this.logger.debug(`Collection ${table.qualifiedName} not used in sync config - skipping`);
       return null;
     }
 
     this.metrics.getCounter(ReplicationMetric.ROWS_REPLICATED).add(1);
     if (change.operationType == 'insert') {
-      const baseRecord = this.constructAfterRecord(change.fullDocument);
+      const { row: baseRecord, replicaId: _replicaId } = this.rawToSqliteRow(change.fullDocument);
       return await batch.save({
         tag: SaveOperationTag.INSERT,
         sourceTable: table,
         before: undefined,
         beforeReplicaId: undefined,
         after: baseRecord,
+        // Same as _replicaId
+        // We specifically need to use the source _id, not the converted one in baseRecord,
+        // to preserve _id uniqueness properties.
         afterReplicaId: change.documentKey._id
       });
     } else if (change.operationType == 'update' || change.operationType == 'replace') {
@@ -685,14 +342,14 @@ export class ChangeStream {
           beforeReplicaId: change.documentKey._id
         });
       }
-      const after = this.constructAfterRecord(change.fullDocument!);
+      const { row: after, replicaId: _replicaId } = this.rawToSqliteRow(change.fullDocument!);
       return await batch.save({
         tag: SaveOperationTag.UPDATE,
         sourceTable: table,
         before: undefined,
         beforeReplicaId: undefined,
         after: after,
-        afterReplicaId: change.documentKey._id
+        afterReplicaId: change.documentKey._id // Same as _replicaId
       });
     } else if (change.operationType == 'delete') {
       return await batch.save({
@@ -707,38 +364,65 @@ export class ChangeStream {
   }
 
   async replicate() {
+    let streamPromise: Promise<void> | null = null;
+    let loopPromise: Promise<void> | null = null;
     try {
       // If anything errors here, the entire replication process is halted, and
       // all connections automatically closed, including this one.
-      await this.initReplication();
-      await this.streamChanges();
+      this.initPromise = this.initReplication();
+      await this.initPromise;
+      loopPromise = this.snapshotter
+        .replicationLoop()
+        .then(() => {
+          throw new ReplicationAssertionError(`Replication snapshotter exited unexpectedly`);
+        })
+        .catch((e) => {
+          this.abortController.abort(e);
+          throw e;
+        });
+      if (!this.snapshotter.supportsConcurrentSnapshots) {
+        await Promise.race([this.snapshotter.waitForInitialSnapshot(), loopPromise]);
+      }
+      streamPromise = this.streamChanges()
+        .then(() => {
+          throw new ReplicationAssertionError(`Replication stream exited unexpectedly`);
+        })
+        .catch((e) => {
+          this.abortController.abort(e);
+          throw e;
+        });
+
+      const results = await Promise.allSettled([loopPromise, streamPromise]);
+      throw replicationLoopError(results);
     } catch (e) {
       await this.storage.reportError(e);
       throw e;
+    } finally {
+      this.abortController.abort();
+    }
+  }
+
+  public async waitForInitialSnapshot() {
+    if (this.initPromise == null) {
+      throw new ReplicationAssertionError('replicate() must be called before waitForInitialSnapshot()');
     }
+    await this.initPromise;
+    await this.snapshotter.waitForInitialSnapshot();
   }
 
-  async initReplication() {
-    const result = await this.initSlot();
-    await this.setupCheckpointsCollection();
+  private async initReplication() {
+    const result = await this.snapshotter.checkSlot();
+    await this.snapshotter.setupCheckpointsCollection();
     if (result.needsInitialSync) {
       if (result.snapshotLsn == null) {
         // Snapshot LSN is not present, so we need to start replication from scratch.
-        await this.storage.clear({ signal: this.abort_signal });
-      }
-      const { lastOpId } = await this.initialReplication(result.snapshotLsn);
-      if (lastOpId != null) {
-        // Populate the cache _after_ initial replication, but _before_ we switch to this sync rules.
-        await this.storage.populatePersistentChecksumCache({
-          signal: this.abort_signal,
-          // No checkpoint yet, but we do have the opId.
-          maxOpId: lastOpId
-        });
+        await this.storage.clear({ signal: this.abortSignal });
       }
+      await this.snapshotter.queueSnapshotTables(result.snapshotLsn);
     }
   }
 
-  async streamChanges() {
+  private async streamChanges() {
     try {
       await this.streamChangesInternal();
     } catch (e) {
@@ -753,19 +437,19 @@ export class ChangeStream {
     }
   }
 
-  private openChangeStream(options: { lsn: string | null; maxAwaitTimeMs?: number }) {
+  private rawChangeStreamBatches(options: {
+    lsn: string | null;
+    maxAwaitTimeMS?: number;
+    batchSize?: number;
+    filters: { $match: any; multipleDatabases: boolean };
+    signal?: AbortSignal;
+    tracer?: PerformanceTracer<'changestream'>;
+  }): AsyncIterableIterator<ChangeStreamBatch> {
     const lastLsn = options.lsn ? MongoLSN.fromSerialized(options.lsn) : null;
     const startAfter = lastLsn?.timestamp;
     const resumeAfter = lastLsn?.resumeToken;
 
-    const filters = this.getSourceNamespaceFilters();
-
-    const pipeline: mongo.Document[] = [
-      {
-        $match: filters.$match
-      },
-      { $changeStreamSplitLargeEvent: {} }
-    ];
+    const filters = options.filters;
 
     let fullDocument: 'required' | 'updateLookup';
 
@@ -777,12 +461,19 @@ export class ChangeStream {
     } else {
       fullDocument = 'updateLookup';
     }
-    const streamOptions: mongo.ChangeStreamOptions = {
+    const streamOptions: mongo.ChangeStreamOptions & mongo.Document = {
       showExpandedEvents: true,
-      maxAwaitTimeMS: options.maxAwaitTimeMs ?? this.maxAwaitTimeMS,
-      fullDocument: fullDocument,
-      maxTimeMS: this.changeStreamTimeout
+      fullDocument: fullDocument
     };
+    const pipeline: mongo.Document[] = [
+      {
+        $changeStream: streamOptions
+      },
+      {
+        $match: filters.$match
+      },
+      { $changeStreamSplitLargeEvent: {} }
+    ];
 
     /**
      * Only one of these options can be supplied at a time.
@@ -796,45 +487,27 @@ export class ChangeStream {
       streamOptions.startAtOperationTime = startAfter;
     }
 
-    let stream: mongo.ChangeStream<mongo.Document>;
+    let watchDb: mongo.Db;
     if (filters.multipleDatabases) {
-      // Requires readAnyDatabase@admin on Atlas
-      stream = this.client.watch(pipeline, streamOptions);
+      watchDb = this.client.db('admin');
+      streamOptions.allChangesForCluster = true;
     } else {
-      // Same general result, but requires less permissions than the above
-      stream = this.defaultDb.watch(pipeline, streamOptions);
+      watchDb = this.defaultDb;
     }
 
-    this.abort_signal.addEventListener('abort', () => {
-      stream.close();
-    });
+    return rawChangeStream(watchDb, pipeline, {
+      batchSize: options.batchSize ?? this.snapshotChunkLength,
+      maxAwaitTimeMS: options.maxAwaitTimeMS ?? this.maxAwaitTimeMS,
+      maxTimeMS: this.changeStreamTimeout,
 
-    return {
-      stream,
-      filters,
-      [Symbol.asyncDispose]: async () => {
-        return stream.close();
-      }
-    };
+      signal: options.signal,
+      logger: this.logger,
+      tracer: options.tracer
+    });
   }
 
-  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();
+  private rawToSqliteRow(row: Buffer) {
+    return this.sourceRowConverter.rawToSqliteRow(row);
   }
 
   async streamChangesInternal() {
@@ -842,13 +515,18 @@ export class ChangeStream {
     const bytesReplicatedMetric = this.metrics.getCounter(ReplicationMetric.DATA_REPLICATED_BYTES);
     const chunksReplicatedMetric = this.metrics.getCounter(ReplicationMetric.CHUNKS_REPLICATED);
 
+    const tracer = new PerformanceTracer<
+      'storage' | 'evaluate' | 'batch' | 'source_checkpoint' | 'changestream' | 'processing'
+    >('MongoDB streaming replication');
     await this.storage.startBatch(
       {
         logger: this.logger,
         zeroLSN: MongoLSN.ZERO.comparable,
         defaultSchema: this.defaultDb.databaseName,
         // We get a complete postimage for every change, so we don't need to store the current data.
-        storeCurrentData: false
+        storeCurrentData: false,
+        hooks: this.storageHooks,
+        tracer
       },
       async (batch) => {
         const { resumeFromLsn } = batch;
@@ -857,6 +535,7 @@ export class ChangeStream {
         }
         const lastLsn = MongoLSN.fromSerialized(resumeFromLsn);
         const startAfter = lastLsn?.timestamp;
+        let outerSpan = tracer.span('batch');
 
         // It is normal for this to be a minute or two old when there is a low volume
         // of ChangeStream events.
@@ -864,16 +543,13 @@ export class ChangeStream {
 
         this.logger.info(`Resume streaming at ${startAfter?.inspect()} / ${lastLsn}  | Token age: ${tokenAgeSeconds}s`);
 
-        await using streamManager = this.openChangeStream({ lsn: resumeFromLsn });
-        const { stream, filters } = streamManager;
-        if (this.abort_signal.aborted) {
-          await stream.close();
-          return;
-        }
-        trackChangeStreamBsonBytes(stream, (bytes) => {
-          bytesReplicatedMetric.add(bytes);
-          // Each of these represent a single response message from MongoDB.
-          chunksReplicatedMetric.add(1);
+        const filters = this.getSourceNamespaceFilters();
+        // This is closed when the for loop below returns/breaks/throws
+        const batchStream = this.rawChangeStreamBatches({
+          lsn: resumeFromLsn,
+          filters,
+          signal: this.abortSignal,
+          tracer
         });
 
         // Always start with a checkpoint.
@@ -885,44 +561,30 @@ export class ChangeStream {
           this.checkpointStreamId
         );
 
-        let splitDocument: mongo.ChangeStreamDocument | null = null;
+        let splitDocument: ProjectedChangeStreamDocument | null = null;
 
         let flexDbNameWorkaroundLogged = false;
-        let changesSinceLastCheckpoint = 0;
 
         let lastEmptyResume = performance.now();
         let lastTxnKey: string | null = null;
 
-        while (true) {
-          if (this.abort_signal.aborted) {
-            break;
-          }
+        for await (let eventBatch of batchStream) {
+          const { events, resumeToken } = eventBatch;
+          using batchSpan = tracer.span('processing');
 
-          const originalChangeDocument = await stream.tryNext().catch((e) => {
-            throw mapChangeStreamError(e);
-          });
-          // The stream was closed, we will only ever receive `null` from it
-          if (!originalChangeDocument && stream.closed) {
-            break;
-          }
-
-          if (this.abort_signal.aborted) {
+          bytesReplicatedMetric.add(eventBatch.byteSize);
+          chunksReplicatedMetric.add(1);
+          if (this.abortSignal.aborted) {
             break;
           }
-
-          if (originalChangeDocument == null) {
-            // We get a new null document after `maxAwaitTimeMS` if there were no other events.
-            // In this case, stream.resumeToken is the resume token associated with the last response.
-            // stream.resumeToken is not updated if stream.tryNext() returns data, while stream.next()
-            // does update it.
-            // From observed behavior, the actual resumeToken changes around once every 10 seconds.
+          this.touch();
+          if (events.length == 0) {
+            // No changes in this batch, but we still want to keep the connection alive.
+            // We do this by persisting a keepalive checkpoint.
             // If we don't update it on empty events, we do keep consistency, but resuming the stream
             // with old tokens may cause connection timeouts.
-            // We throttle this further by only persisting a keepalive once a minute.
-            // We add an additional check for waitForCheckpointLsn == null, to make sure we're not
-            // doing a keepalive in the middle of a transaction.
             if (waitForCheckpointLsn == null && performance.now() - lastEmptyResume > 60_000) {
-              const { comparable: lsn, timestamp } = MongoLSN.fromResumeToken(stream.resumeToken);
+              const { comparable: lsn, timestamp } = MongoLSN.fromResumeToken(resumeToken);
               await batch.keepalive(lsn);
               this.touch();
               lastEmptyResume = performance.now();
@@ -933,229 +595,267 @@ export class ChangeStream {
               );
               this.replicationLag.markStarted();
             }
-            continue;
-          }
 
-          this.touch();
-
-          if (startAfter != null && originalChangeDocument.clusterTime?.lte(startAfter)) {
-            continue;
+            // If we have no changes, we can just persist the keepalive.
+            // This is throttled to once per minute.
+            if (performance.now() - lastEmptyResume < 60_000) {
+              continue;
+            }
           }
 
-          let changeDocument = originalChangeDocument;
-          if (originalChangeDocument?.splitEvent != null) {
-            // Handle split events from $changeStreamSplitLargeEvent.
-            // This is only relevant for very large update operations.
-            const splitEvent = originalChangeDocument?.splitEvent;
+          this.touch();
 
-            if (splitDocument == null) {
-              splitDocument = originalChangeDocument;
-            } else {
-              splitDocument = Object.assign(splitDocument, originalChangeDocument);
+          for (let eventIndex = 0; eventIndex < events.length; eventIndex++) {
+            const rawChangeDocument = events[eventIndex];
+            const originalChangeDocument = parseChangeDocument(rawChangeDocument);
+            if (this.abortSignal.aborted) {
+              break;
             }
 
-            if (splitEvent.fragment == splitEvent.of) {
-              // Got all fragments
-              changeDocument = splitDocument;
-              splitDocument = null;
-            } else {
-              // Wait for more fragments
+            if (startAfter != null && originalChangeDocument.clusterTime?.lte(startAfter)) {
               continue;
             }
-          } else if (splitDocument != null) {
-            // We were waiting for fragments, but got a different event
-            throw new ReplicationAssertionError(`Incomplete splitEvent: ${JSON.stringify(splitDocument.splitEvent)}`);
-          }
 
-          if (
-            !filters.multipleDatabases &&
-            'ns' in changeDocument &&
-            changeDocument.ns.db != this.defaultDb.databaseName &&
-            changeDocument.ns.db.endsWith(`_${this.defaultDb.databaseName}`)
-          ) {
-            // When all of the following conditions are met:
-            // 1. We're replicating from an Atlas Flex instance.
-            // 2. There were changestream events recorded while the PowerSync service is paused.
-            // 3. We're only replicating from a single database.
-            // Then we've observed an ns with for example {db: '67b83e86cd20730f1e766dde_ps'},
-            // instead of the expected {db: 'ps'}.
-            // We correct this.
-            changeDocument.ns.db = this.defaultDb.databaseName;
-
-            if (!flexDbNameWorkaroundLogged) {
-              flexDbNameWorkaroundLogged = true;
-              this.logger.warn(
-                `Incorrect DB name in change stream: ${changeDocument.ns.db}. Changed to ${this.defaultDb.databaseName}.`
-              );
-            }
-          }
+            let changeDocument = originalChangeDocument;
+            if (originalChangeDocument?.splitEvent != null) {
+              // Handle split events from $changeStreamSplitLargeEvent.
+              // This is only relevant for very large update operations.
+              const splitEvent = originalChangeDocument?.splitEvent;
 
-          const ns = 'ns' in changeDocument && 'coll' in changeDocument.ns ? changeDocument.ns : undefined;
-
-          if (ns?.coll == CHECKPOINTS_COLLECTION) {
-            /**
-             * Dropping the database does not provide an `invalidate` event.
-             * We typically would receive `drop` events for the collection which we
-             * would process below.
-             *
-             * However we don't commit the LSN after collections are dropped.
-             * The prevents the `startAfter` or `resumeToken` from advancing past the drop events.
-             * The stream also closes after the drop events.
-             * This causes an infinite loop of processing the collection drop events.
-             *
-             * This check here invalidates the change stream if our `_checkpoints` collection
-             * is dropped. This allows for detecting when the DB is dropped.
-             */
-            if (changeDocument.operationType == 'drop') {
-              throw new ChangeStreamInvalidatedError(
-                'Internal collections have been dropped',
-                new Error('_checkpoints collection was dropped')
-              );
+              if (splitDocument == null) {
+                splitDocument = originalChangeDocument;
+              } else {
+                splitDocument = Object.assign(splitDocument, originalChangeDocument);
+              }
+
+              if (splitEvent.fragment == splitEvent.of) {
+                // Got all fragments
+                changeDocument = splitDocument;
+                splitDocument = null;
+              } else {
+                // Wait for more fragments
+                continue;
+              }
+            } else if (splitDocument != null) {
+              // We were waiting for fragments, but got a different event
+              throw new ReplicationAssertionError(`Incomplete splitEvent: ${JSON.stringify(splitDocument.splitEvent)}`);
             }
 
             if (
-              !(
-                changeDocument.operationType == 'insert' ||
-                changeDocument.operationType == 'update' ||
-                changeDocument.operationType == 'replace'
-              )
+              !filters.multipleDatabases &&
+              'ns' in changeDocument &&
+              changeDocument.ns.db != this.defaultDb.databaseName &&
+              changeDocument.ns.db.endsWith(`_${this.defaultDb.databaseName}`)
             ) {
-              continue;
+              // When all of the following conditions are met:
+              // 1. We're replicating from an Atlas Flex instance.
+              // 2. There were changestream events recorded while the PowerSync service is paused.
+              // 3. We're only replicating from a single database.
+              // Then we've observed an ns with for example {db: '67b83e86cd20730f1e766dde_ps'},
+              // instead of the expected {db: 'ps'}.
+              // We correct this.
+              changeDocument.ns.db = this.defaultDb.databaseName;
+
+              if (!flexDbNameWorkaroundLogged) {
+                flexDbNameWorkaroundLogged = true;
+                this.logger.warn(
+                  `Incorrect DB name in change stream: ${changeDocument.ns.db}. Changed to ${this.defaultDb.databaseName}.`
+                );
+              }
             }
 
-            // We handle two types of checkpoint events:
-            // 1. "Standalone" checkpoints, typically write checkpoints. We want to process these
-            //    immediately, regardless of where they were created.
-            // 2. "Batch" checkpoints for the current stream. This is used as a form of dynamic rate
-            //    limiting of commits, so we specifically want to exclude checkpoints from other streams.
-            //
-            // It may be useful to also throttle commits due to standalone checkpoints in the future.
-            // 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) {
-              // 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) {
+            const ns = 'ns' in changeDocument && 'coll' in changeDocument.ns ? changeDocument.ns : undefined;
+
+            if (ns?.coll == CHECKPOINTS_COLLECTION) {
+              /**
+               * Dropping the database does not provide an `invalidate` event.
+               * We typically would receive `drop` events for the collection which we
+               * would process below.
+               *
+               * However we don't commit the LSN after collections are dropped.
+               * This prevents the `startAfter` or `resumeToken` from advancing past the drop events.
+               * The stream also closes after the drop events.
+               * This causes an infinite loop of processing the collection drop events.
+               *
+               * This check here invalidates the change stream if our `_powersync_checkpoints` collection
+               * is dropped. This allows for detecting when the DB is dropped.
+               */
+              if (changeDocument.operationType == 'drop') {
+                throw new ChangeStreamInvalidatedError(
+                  'Internal collections have been dropped',
+                  new Error('_powersync_checkpoints collection was dropped')
+                );
+              }
+
+              if (
+                !(
+                  changeDocument.operationType == 'insert' ||
+                  changeDocument.operationType == 'update' ||
+                  changeDocument.operationType == 'replace'
+                )
+              ) {
+                continue;
+              }
+
+              // We handle two types of checkpoint events:
+              // 1. "Standalone" checkpoints, typically write checkpoints. We want to process these
+              //    immediately, regardless of where they were created.
+              // 2. "Batch" checkpoints for the current stream. This is used as a form of dynamic rate
+              //    limiting of commits, so we specifically want to exclude checkpoints from other streams.
+              //
+              // It may be useful to also throttle commits due to standalone checkpoints in the future.
+              // 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) {
+                // 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.
+                const hasBufferedChanges = eventIndex < events.length - 1;
+                if (hasBufferedChanges && waitForCheckpointLsn == null) {
+                  // Buffered changes - create a new batch checkpoint to rate limit commits
+                  using _ = tracer.span('source_checkpoint');
                   waitForCheckpointLsn = await createCheckpoint(this.client, this.defaultDb, this.checkpointStreamId);
+                  continue;
+                } else if (waitForCheckpointLsn != null) {
+                  // Skip this checkpoint - wait for the batch checkpoint.
+                  continue;
+                } else {
+                  // No buffered changes, and no batch checkpoint pending - commit immediately.
                 }
+              } else if (!this.checkpointStreamId.equals(checkpointId)) {
                 continue;
               }
-            } else if (!this.checkpointStreamId.equals(checkpointId)) {
-              continue;
-            }
-            const { comparable: lsn } = new MongoLSN({
-              timestamp: changeDocument.clusterTime!,
-              resume_token: changeDocument._id
-            });
-            if (batch.lastCheckpointLsn != null && lsn < batch.lastCheckpointLsn) {
-              // Checkpoint out of order - should never happen with MongoDB.
-              // If it does happen, we throw an error to stop the replication - restarting should recover.
-              // Since we use batch.lastCheckpointLsn for the next resumeAfter, this should not result in an infinite loop.
-              // Originally a workaround for https://jira.mongodb.org/browse/NODE-7042.
-              // This has been fixed in the driver in the meantime, but we still keep this as a safety-check.
-              throw new ReplicationAssertionError(
-                `Change resumeToken ${(changeDocument._id as any)._data} (${timestampToDate(changeDocument.clusterTime!).toISOString()}) is less than last checkpoint LSN ${batch.lastCheckpointLsn}. Restarting replication.`
-              );
-            }
-
-            if (waitForCheckpointLsn != null && lsn >= waitForCheckpointLsn) {
-              waitForCheckpointLsn = null;
-            }
-            const { checkpointBlocked } = await batch.commit(lsn, {
-              oldestUncommittedChange: this.replicationLag.oldestUncommittedChange
-            });
+              const { comparable: lsn } = new MongoLSN({
+                timestamp: changeDocument.clusterTime!,
+                resume_token: changeDocument._id
+              });
+              if (batch.lastCheckpointLsn != null && lsn < batch.lastCheckpointLsn) {
+                // Checkpoint out of order - should never happen with MongoDB.
+                // If it does happen, we throw an error to stop the replication - restarting should recover.
+                // Since we use batch.lastCheckpointLsn for the next resumeAfter, this should not result in an infinite loop.
+                // Originally a workaround for https://jira.mongodb.org/browse/NODE-7042.
+                // This has been fixed in the driver in the meantime, but we still keep this as a safety-check.
+                throw new ReplicationAssertionError(
+                  `Change resumeToken ${(changeDocument._id as any)._data} (${timestampToDate(changeDocument.clusterTime!).toISOString()}) is less than last checkpoint LSN ${batch.lastCheckpointLsn}. Restarting replication.`
+                );
+              }
 
-            if (!checkpointBlocked) {
-              this.replicationLag.markCommitted();
-              changesSinceLastCheckpoint = 0;
-            }
-          } else if (
-            changeDocument.operationType == 'insert' ||
-            changeDocument.operationType == 'update' ||
-            changeDocument.operationType == 'replace' ||
-            changeDocument.operationType == 'delete'
-          ) {
-            if (waitForCheckpointLsn == null) {
-              waitForCheckpointLsn = await createCheckpoint(this.client, this.defaultDb, this.checkpointStreamId);
-            }
+              if (waitForCheckpointLsn != null && lsn >= waitForCheckpointLsn) {
+                waitForCheckpointLsn = null;
+              }
+              const { checkpointBlocked } = await batch.commit(lsn, {
+                oldestUncommittedChange: this.replicationLag.oldestUncommittedChange
+              });
 
-            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
-              // for whatever reason, then we do need to snapshot it.
-              // This may result in some duplicate operations when a collection is created for the first time after
-              // sync rules was deployed.
-              snapshot: true
-            });
-            if (table.syncAny) {
-              this.replicationLag.trackUncommittedChange(
-                changeDocument.clusterTime == null ? null : timestampToDate(changeDocument.clusterTime)
-              );
+              if (!checkpointBlocked) {
+                this.replicationLag.markCommitted();
+              }
+            } else if (
+              changeDocument.operationType == 'insert' ||
+              changeDocument.operationType == 'update' ||
+              changeDocument.operationType == 'replace' ||
+              changeDocument.operationType == 'delete'
+            ) {
+              if (waitForCheckpointLsn == null) {
+                using _ = tracer.span('source_checkpoint');
+                waitForCheckpointLsn = await createCheckpoint(this.client, this.defaultDb, this.checkpointStreamId);
+              }
 
-              const transactionKeyValue = transactionKey(changeDocument);
+              const rel = getMongoRelation(changeDocument.ns, this.connections.connectionTag);
+              const tables = await this.getRelations(batch, rel, {
+                // In most cases, we should not need to snapshot this. But if this is the first time we see the collection
+                // for whatever reason, then we do need to snapshot it.
+                // This may result in some duplicate operations when a collection is created for the first time after
+                // sync config was deployed.
+                snapshot: true
+              });
+              const tablesToReplicate = tables.filter((table) => table.syncAny);
+              if (tablesToReplicate.length > 0) {
+                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);
+                }
 
-              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);
+                for (const table of tablesToReplicate) {
+                  await this.writeChange(batch, table, changeDocument);
+                }
               }
-
-              const flushResult = await this.writeChange(batch, table, changeDocument);
-              changesSinceLastCheckpoint += 1;
-              if (flushResult != null && changesSinceLastCheckpoint >= 20_000) {
-                // When we are catching up replication after an initial snapshot, there may be a very long delay
-                // before we do a commit(). In that case, we need to periodically persist the resume LSN, so
-                // we don't restart from scratch if we restart replication.
-                // The same could apply if we need to catch up on replication after some downtime.
-                const { comparable: lsn } = new MongoLSN({
-                  timestamp: changeDocument.clusterTime!,
-                  resume_token: changeDocument._id
-                });
-                this.logger.info(`Updating resume LSN to ${lsn} after ${changesSinceLastCheckpoint} changes`);
-                await batch.setResumeLsn(lsn);
-                changesSinceLastCheckpoint = 0;
+            } else if (changeDocument.operationType == 'drop') {
+              const rel = getMongoRelation(changeDocument.ns, this.connections.connectionTag);
+              const tables = await this.getRelations(batch, rel, {
+                // We're "dropping" this collection, so never snapshot it.
+                snapshot: false
+              });
+              const tablesToDrop = tables.filter((table) => table.syncAny);
+              if (tablesToDrop.length > 0) {
+                await batch.drop(tablesToDrop);
+              }
+              this.relationCache.delete(rel);
+            } else if (changeDocument.operationType == 'rename') {
+              const relFrom = getMongoRelation(changeDocument.ns, this.connections.connectionTag);
+              const relTo = getMongoRelation(changeDocument.to, this.connections.connectionTag);
+              const tablesFrom = await this.getRelations(batch, relFrom, {
+                // We're "dropping" this collection, so never snapshot it.
+                snapshot: false
+              });
+              const tablesToDrop = tablesFrom.filter((table) => table.syncAny);
+              if (tablesToDrop.length > 0) {
+                await batch.drop(tablesToDrop);
               }
-            }
-          } else if (changeDocument.operationType == 'drop') {
-            const rel = getMongoRelation(changeDocument.ns);
-            const table = await this.getRelation(batch, rel, {
-              // We're "dropping" this collection, so never snapshot it.
-              snapshot: false
-            });
-            if (table.syncAny) {
-              await batch.drop([table]);
-              this.relationCache.delete(table);
-            }
-          } else if (changeDocument.operationType == 'rename') {
-            const relFrom = getMongoRelation(changeDocument.ns);
-            const relTo = getMongoRelation(changeDocument.to);
-            const tableFrom = await this.getRelation(batch, relFrom, {
-              // We're "dropping" this collection, so never snapshot it.
-              snapshot: false
-            });
-            if (tableFrom.syncAny) {
-              await batch.drop([tableFrom]);
               this.relationCache.delete(relFrom);
+              // Here we do need to snapshot the new table
+              const collection = await this.getCollectionInfo(relTo.schema, relTo.name);
+              await this.handleRelation(batch, relTo, {
+                // This is a new (renamed) collection, so always snapshot it.
+                snapshot: true,
+                collectionInfo: collection
+              });
             }
-            // Here we do need to snapshot the new table
-            const collection = await this.getCollectionInfo(relTo.schema, relTo.name);
-            await this.handleRelation(batch, relTo, {
-              // This is a new (renamed) collection, so always snapshot it.
-              snapshot: true,
-              collectionInfo: collection
-            });
           }
+
+          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);
+          }
+
+          batchSpan.end();
+          const durationsMicroseconds = outerSpan.end();
+          const duration = batchSpan.durationMillis;
+
+          this.logger.info(
+            `Processed batch of ${events.length} changes / ${eventBatch.byteSize} bytes in ${duration}ms`,
+            {
+              count: events.length,
+              bytes: eventBatch.byteSize,
+              duration,
+              t: durationsMicroseconds
+            }
+          );
+          outerSpan = tracer.span('batch');
         }
       }
     );
+
+    throw new ReplicationAbortedError(`Replication stream aborted`, this.abortSignal.reason);
   }
 
   getReplicationLagMillis(): number | undefined {
@@ -1175,34 +875,35 @@ export class ChangeStream {
   }
 }
 
-function mapChangeStreamError(e: any) {
-  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
-    throw new DatabaseConnectionError(ErrorCode.PSYNC_S1345, `Timeout while reading MongoDB ChangeStream`, e);
-  } else if (
-    isMongoServerError(e) &&
-    e.codeName == 'NoMatchingDocument' &&
-    e.errmsg?.includes('post-image was not found')
-  ) {
-    throw new ChangeStreamInvalidatedError(e.errmsg, e);
-  } else if (isMongoServerError(e) && e.hasErrorLabel('NonResumableChangeStreamError')) {
-    throw new ChangeStreamInvalidatedError(e.message, e);
-  } else {
-    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 {
+function transactionKey(doc: Pick<mongo.ChangeStreamDocument, 'lsid' | 'txnNumber'>): string | null {
   if (doc.txnNumber == null || doc.lsid == null) {
     return null;
   }
   return `${doc.lsid.id.toString('hex')}:${doc.txnNumber}`;
 }
+
+/**
+ * Prioritize errors that are _not_ ReplicationAbortedError. Any error on either loopPromise or
+ * streamPromise aborts the other one, which then results in a ReplicationAbortedError, hiding the
+ * original cause.
+ */
+function replicationLoopError(results: PromiseSettledResult<any>[]): unknown {
+  // 1. Prioritize not ReplicationAbortedError.
+  for (const result of results) {
+    if (result.status == 'rejected' && !(result.reason instanceof ReplicationAbortedError)) {
+      return result.reason;
+    }
+  }
+  // 2. Fallback to ReplicationAbortedError.
+  for (const result of results) {
+    if (result.status == 'rejected') {
+      // At this point only ReplicationAbortedError remains
+      return result.reason;
+    }
+  }
+  // 3. Should never happen, but we cover this case.
+  return new ReplicationAssertionError(`Replication loop exited unexpectedly`);
+}