Handover WIP and tests

Author: ericallam

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

@@ -61,6 +61,9 @@ function initializeRunsReplicationInstance() {
           error,
         });
       });
+
+    process.on("SIGTERM", service.shutdown.bind(service));
+    process.on("SIGINT", service.shutdown.bind(service));
   }
 
   return service;

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

@@ -1,9 +1,8 @@
-import type { ClickHouse, TaskRunV1, RawTaskRunPayloadV1 } from "@internal/clickhouse";
+import type { ClickHouse, RawTaskRunPayloadV1, TaskRunV1 } from "@internal/clickhouse";
 import { RedisOptions } from "@internal/redis";
 import { LogicalReplicationClient, Transaction, type PgoutputMessage } from "@internal/replication";
 import { Logger } from "@trigger.dev/core/logger";
 import { tryCatch } from "@trigger.dev/core/utils";
-import { TaskRunError } from "@trigger.dev/core/v3/schemas";
 import { parsePacket } from "@trigger.dev/core/v3/utils/ioSerialization";
 import { TaskRun } from "@trigger.dev/database";
 import { nanoid } from "nanoid";
@@ -46,6 +45,8 @@ export class RunsReplicationService {
   private _lastReplicationLagMs: number | null = null;
   private _transactionCounter?: Counter;
   private _insertStrategy: "streaming" | "batching";
+  private _isShuttingDown = false;
+  private _isShutDownComplete = false;
 
   constructor(private readonly options: RunsReplicationServiceOptions) {
     this.logger = new Logger("RunsReplicationService", "debug");
@@ -62,7 +63,7 @@ export class RunsReplicationService {
       table: "TaskRun",
       redisOptions: options.redisOptions,
       autoAcknowledge: false,
-      publicationActions: ["insert", "update"],
+      publicationActions: ["insert", "update", "delete"],
       logger: new Logger("RunsReplicationService", "debug"),
       leaderLockTimeoutMs: options.leaderLockTimeoutMs ?? 30_000,
       leaderLockExtendIntervalMs: options.leaderLockExtendIntervalMs ?? 10_000,
@@ -84,6 +85,9 @@ export class RunsReplicationService {
     });
 
     this._replicationClient.events.on("heartbeat", async ({ lsn, shouldRespond }) => {
+      if (this._isShuttingDown) return;
+      if (this._isShutDownComplete) return;
+
       if (shouldRespond) {
         await this._replicationClient.acknowledge(lsn);
       }
@@ -130,6 +134,11 @@ export class RunsReplicationService {
     }
   }
 
+  public shutdown() {
+    this.logger.info("Initiating shutdown of runs replication service");
+    this._isShuttingDown = true;
+  }
+
   async start(insertStrategy?: "streaming" | "batching") {
     this._insertStrategy = insertStrategy ?? this._insertStrategy;
 
@@ -201,11 +210,27 @@ export class RunsReplicationService {
   }
 
   async #handleTransaction(transaction: Transaction<TaskRun>) {
+    if (this._isShutDownComplete) return;
+
+    let alreadyAcknowledged = false;
+
+    if (this._isShuttingDown) {
+      // We need to immediately acknowledge the transaction
+      // And then try and handle this transaction
+      if (transaction.commitEndLsn) {
+        await this._replicationClient.acknowledge(transaction.commitEndLsn);
+        alreadyAcknowledged = true;
+      }
+
+      await this._replicationClient.stop();
+      this._isShutDownComplete = true;
+    }
+
     this._lastReplicationLagMs = transaction.replicationLagMs;
 
     // If there are no events, do nothing
     if (transaction.events.length === 0) {
-      if (transaction.commitEndLsn) {
+      if (transaction.commitEndLsn && !alreadyAcknowledged) {
         await this._replicationClient.acknowledge(transaction.commitEndLsn);
       }
 
@@ -222,6 +247,7 @@ export class RunsReplicationService {
 
     this.logger.debug("Handling transaction", {
       transaction,
+      alreadyAcknowledged,
     });
 
     // If there are events, we need to handle them
@@ -230,13 +256,19 @@ export class RunsReplicationService {
     this._transactionCounter?.inc();
 
     if (this._insertStrategy === "streaming") {
-      await this._concurrentFlushScheduler.addToBatch(
-        transaction.events.map((event) => ({
-          _version,
-          run: event.data,
-          event: event.tag,
-        }))
-      );
+      this._concurrentFlushScheduler
+        .addToBatch(
+          transaction.events.map((event) => ({
+            _version,
+            run: event.data,
+            event: event.tag,
+          }))
+        )
+        .catch((error) => {
+          this.logger.error("Error adding to batch", {
+            error,
+          });
+        });
     } else {
       const [flushError] = await tryCatch(
         this.#flushBatch(
@@ -256,7 +288,9 @@ export class RunsReplicationService {
       }
     }
 
-    await this._replicationClient.acknowledge(transaction.commitEndLsn);
+    if (!alreadyAcknowledged) {
+      await this._replicationClient.acknowledge(transaction.commitEndLsn);
+    }
   }
 
   async #flushBatch(flushId: string, batch: Array<TaskRunInsert>) {
@@ -497,7 +531,6 @@ export class ConcurrentFlushScheduler<T> {
   private readonly MAX_CONCURRENCY: number;
   private readonly concurrencyLimiter: ReturnType<typeof pLimit>;
   private flushTimer: NodeJS.Timeout | null;
-  private isShuttingDown;
   private failedBatchCount;
   private metricsRegister?: MetricsRegister;
   private logger: Logger;
@@ -510,7 +543,6 @@ export class ConcurrentFlushScheduler<T> {
     this.MAX_CONCURRENCY = config.maxConcurrency || 1;
     this.concurrencyLimiter = pLimit(this.MAX_CONCURRENCY);
     this.flushTimer = null;
-    this.isShuttingDown = false;
     this.failedBatchCount = 0;
 
     this.logger.info("Initializing ConcurrentFlushScheduler", {
@@ -520,7 +552,6 @@ export class ConcurrentFlushScheduler<T> {
     });
 
     this.startFlushTimer();
-    this.setupShutdownHandlers();
 
     if (!process.env.VITEST && config.metricsRegister) {
       this.metricsRegister = config.metricsRegister;
@@ -592,27 +623,6 @@ export class ConcurrentFlushScheduler<T> {
     this.logger.debug("Started flush timer", { interval: this.FLUSH_INTERVAL });
   }
 
-  private setupShutdownHandlers() {
-    process.on("SIGTERM", this.shutdown.bind(this));
-    process.on("SIGINT", this.shutdown.bind(this));
-    this.logger.debug("Shutdown handlers configured");
-  }
-
-  private async shutdown(): Promise<void> {
-    if (this.isShuttingDown) return;
-    this.isShuttingDown = true;
-    this.logger.info("Initiating shutdown of dynamic flush scheduler", {
-      remainingItems: this.currentBatch.length,
-    });
-
-    await this.checkAndFlush();
-    this.clearTimer();
-
-    this.logger.info("Dynamic flush scheduler shutdown complete", {
-      totalFailedBatches: this.failedBatchCount,
-    });
-  }
-
   private clearTimer(): void {
     if (this.flushTimer) {
       clearInterval(this.flushTimer);

apps/webapp/test/engine/triggerTask.test.ts

@@ -32,8 +32,6 @@ import {
   ValidationResult,
 } from "~/runEngine/types";
 import { RunEngineTriggerTaskService } from "../../app/runEngine/services/triggerTask.server";
-import { ClickHouse } from "@internal/clickhouse";
-import { RunsDashboardService } from "~/services/runsDashboardService.server";
 
 vi.setConfig({ testTimeout: 30_000 }); // 30 seconds timeout
 

apps/webapp/test/runsReplicationService.test.ts

@@ -0,0 +1,123 @@
+import { describe, vi } from "vitest";
+
+// Mock the db prisma client
+vi.mock("~/db.server", () => ({
+  prisma: {},
+}));
+
+vi.mock("~/services/platform.v3.server", () => ({
+  getEntitlement: vi.fn(),
+}));
+
+import { ClickHouse } from "@internal/clickhouse";
+import { containerTest } from "@internal/testcontainers";
+import { setTimeout } from "node:timers/promises";
+import { RunsReplicationService } from "~/services/runsReplicationService.server";
+import { z } from "zod";
+
+vi.setConfig({ testTimeout: 60_000 });
+
+describe("RunsReplicationService", () => {
+  containerTest(
+    "should replicate runs to clickhouse",
+    async ({ clickhouseContainer, redisOptions, postgresContainer, prisma }) => {
+      await prisma.$executeRawUnsafe(`ALTER TABLE public."TaskRun" REPLICA IDENTITY FULL;`);
+
+      const clickhouse = new ClickHouse({
+        url: clickhouseContainer.getConnectionUrl(),
+        name: "runs-replication",
+      });
+
+      const runsReplicationService = new RunsReplicationService({
+        clickhouse,
+        pgConnectionUrl: postgresContainer.getConnectionUri(),
+        serviceName: "runs-replication",
+        slotName: "task_runs_to_clickhouse_v1",
+        publicationName: "task_runs_to_clickhouse_v1_publication",
+        redisOptions,
+        maxFlushConcurrency: 1,
+        flushIntervalMs: 100,
+        flushBatchSize: 1,
+        insertStrategy: "batching",
+        leaderLockTimeoutMs: 1000,
+        leaderLockExtendIntervalMs: 1000,
+        ackIntervalSeconds: 5,
+      });
+
+      await runsReplicationService.start();
+
+      const organization = await prisma.organization.create({
+        data: {
+          title: "test",
+          slug: "test",
+        },
+      });
+
+      const project = await prisma.project.create({
+        data: {
+          name: "test",
+          slug: "test",
+          organizationId: organization.id,
+          externalRef: "test",
+        },
+      });
+
+      const runtimeEnvironment = await prisma.runtimeEnvironment.create({
+        data: {
+          slug: "test",
+          type: "DEVELOPMENT",
+          projectId: project.id,
+          organizationId: organization.id,
+          apiKey: "test",
+          pkApiKey: "test",
+          shortcode: "test",
+        },
+      });
+
+      // Now we insert a row into the table
+      const taskRun = await prisma.taskRun.create({
+        data: {
+          friendlyId: "run_1234",
+          taskIdentifier: "my-task",
+          payload: JSON.stringify({ foo: "bar" }),
+          traceId: "1234",
+          spanId: "1234",
+          queue: "test",
+          runtimeEnvironmentId: runtimeEnvironment.id,
+          projectId: project.id,
+          organizationId: organization.id,
+          environmentType: "DEVELOPMENT",
+          engine: "V2",
+        },
+      });
+
+      await setTimeout(1000);
+
+      // Check that the row was replicated to clickhouse
+      const queryRuns = clickhouse.reader.query({
+        name: "runs-replication",
+        query: "SELECT * FROM trigger_dev.task_runs_v1",
+        schema: z.any(),
+      });
+
+      const [queryError, result] = await queryRuns({});
+
+      expect(queryError).toBeNull();
+      expect(result?.length).toBe(1);
+      expect(result?.[0]).toEqual(
+        expect.objectContaining({
+          run_id: taskRun.id,
+          friendly_id: taskRun.friendlyId,
+          task_identifier: taskRun.taskIdentifier,
+          environment_id: runtimeEnvironment.id,
+          project_id: project.id,
+          organization_id: organization.id,
+          environment_type: "DEVELOPMENT",
+          engine: "V2",
+        })
+      );
+
+      await runsReplicationService.stop();
+    }
+  );
+});

internal-packages/replication/src/client.ts

@@ -50,6 +50,17 @@ export interface LogicalReplicationClientOptions {
    * The interval in ms to extend the leader lock (default: 10000)
    */
   leaderLockExtendIntervalMs?: number;
+
+  /**
+   * The number of times to retry acquiring the leader lock (default: 120)
+   */
+  leaderLockRetryCount?: number;
+
+  /**
+   * The interval in ms to retry acquiring the leader lock (default: 500)
+   */
+  leaderLockRetryIntervalMs?: number;
+
   /**
    * The interval in seconds to automatically acknowledge the last LSN if no ack has been sent (default: 10)
    */
@@ -83,6 +94,8 @@ export class LogicalReplicationClient {
   private lastAcknowledgedLsn: string | null = null;
   private leaderLockTimeoutMs: number;
   private leaderLockExtendIntervalMs: number;
+  private leaderLockRetryCount: number;
+  private leaderLockRetryIntervalMs: number;
   private leaderLockHeartbeatTimer: NodeJS.Timeout | null = null;
   private ackIntervalSeconds: number;
   private lastAckTimestamp: number = 0;
@@ -106,6 +119,8 @@ export class LogicalReplicationClient {
 
     this.leaderLockTimeoutMs = options.leaderLockTimeoutMs ?? 30000;
     this.leaderLockExtendIntervalMs = options.leaderLockExtendIntervalMs ?? 10000;
+    this.leaderLockRetryCount = options.leaderLockRetryCount ?? 120;
+    this.leaderLockRetryIntervalMs = options.leaderLockRetryIntervalMs ?? 500;
     this.ackIntervalSeconds = options.ackIntervalSeconds ?? 10;
 
     this.redis = createRedisClient(
@@ -544,9 +559,8 @@ export class LogicalReplicationClient {
         [`logical-replication-client:${this.options.name}`],
         this.leaderLockTimeoutMs,
         {
-          retryCount: 60,
-          retryDelay: 1000,
-          retryJitter: 100,
+          retryCount: this.leaderLockRetryCount,
+          retryDelay: this.leaderLockRetryIntervalMs,
         }
       );
     } catch (err) {