Fix new internal package typecheck issues and start adding telemetry to the replication service

Author: ericallam

apps/webapp/app/services/runsReplicationService.server.ts

@@ -10,6 +10,7 @@ import EventEmitter from "node:events";
 import pLimit from "p-limit";
 import { Counter, Gauge } from "prom-client";
 import type { MetricsRegister } from "~/metrics.server";
+import { Span, type Tracer, recordSpanError, trace } from "@internal/tracing";
 
 export type RunsReplicationServiceOptions = {
   clickhouse: ClickHouse;
@@ -27,6 +28,7 @@ export type RunsReplicationServiceOptions = {
   leaderLockExtendIntervalMs?: number;
   ackIntervalSeconds?: number;
   logger?: Logger;
+  tracer?: Tracer;
 };
 
 type TaskRunInsert = { _version: bigint; run: TaskRun; event: "insert" | "update" | "delete" };
@@ -54,12 +56,15 @@ export class RunsReplicationService {
   private _insertStrategy: "streaming" | "batching";
   private _isShuttingDown = false;
   private _isShutDownComplete = false;
+  private _tracer: Tracer;
+  private _currentSpan: Span | null = null;
 
   public readonly events: EventEmitter<RunsReplicationServiceEvents>;
 
   constructor(private readonly options: RunsReplicationServiceOptions) {
     this.logger = options.logger ?? new Logger("RunsReplicationService", "debug");
     this.events = new EventEmitter();
+    this._tracer = options.tracer ?? trace.getTracer("runs-replication-service");
 
     this._insertStrategy = options.insertStrategy ?? "streaming";
 
@@ -206,6 +211,13 @@ export class RunsReplicationService {
           xid: message.xid,
           events: [],
         };
+
+        this._currentSpan = this._tracer.startSpan("handle_transaction", {
+          attributes: {
+            "transaction.xid": message.xid,
+          },
+        });
+
         break;
       }
       case "insert": {
@@ -269,6 +281,8 @@ export class RunsReplicationService {
       // We need to immediately acknowledge the transaction
       // And then try and handle this transaction
       if (transaction.commitEndLsn) {
+        this._currentSpan?.setAttribute("transaction.shutdown", true);
+
         await this._replicationClient.acknowledge(transaction.commitEndLsn);
         alreadyAcknowledged = true;
       }
@@ -279,12 +293,23 @@ export class RunsReplicationService {
 
     this._lastReplicationLagMs = transaction.replicationLagMs;
 
+    this._currentSpan?.setAttribute("transaction.replication_lag_ms", transaction.replicationLagMs);
+    this._currentSpan?.setAttribute("transaction.xid", transaction.xid);
+
+    if (transaction.commitEndLsn) {
+      this._currentSpan?.setAttribute("transaction.commit_end_lsn", transaction.commitEndLsn);
+    }
+
+    this._currentSpan?.setAttribute("transaction.events", transaction.events.length);
+
     // If there are no events, do nothing
     if (transaction.events.length === 0) {
       if (transaction.commitEndLsn && !alreadyAcknowledged) {
         await this._replicationClient.acknowledge(transaction.commitEndLsn);
       }
 
+      this._currentSpan?.end();
+
       return;
     }
 
@@ -293,6 +318,8 @@ export class RunsReplicationService {
         transaction,
       });
 
+      this._currentSpan?.end();
+
       return;
     }
 
@@ -301,12 +328,19 @@ export class RunsReplicationService {
       alreadyAcknowledged,
     });
 
+    const lsnToUInt64Start = process.hrtime.bigint();
+
     // If there are events, we need to handle them
     const _version = lsnToUInt64(transaction.commitEndLsn);
 
+    this._currentSpan?.setAttribute(
+      "transaction.lsn_to_uint64_ms",
+      Number(process.hrtime.bigint() - lsnToUInt64Start) / 1_000_000
+    );
+
     this._transactionCounter?.inc();
 
-    if (this._insertStrategy === "streaming") {
+    if (this._insertStrategy === "batching") {
       this._concurrentFlushScheduler
         .addToBatch(
           transaction.events.map((event) => ({
@@ -336,12 +370,23 @@ export class RunsReplicationService {
         this.logger.error("Error flushing batch", {
           error: flushError,
         });
+
+        if (this._currentSpan) {
+          recordSpanError(this._currentSpan, flushError);
+        }
       }
     }
 
     if (!alreadyAcknowledged) {
+      const acknowledgeStart = process.hrtime.bigint();
       await this._replicationClient.acknowledge(transaction.commitEndLsn);
+      this._currentSpan?.setAttribute(
+        "transaction.acknowledge_ms",
+        Number(process.hrtime.bigint() - acknowledgeStart) / 1_000_000
+      );
     }
+
+    this._currentSpan?.end();
   }
 
   async #flushBatch(flushId: string, batch: Array<TaskRunInsert>) {

apps/webapp/package.json

@@ -105,6 +105,7 @@
     "@trigger.dev/otlp-importer": "workspace:*",
     "@trigger.dev/platform": "1.0.14",
     "@trigger.dev/sdk": "workspace:*",
+    "@internal/tracing": "workspace:*",
     "@types/pg": "8.6.6",
     "@uiw/react-codemirror": "^4.19.5",
     "@unkey/cache": "^1.5.0",

internal-packages/clickhouse/package.json

@@ -20,7 +20,7 @@
   },
   "scripts": {
     "clean": "rimraf dist",
-    "typecheck": "tsc --noEmit",
+    "typecheck": "tsc --noEmit -p tsconfig.build.json",
     "build": "pnpm run clean && tsc -p tsconfig.build.json",
     "dev": "tsc --watch  -p tsconfig.build.json",
     "db:migrate": "docker compose -p triggerdotdev-docker -f ../../docker/docker-compose.yml up clickhouse_migrator --build",

internal-packages/replication/package.json

@@ -21,7 +21,7 @@
   },
   "scripts": {
     "clean": "rimraf dist",
-    "typecheck": "tsc --noEmit",
+    "typecheck": "tsc --noEmit -p tsconfig.build.json",
     "build": "pnpm run clean && tsc -p tsconfig.build.json",
     "dev": "tsc --watch  -p tsconfig.build.json",
     "test": "vitest --sequence.concurrent=false --no-file-parallelism",

internal-packages/replication/src/client.ts

@@ -7,6 +7,7 @@ import { createRedisClient } from "@internal/redis";
 import { Logger } from "@trigger.dev/core/logger";
 import { LogicalReplicationClientError } from "./errors.js";
 import { PgoutputMessage, PgoutputParser, getPgoutputStartReplicationSQL } from "./pgoutput.js";
+import { startSpan, trace, Tracer } from "@internal/tracing";
 
 export interface LogicalReplicationClientOptions {
   /**
@@ -70,6 +71,8 @@ export interface LogicalReplicationClientOptions {
    * The actions to publish to the publication.
    */
   publicationActions?: Array<"insert" | "update" | "delete" | "truncate">;
+
+  tracer?: Tracer;
 }
 
 export type LogicalReplicationClientEvents = {
@@ -101,6 +104,7 @@ export class LogicalReplicationClient {
   private lastAckTimestamp: number = 0;
   private ackIntervalTimer: NodeJS.Timeout | null = null;
   private _isStopped: boolean = false;
+  private _tracer: Tracer;
 
   public get lastLsn(): string {
     return this.lastAcknowledgedLsn ?? "0/00000000";
@@ -113,6 +117,7 @@ export class LogicalReplicationClient {
   constructor(options: LogicalReplicationClientOptions) {
     this.options = options;
     this.logger = options.logger ?? new Logger("LogicalReplicationClient", "info");
+    this._tracer = options.tracer ?? trace.getTracer("logical-replication-client");
 
     this.autoAcknowledge =
       typeof options.autoAcknowledge === "boolean" ? options.autoAcknowledge : true;
@@ -145,54 +150,62 @@ export class LogicalReplicationClient {
   }
 
   public async stop(): Promise<this> {
-    if (this._isStopped) return this;
-    this._isStopped = true;
-    // Clean up leader lock heartbeat
-    if (this.leaderLockHeartbeatTimer) {
-      clearInterval(this.leaderLockHeartbeatTimer);
-      this.leaderLockHeartbeatTimer = null;
-    }
-    // Clean up ack interval
-    if (this.ackIntervalTimer) {
-      clearInterval(this.ackIntervalTimer);
-      this.ackIntervalTimer = null;
-    }
-    // Release leader lock if held
-    await this.#releaseLeaderLock();
+    return await startSpan(this._tracer, "logical_replication_client.stop", async (span) => {
+      if (this._isStopped) return this;
+
+      span.setAttribute("replication_client.name", this.options.name);
+      span.setAttribute("replication_client.table", this.options.table);
+      span.setAttribute("replication_client.slot_name", this.options.slotName);
+      span.setAttribute("replication_client.publication_name", this.options.publicationName);
+
+      this._isStopped = true;
+      // Clean up leader lock heartbeat
+      if (this.leaderLockHeartbeatTimer) {
+        clearInterval(this.leaderLockHeartbeatTimer);
+        this.leaderLockHeartbeatTimer = null;
+      }
+      // Clean up ack interval
+      if (this.ackIntervalTimer) {
+        clearInterval(this.ackIntervalTimer);
+        this.ackIntervalTimer = null;
+      }
+      // Release leader lock if held
+      await this.#releaseLeaderLock();
 
-    this.connection?.removeAllListeners();
-    this.connection = null;
+      this.connection?.removeAllListeners();
+      this.connection = null;
 
-    if (this.client) {
-      this.client.removeAllListeners();
+      if (this.client) {
+        this.client.removeAllListeners();
 
-      const [endError] = await tryCatch(this.client.end());
+        const [endError] = await tryCatch(this.client.end());
 
-      if (endError) {
-        this.logger.error("Failed to end client", {
-          name: this.options.name,
-          error: endError,
-        });
-      } else {
-        this.logger.info("Ended client", {
-          name: this.options.name,
-        });
+        if (endError) {
+          this.logger.error("Failed to end client", {
+            name: this.options.name,
+            error: endError,
+          });
+        } else {
+          this.logger.info("Ended client", {
+            name: this.options.name,
+          });
+        }
+        this.client = null;
       }
-      this.client = null;
-    }
 
-    // clear any intervals
-    if (this.leaderLockHeartbeatTimer) {
-      clearInterval(this.leaderLockHeartbeatTimer);
-      this.leaderLockHeartbeatTimer = null;
-    }
+      // clear any intervals
+      if (this.leaderLockHeartbeatTimer) {
+        clearInterval(this.leaderLockHeartbeatTimer);
+        this.leaderLockHeartbeatTimer = null;
+      }
 
-    if (this.ackIntervalTimer) {
-      clearInterval(this.ackIntervalTimer);
-      this.ackIntervalTimer = null;
-    }
+      if (this.ackIntervalTimer) {
+        clearInterval(this.ackIntervalTimer);
+        this.ackIntervalTimer = null;
+      }
 
-    return this;
+      return this;
+    });
   }
 
   public async teardown(): Promise<boolean> {
@@ -523,34 +536,43 @@ export class LogicalReplicationClient {
   public async acknowledge(lsn: string): Promise<boolean> {
     if (this._isStopped) return false;
     if (!this.connection) return false;
-    // WAL LSN split
-    const slice = lsn.split("/");
-    let [upperWAL, lowerWAL]: [number, number] = [parseInt(slice[0], 16), parseInt(slice[1], 16)];
-    // Timestamp as microseconds since midnight 2000-01-01
-    const now = Date.now() - 946080000000;
-    const upperTimestamp = Math.floor(now / 4294967.296);
-    const lowerTimestamp = Math.floor(now - upperTimestamp * 4294967.296);
-    if (lowerWAL === 4294967295) {
-      upperWAL = upperWAL + 1;
-      lowerWAL = 0;
-    } else {
-      lowerWAL = lowerWAL + 1;
-    }
-    const response = Buffer.alloc(34);
-    response.fill(0x72); // 'r'
-    response.writeUInt32BE(upperWAL, 1);
-    response.writeUInt32BE(lowerWAL, 5);
-    response.writeUInt32BE(upperWAL, 9);
-    response.writeUInt32BE(lowerWAL, 13);
-    response.writeUInt32BE(upperWAL, 17);
-    response.writeUInt32BE(lowerWAL, 21);
-    response.writeUInt32BE(upperTimestamp, 25);
-    response.writeUInt32BE(lowerTimestamp, 29);
-    response.writeInt8(0, 33);
-    // @ts-ignore
-    this.connection.sendCopyFromChunk(response);
-    this.lastAckTimestamp = Date.now();
-    return true;
+
+    return await startSpan(this._tracer, "logical_replication_client.acknowledge", async (span) => {
+      span.setAttribute("replication_client.lsn", lsn);
+      span.setAttribute("replication_client.name", this.options.name);
+      span.setAttribute("replication_client.table", this.options.table);
+      span.setAttribute("replication_client.slot_name", this.options.slotName);
+      span.setAttribute("replication_client.publication_name", this.options.publicationName);
+
+      // WAL LSN split
+      const slice = lsn.split("/");
+      let [upperWAL, lowerWAL]: [number, number] = [parseInt(slice[0], 16), parseInt(slice[1], 16)];
+      // Timestamp as microseconds since midnight 2000-01-01
+      const now = Date.now() - 946080000000;
+      const upperTimestamp = Math.floor(now / 4294967.296);
+      const lowerTimestamp = Math.floor(now - upperTimestamp * 4294967.296);
+      if (lowerWAL === 4294967295) {
+        upperWAL = upperWAL + 1;
+        lowerWAL = 0;
+      } else {
+        lowerWAL = lowerWAL + 1;
+      }
+      const response = Buffer.alloc(34);
+      response.fill(0x72); // 'r'
+      response.writeUInt32BE(upperWAL, 1);
+      response.writeUInt32BE(lowerWAL, 5);
+      response.writeUInt32BE(upperWAL, 9);
+      response.writeUInt32BE(lowerWAL, 13);
+      response.writeUInt32BE(upperWAL, 17);
+      response.writeUInt32BE(lowerWAL, 21);
+      response.writeUInt32BE(upperTimestamp, 25);
+      response.writeUInt32BE(lowerTimestamp, 29);
+      response.writeInt8(0, 33);
+      // @ts-ignore
+      this.connection.sendCopyFromChunk(response);
+      this.lastAckTimestamp = Date.now();
+      return true;
+    });
   }
 
   async #acquireLeaderLock(): Promise<boolean> {