MongoDB replication tracing

.changeset/nine-oranges-destroy.md

@@ -0,0 +1,8 @@
+---
+'@powersync/service-module-postgres-storage': patch
+'@powersync/service-module-mongodb-storage': patch
+'@powersync/service-module-mongodb': patch
+'@powersync/service-core': patch
+---
+
+[MongoDB] Log replication timing info per batch

modules/module-mongodb-storage/src/storage/implementation/MongoBucketBatch.ts

@@ -18,6 +18,7 @@ import {
   deserializeBson,
   InternalOpId,
   isCompleteRow,
+  PerformanceTracer,
   SaveOperationTag,
   storage,
   SyncRuleState,
@@ -60,6 +61,7 @@ export interface MongoBucketBatchOptions {
   markRecordUnavailable: BucketStorageMarkRecordUnavailable | undefined;
 
   logger?: Logger;
+  tracer?: PerformanceTracer<'storage' | 'evaluate'>;
 }
 
 export abstract class MongoBucketBatch
@@ -85,6 +87,8 @@ export abstract class MongoBucketBatch
   private markRecordUnavailable: BucketStorageMarkRecordUnavailable | undefined;
   private clearedError = false;
 
+  private tracer: PerformanceTracer<'storage' | 'evaluate'>;
+
   /**
    * Last LSN received associated with a checkpoint.
    *
@@ -133,6 +137,7 @@ export abstract class MongoBucketBatch
     this.batch = new OperationBatch();
 
     this.persisted_op = options.keepaliveOp ?? null;
+    this.tracer = options.tracer ?? new PerformanceTracer('MongoDB storage');
   }
 
   addCustomWriteCheckpoint(checkpoint: storage.BatchedCustomWriteCheckpointOptions): void {
@@ -170,6 +175,8 @@ export abstract class MongoBucketBatch
     let last_op: InternalOpId | null = null;
     let resumeBatch: OperationBatch | null = null;
 
+    using _ = this.tracer.span('storage', 'flush');
+
     await this.withReplicationTransaction(`Flushing ${batch?.length ?? 0} ops`, async (session, opSeq) => {
       if (batch != null) {
         resumeBatch = await this.replicateBatch(session, batch, opSeq, options);
@@ -203,6 +210,7 @@ export abstract class MongoBucketBatch
     options?: storage.BucketBatchCommitOptions
   ): Promise<OperationBatch | null> {
     let sizes: Map<string, number> | undefined = undefined;
+    using _ = this.tracer.span('storage', 'replicate_batch');
     if (this.storeCurrentData && !this.skipExistingRows) {
       // We skip this step if we don't store current_data, since the sizes will
       // always be small in that case.
@@ -240,14 +248,19 @@ export abstract class MongoBucketBatch
         }
         continue;
       }
+      using lookupSpan = this.tracer.span('storage', 'lookup');
       const lookups = b.map((r) => ({
         sourceTableId: mongoTableId(r.record.sourceTable.id),
         replicaId: r.beforeId
       }));
       let sourceRecordLookup = await this.sourceRecordStore.loadDocuments(session, lookups, this.skipExistingRows);
+      lookupSpan.end();
 
       let persistedBatch: PersistedBatch | null = this.createPersistedBatch(transactionSize);
 
+      // The current code structure makes it tricky to cleanly split this span from the one
+      // where fluhsing. So we manually end and re-create this span whenever we flush.
+      let evalSpan = this.tracer.span('evaluate');
       for (let op of b) {
         if (resumeBatch) {
           resumeBatch.push(op);
@@ -266,26 +279,34 @@ export abstract class MongoBucketBatch
         }
 
         if (persistedBatch!.shouldFlushTransaction()) {
+          evalSpan.end();
           // Transaction is getting big.
           // Flush, and resume in a new transaction.
+          using persistSpan = this.tracer.span('storage', 'persist_flush');
           const { flushedAny } = await persistedBatch!.flush(this.session, options);
+
           didFlush ||= flushedAny;
           persistedBatch = null;
           // Computing our current progress is a little tricky here, since
           // we're stopping in the middle of a batch.
           // We create a new batch, and push any remaining operations to it.
           resumeBatch = new OperationBatch();
+          persistSpan.end();
+          evalSpan = this.tracer.span('evaluate');
         }
       }
+      evalSpan.end();
 
       if (persistedBatch) {
         transactionSize = persistedBatch.currentSize;
+        using _ = this.tracer.span('storage', 'persist_flush');
         const { flushedAny } = await persistedBatch.flush(this.session, options);
         didFlush ||= flushedAny;
       }
     }
 
     if (didFlush) {
+      using _ = this.tracer.span('storage', 'clear_error');
       await this.clearError();
     }
 
@@ -540,7 +561,9 @@ export abstract class MongoBucketBatch
   }
 
   private async withTransaction(cb: () => Promise<void>) {
+    using lockSpan = this.tracer.span('storage', 'internal_lock');
     await replicationMutex.exclusiveLock(async () => {
+      lockSpan.end();
       await this.session.withTransaction(
         async () => {
           try {
@@ -551,7 +574,9 @@ export abstract class MongoBucketBatch
             } else {
               this.logger.warn('Transaction error', e as Error);
             }
-            await timers.setTimeout(Math.random() * 50);
+            const delay = Math.random() * 50;
+            using _ = this.tracer.span('storage', 'retry_delay');
+            await timers.setTimeout(delay);
             throw e;
           }
         },
@@ -979,6 +1004,7 @@ export abstract class MongoBucketBatch
     let lastBatchCount = BATCH_LIMIT;
     while (lastBatchCount == BATCH_LIMIT) {
       await this.withReplicationTransaction(`Truncate ${sourceTable.qualifiedName}`, async (session, opSeq) => {
+        using evalSpan = this.tracer.span('evaluate');
         const sourceTableId = mongoTableId(sourceTable.id);
         const batch = await this.sourceRecordStore.loadTruncateBatch(session, sourceTableId, BATCH_LIMIT);
         const persistedBatch = this.createPersistedBatch(0);
@@ -1002,6 +1028,9 @@ export abstract class MongoBucketBatch
           // Since this is not from streaming replication, we can do a hard delete
           persistedBatch.hardDeleteCurrentData(sourceTableId, value.replicaId);
         }
+        evalSpan.end();
+
+        using _ = this.tracer.span('storage', 'persist_flush');
         await persistedBatch.flush(session);
         lastBatchCount = batch.length;
 

modules/module-mongodb-storage/src/storage/implementation/MongoSyncBucketStorage.ts

@@ -198,7 +198,7 @@ export abstract class MongoSyncBucketStorage
     );
     const checkpoint_lsn = doc?.last_checkpoint_lsn ?? null;
 
-    const batchOptions = {
+    const batchOptions: MongoBucketBatchOptions = {
       logger: options.logger,
       db: this.db,
       syncRules: this.sync_rules.parsed(options).hydratedSyncRules(),
@@ -210,7 +210,8 @@ export abstract class MongoSyncBucketStorage
       keepaliveOp: doc?.keepalive_op ? BigInt(doc.keepalive_op) : null,
       storeCurrentData: options.storeCurrentData,
       skipExistingRows: options.skipExistingRows ?? false,
-      markRecordUnavailable: options.markRecordUnavailable
+      markRecordUnavailable: options.markRecordUnavailable,
+      tracer: options.tracer
     };
     const writer = this.createWriterImpl(batchOptions);
     this.iterateListeners((cb) => cb.batchStarted?.(writer));

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

@@ -11,6 +11,7 @@ import {
 } from '@powersync/lib-services-framework';
 import {
   MetricsEngine,
+  PerformanceTracer,
   RelationCache,
   ReplicationLagTracker,
   SaveOperationTag,
@@ -20,6 +21,7 @@ import {
 } from '@powersync/service-core';
 import { HydratedSyncRules, TablePattern } 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';
@@ -34,7 +36,6 @@ import {
 } 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;
@@ -328,14 +329,16 @@ export class ChangeStream {
   async initialReplication(snapshotLsn: string | null) {
     const sourceTables = this.sync_rules.getSourceTables();
     await this.client.connect();
+    const tracer = new PerformanceTracer('MongoDB initial replication');
 
     const flushResult = await this.storage.startBatch(
       {
         logger: this.logger,
         zeroLSN: MongoLSN.ZERO.comparable,
         defaultSchema: this.defaultDb.databaseName,
         storeCurrentData: false,
-        skipExistingRows: true
+        skipExistingRows: true,
+        tracer
       },
       async (batch) => {
         if (snapshotLsn == null) {
@@ -510,7 +513,7 @@ export class ChangeStream {
       // Pre-fetch next batch, so that we can read and write concurrently
       nextChunkPromise = query.nextChunk();
       for (let buffer of docBatch) {
-        const { row: record, replicaId: replicaId } = this.sourceRowConverter.rawToSqliteRow(buffer);
+        const { row: record, replicaId: replicaId } = this.rawToSqliteRow(buffer);
 
         // This auto-flushes when the batch reaches its size limit
         await batch.save({
@@ -660,7 +663,7 @@ export class ChangeStream {
 
     this.metrics.getCounter(ReplicationMetric.ROWS_REPLICATED).add(1);
     if (change.operationType == 'insert') {
-      const { row: baseRecord, replicaId: _replicaId } = this.sourceRowConverter.rawToSqliteRow(change.fullDocument);
+      const { row: baseRecord, replicaId: _replicaId } = this.rawToSqliteRow(change.fullDocument);
       return await batch.save({
         tag: SaveOperationTag.INSERT,
         sourceTable: table,
@@ -682,7 +685,7 @@ export class ChangeStream {
           beforeReplicaId: change.documentKey._id
         });
       }
-      const { row: after, replicaId: _replicaId } = this.sourceRowConverter.rawToSqliteRow(change.fullDocument!);
+      const { row: after, replicaId: _replicaId } = this.rawToSqliteRow(change.fullDocument!);
       return await batch.save({
         tag: SaveOperationTag.UPDATE,
         sourceTable: table,
@@ -756,6 +759,7 @@ export class ChangeStream {
     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;
@@ -813,22 +817,29 @@ export class ChangeStream {
       maxTimeMS: this.changeStreamTimeout,
 
       signal: options.signal,
-      logger: this.logger
+      logger: this.logger,
+      tracer: options.tracer
     });
   }
 
+  private rawToSqliteRow(row: Buffer) {
+    return this.sourceRowConverter.rawToSqliteRow(row);
+  }
+
   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);
 
+    const tracer = new PerformanceTracer('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,
+        tracer
       },
       async (batch) => {
         const { resumeFromLsn } = batch;
@@ -837,6 +848,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.
@@ -849,7 +861,8 @@ export class ChangeStream {
         const batchStream = this.rawChangeStreamBatches({
           lsn: resumeFromLsn,
           filters,
-          signal: this.abort_signal
+          signal: this.abort_signal,
+          tracer
         });
 
         // Always start with a checkpoint.
@@ -864,13 +877,14 @@ export class ChangeStream {
         let splitDocument: ProjectedChangeStreamDocument | null = null;
 
         let flexDbNameWorkaroundLogged = false;
-        let changesSinceLastCheckpoint = 0;
 
         let lastEmptyResume = performance.now();
         let lastTxnKey: string | null = null;
 
         for await (let eventBatch of batchStream) {
           const { events, resumeToken } = eventBatch;
+          using batchSpan = tracer.span('processing');
+
           bytesReplicatedMetric.add(eventBatch.byteSize);
           chunksReplicatedMetric.add(1);
           if (this.abort_signal.aborted) {
@@ -904,7 +918,6 @@ export class ChangeStream {
 
           this.touch();
 
-          const batchStart = Date.now();
           for (let eventIndex = 0; eventIndex < events.length; eventIndex++) {
             const rawChangeDocument = events[eventIndex];
             const originalChangeDocument = parseChangeDocument(rawChangeDocument);
@@ -1048,7 +1061,6 @@ export class ChangeStream {
 
               if (!checkpointBlocked) {
                 this.replicationLag.markCommitted();
-                changesSinceLastCheckpoint = 0;
               }
             } else if (
               changeDocument.operationType == 'insert' ||
@@ -1083,21 +1095,7 @@ export class ChangeStream {
                   transactionsReplicatedMetric.add(1);
                 }
 
-                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;
-                }
+                await this.writeChange(batch, table, changeDocument);
               }
             } else if (changeDocument.operationType == 'drop') {
               const rel = getMongoRelation(changeDocument.ns);
@@ -1140,7 +1138,21 @@ export class ChangeStream {
             // 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`);
+
+          batchSpan.end();
+          const durations = outerSpan.end();
+          const duration = batchSpan.endAt - batchSpan.startAt;
+
+          this.logger.info(
+            `Processed batch of ${events.length} changes / ${eventBatch.byteSize} bytes in ${duration}ms`,
+            {
+              count: events.length,
+              bytes: eventBatch.byteSize,
+              duration,
+              t: durations
+            }
+          );
+          outerSpan = tracer.span('batch');
         }
       }
     );