Replication logging improvements.

Author: rkistner

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

@@ -290,6 +290,9 @@ export class ChangeStream {
           });
           tablesWithStatus.push(updated);
           this.relationCache.update(updated);
+          logger.info(
+            `${this.logPrefix} To replicate: ${table.qualifiedName}: ${updated.snapshotStatus?.replicatedCount}/~${updated.snapshotStatus?.totalEstimatedCount}`
+          );
         }
 
         for (let table of tablesWithStatus) {
@@ -366,14 +369,17 @@ export class ChangeStream {
   }
 
   private async snapshotTable(batch: storage.BucketStorageBatch, table: storage.SourceTable) {
-    logger.info(`${this.logPrefix} Replicating ${table.qualifiedName}`);
     const totalEstimatedCount = await this.estimatedCountNumber(table);
     let at = table.snapshotStatus?.replicatedCount ?? 0;
     const db = this.client.db(table.schema);
     const collection = db.collection(table.table);
     await using query = new ChunkedSnapshotQuery({ collection, key: table.snapshotStatus?.lastKey });
     if (query.lastKey != null) {
-      logger.info(`${this.logPrefix} Resuming snapshot for ${table.qualifiedName} at ${query.lastKey}`);
+      logger.info(
+        `${this.logPrefix} Replicating ${table.qualifiedName} ${table.formatSnapshotProgress()} - resuming at _id > ${query.lastKey}`
+      );
+    } else {
+      logger.info(`${this.logPrefix} Replicating ${table.qualifiedName} ${table.formatSnapshotProgress()}`);
     }
 
     let lastBatch = performance.now();
@@ -418,7 +424,7 @@ export class ChangeStream {
       const duration = performance.now() - lastBatch;
       lastBatch = performance.now();
       logger.info(
-        `${this.logPrefix} Replicating ${table.qualifiedName} ${at}/${totalEstimatedCount} in ${duration.toFixed(0)}ms`
+        `${this.logPrefix} Replicating ${table.qualifiedName} ${table.formatSnapshotProgress()} in ${duration.toFixed(0)}ms`
       );
       await touch();
     }

modules/module-postgres/src/replication/SnapshotQuery.ts

@@ -85,16 +85,21 @@ export class ChunkedSnapshotQuery implements SnapshotQuery {
   public constructor(
     private readonly connection: PgConnection,
     private readonly table: SourceTable,
-    private readonly chunkSize: number = 10_000
+    private readonly chunkSize: number = 10_000,
+    lastKeySerialized: Uint8Array | null
   ) {
     this.key = table.replicaIdColumns[0];
+
+    if (lastKeySerialized != null) {
+      this.lastKey = this.deserializeKey(lastKeySerialized);
+    }
   }
 
   public async initialize(): Promise<void> {
     // No-op
   }
 
-  public setLastKeySerialized(key: Uint8Array) {
+  private deserializeKey(key: Uint8Array) {
     const decoded = bson.deserialize(key, { useBigInt64: true });
     const keys = Object.keys(decoded);
     if (keys.length != 1) {
@@ -104,7 +109,7 @@ export class ChunkedSnapshotQuery implements SnapshotQuery {
       throw new ServiceAssertionError(`Key name mismatch: expected ${this.key.name}, got ${keys[0]}`);
     }
     const value = decoded[this.key.name];
-    this.lastKey = value;
+    return value;
   }
 
   public getLastKeySerialized(): Uint8Array {
@@ -113,17 +118,18 @@ export class ChunkedSnapshotQuery implements SnapshotQuery {
 
   public async *nextChunk(): AsyncIterableIterator<PgChunk> {
     let stream: AsyncIterableIterator<PgChunk>;
+    const escapedKeyName = escapeIdentifier(this.key.name);
     if (this.lastKey == null) {
       stream = this.connection.stream(
-        `SELECT * FROM ${this.table.escapedIdentifier} ORDER BY ${escapeIdentifier(this.key.name)} LIMIT ${this.chunkSize}`
+        `SELECT * FROM ${this.table.escapedIdentifier} ORDER BY ${escapedKeyName} LIMIT ${this.chunkSize}`
       );
     } else {
       if (this.key.typeId == null) {
         throw new Error(`typeId required for primary key ${this.key.name}`);
       }
       let type: StatementParam['type'] = Number(this.key.typeId);
       stream = this.connection.stream({
-        statement: `SELECT * FROM ${this.table.escapedIdentifier} WHERE ${escapeIdentifier(this.key.name)} > $1 ORDER BY ${escapeIdentifier(this.key.name)} LIMIT ${this.chunkSize}`,
+        statement: `SELECT * FROM ${this.table.escapedIdentifier} WHERE ${escapedKeyName} > $1 ORDER BY ${escapedKeyName} LIMIT ${this.chunkSize}`,
         params: [{ value: this.lastKey, type }]
       });
     }

modules/module-postgres/src/replication/WalStream.ts

@@ -434,8 +434,10 @@ WHERE  oid = $1::regclass`,
     await this.storage.startBatch(
       { zeroLSN: ZERO_LSN, defaultSchema: POSTGRES_DEFAULT_SCHEMA, storeCurrentData: true, skipExistingRows: true },
       async (batch) => {
+        let tablesWithStatus: SourceTable[] = [];
         for (let tablePattern of sourceTables) {
           const tables = await this.getQualifiedTableNames(batch, db, tablePattern);
+          // Pre-get counts
           for (let table of tables) {
             if (table.snapshotComplete) {
               logger.info(`${this.slot_name} Skipping ${table.qualifiedName} - snapshot already done`);
@@ -444,11 +446,17 @@ WHERE  oid = $1::regclass`,
             const count = await this.estimatedCountNumber(db, table);
             table = await batch.updateTableProgress(table, { totalEstimatedCount: count });
             this.relationCache.update(table);
-            await this.snapshotTableInTx(batch, db, table);
-            await touch();
+            tablesWithStatus.push(table);
+
+            logger.info(`${this.slot_name} To replicate: ${table.qualifiedName} ${table.formatSnapshotProgress()}`);
           }
         }
 
+        for (let table of tablesWithStatus) {
+          await this.snapshotTableInTx(batch, db, table);
+          await touch();
+        }
+
         // Always commit the initial snapshot at zero.
         // This makes sure we don't skip any changes applied before starting this snapshot,
         // in the case of snapshot retries.
@@ -505,11 +513,9 @@ WHERE  oid = $1::regclass`,
     table: storage.SourceTable,
     limited?: PrimaryKeyValue[]
   ) {
-    logger.info(`${this.slot_name} Replicating ${table.qualifiedName}`);
     let totalEstimatedCount = table.snapshotStatus?.totalEstimatedCount;
     let at = table.snapshotStatus?.replicatedCount ?? 0;
     let lastCountTime = 0;
-    let lastLogIndex = 0;
     let q: SnapshotQuery;
     // We do streaming on two levels:
     // 1. Coarse level: DELCARE CURSOR, FETCH 10000 at a time.
@@ -519,15 +525,21 @@ WHERE  oid = $1::regclass`,
     } else if (ChunkedSnapshotQuery.supports(table)) {
       // Single primary key - we can use the primary key for chunking
       const orderByKey = table.replicaIdColumns[0];
-      logger.info(`Chunking ${table.qualifiedName} by ${orderByKey.name}`);
-      q = new ChunkedSnapshotQuery(db, table, this.snapshotChunkSize);
+      q = new ChunkedSnapshotQuery(db, table, this.snapshotChunkSize, table.snapshotStatus?.lastKey ?? null);
       if (table.snapshotStatus?.lastKey != null) {
-        (q as ChunkedSnapshotQuery).setLastKeySerialized(table.snapshotStatus!.lastKey);
-        logger.info(`Resuming from ${(q as ChunkedSnapshotQuery).lastKey}`);
+        logger.info(
+          `${this.slot_name} Replicating ${table.qualifiedName} ${table.formatSnapshotProgress()} - resuming from ${orderByKey.name} > ${(q as ChunkedSnapshotQuery).lastKey}`
+        );
+      } else {
+        logger.info(
+          `${this.slot_name} Replicating ${table.qualifiedName} ${table.formatSnapshotProgress()} - resumable`
+        );
       }
     } else {
       // Fallback case - query the entire table
-      logger.info(`Snapshot ${table.qualifiedName} without chunking`);
+      logger.info(
+        `${this.slot_name} Replicating ${table.qualifiedName} ${table.formatSnapshotProgress()} - not resumable`
+      );
       q = new SimpleSnapshotQuery(db, table, this.snapshotChunkSize);
       at = 0;
     }
@@ -566,12 +578,6 @@ WHERE  oid = $1::regclass`,
         if (rows.length > 0) {
           hasRemainingData = true;
         }
-        if (rows.length > 0 && at - lastLogIndex >= 5000) {
-          logger.info(
-            `${this.slot_name} Replicating ${table.qualifiedName} ${at}/${limited ? `${limited.length}L` : this.formatCount(totalEstimatedCount)}`
-          );
-          lastLogIndex = at;
-        }
         if (this.abort_signal.aborted) {
           throw new Error(`Aborted initial replication of ${this.slot_name}`);
         }
@@ -593,9 +599,6 @@ WHERE  oid = $1::regclass`,
 
         await touch();
       }
-      logger.info(
-        `${this.slot_name} Replicated ${table.qualifiedName} ${at}/${limited ? `${limited.length}L` : this.formatCount(totalEstimatedCount)}`
-      );
 
       if (limited == null) {
         // Important: flush before marking progress
@@ -614,6 +617,10 @@ WHERE  oid = $1::regclass`,
           totalEstimatedCount: totalEstimatedCount
         });
         this.relationCache.update(table);
+
+        logger.info(`${this.slot_name} Replicating ${table.qualifiedName} ${table.formatSnapshotProgress()}`);
+      } else {
+        logger.info(`${this.slot_name} Replicating ${table.qualifiedName} ${at}/${limited.length} for resnapshot`);
       }
     }
     await batch.flush();

packages/service-core/src/storage/SourceTable.ts

@@ -100,4 +100,15 @@ export class SourceTable {
     copy.snapshotStatus = this.snapshotStatus;
     return copy;
   }
+
+  formatSnapshotProgress() {
+    if (this.snapshotComplete || this.snapshotStatus == null) {
+      // Should not happen
+      return '-';
+    } else if (this.snapshotStatus.totalEstimatedCount < 0) {
+      return `${this.snapshotStatus.replicatedCount}/unknown`;
+    } else {
+      return `${this.snapshotStatus.replicatedCount}/~${this.snapshotStatus.totalEstimatedCount}`;
+    }
+  }
 }