From 1b4dabb8ba4b8ee2a0d46147ad98740f6bd59265 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Mon, 30 Jun 2025 16:22:19 +0100 Subject: [PATCH 1/3] v4: current concurrency sweeper (#2206) * v4: current concurrency sweeper * Fix webapp tests * Ensure only a single instance performs concurrency sweeping by using redis-worker cron jobs * Improved the mark phase * Ensure cron jobs get rescheduled even if the handler throws an error * Better property names --- apps/webapp/app/env.server.ts | 5 + apps/webapp/app/v3/runEngine.server.ts | 14 +- apps/webapp/test/engine/triggerTask.test.ts | 1 + .../run-engine/src/engine/index.ts | 53 +- .../run-engine/src/engine/statuses.ts | 26 +- .../run-engine/src/engine/types.ts | 10 +- .../run-engine/src/run-queue/index.ts | 471 ++++++++++++++++-- .../run-engine/src/run-queue/keyProducer.ts | 11 +- .../tests/concurrencySweeper.test.ts | 171 +++++++ .../run-engine/src/run-queue/types.ts | 4 + packages/redis-worker/package.json | 3 +- packages/redis-worker/src/cron.test.ts | 130 +++++ packages/redis-worker/src/worker.ts | 143 ++++++ pnpm-lock.yaml | 3 + references/hello-world/trigger.config.ts | 2 +- 15 files changed, 986 insertions(+), 61 deletions(-) create mode 100644 internal-packages/run-engine/src/run-queue/tests/concurrencySweeper.test.ts create mode 100644 packages/redis-worker/src/cron.test.ts diff --git a/apps/webapp/app/env.server.ts b/apps/webapp/app/env.server.ts index 6f742faa2..05aed6f3d 100644 --- a/apps/webapp/app/env.server.ts +++ b/apps/webapp/app/env.server.ts @@ -428,6 +428,10 @@ const EnvironmentSchema = z.object({ RUN_ENGINE_PROCESS_WORKER_QUEUE_DEBOUNCE_MS: z.coerce.number().int().default(200), RUN_ENGINE_DEQUEUE_BLOCKING_TIMEOUT_SECONDS: z.coerce.number().int().default(10), RUN_ENGINE_MASTER_QUEUE_CONSUMERS_INTERVAL_MS: z.coerce.number().int().default(500), + RUN_ENGINE_CONCURRENCY_SWEEPER_SCAN_SCHEDULE: z.string().optional(), + RUN_ENGINE_CONCURRENCY_SWEEPER_PROCESS_MARKED_SCHEDULE: z.string().optional(), + RUN_ENGINE_CONCURRENCY_SWEEPER_SCAN_JITTER_IN_MS: z.coerce.number().int().optional(), + RUN_ENGINE_CONCURRENCY_SWEEPER_PROCESS_MARKED_JITTER_IN_MS: z.coerce.number().int().optional(), RUN_ENGINE_RUN_LOCK_DURATION: z.coerce.number().int().default(5000), RUN_ENGINE_RUN_LOCK_AUTOMATIC_EXTENSION_THRESHOLD: z.coerce.number().int().default(1000), @@ -593,6 +597,7 @@ const EnvironmentSchema = z.object({ RUN_ENGINE_WORKER_ENABLED: z.string().default("1"), RUN_ENGINE_WORKER_LOG_LEVEL: z.enum(["log", "error", "warn", "info", "debug"]).default("info"), + RUN_ENGINE_RUN_QUEUE_LOG_LEVEL: z.enum(["log", "error", "warn", "info", "debug"]).default("info"), /** How long should the presence ttl last */ DEV_PRESENCE_SSE_TIMEOUT: z.coerce.number().int().default(30_000), diff --git a/apps/webapp/app/v3/runEngine.server.ts b/apps/webapp/app/v3/runEngine.server.ts index 559c36398..6c9a11c2a 100644 --- a/apps/webapp/app/v3/runEngine.server.ts +++ b/apps/webapp/app/v3/runEngine.server.ts @@ -1,10 +1,10 @@ import { RunEngine } from "@internal/run-engine"; -import { defaultMachine } from "~/services/platform.v3.server"; -import { prisma } from "~/db.server"; +import { $replica, prisma } from "~/db.server"; import { env } from "~/env.server"; +import { defaultMachine } from "~/services/platform.v3.server"; import { singleton } from "~/utils/singleton"; import { allMachines } from "./machinePresets.server"; -import { tracer, meter } from "./tracer.server"; +import { meter, tracer } from "./tracer.server"; export const engine = singleton("RunEngine", createRunEngine); @@ -13,6 +13,7 @@ export type { RunEngine }; function createRunEngine() { const engine = new RunEngine({ prisma, + readOnlyPrisma: $replica, logLevel: env.RUN_ENGINE_WORKER_LOG_LEVEL, worker: { disabled: env.RUN_ENGINE_WORKER_ENABLED === "0", @@ -39,6 +40,7 @@ function createRunEngine() { }, queue: { defaultEnvConcurrency: env.DEFAULT_ENV_EXECUTION_CONCURRENCY_LIMIT, + logLevel: env.RUN_ENGINE_RUN_QUEUE_LOG_LEVEL, redis: { keyPrefix: "engine:", port: env.RUN_ENGINE_RUN_QUEUE_REDIS_PORT ?? undefined, @@ -64,6 +66,12 @@ function createRunEngine() { dequeueBlockingTimeoutSeconds: env.RUN_ENGINE_DEQUEUE_BLOCKING_TIMEOUT_SECONDS, masterQueueConsumersIntervalMs: env.RUN_ENGINE_MASTER_QUEUE_CONSUMERS_INTERVAL_MS, masterQueueConsumersDisabled: env.RUN_ENGINE_WORKER_ENABLED === "0", + concurrencySweeper: { + scanSchedule: env.RUN_ENGINE_CONCURRENCY_SWEEPER_SCAN_SCHEDULE, + processMarkedSchedule: env.RUN_ENGINE_CONCURRENCY_SWEEPER_PROCESS_MARKED_SCHEDULE, + scanJitterInMs: env.RUN_ENGINE_CONCURRENCY_SWEEPER_SCAN_JITTER_IN_MS, + processMarkedJitterInMs: env.RUN_ENGINE_CONCURRENCY_SWEEPER_PROCESS_MARKED_JITTER_IN_MS, + }, }, runLock: { redis: { diff --git a/apps/webapp/test/engine/triggerTask.test.ts b/apps/webapp/test/engine/triggerTask.test.ts index 8aa0fddb7..d90ae99b8 100644 --- a/apps/webapp/test/engine/triggerTask.test.ts +++ b/apps/webapp/test/engine/triggerTask.test.ts @@ -3,6 +3,7 @@ import { describe, expect, vi } from "vitest"; // Mock the db prisma client vi.mock("~/db.server", () => ({ prisma: {}, + $replica: {}, })); vi.mock("~/services/platform.v3.server", async (importOriginal) => { diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts index aca188cd2..8ed9febb3 100644 --- a/internal-packages/run-engine/src/engine/index.ts +++ b/internal-packages/run-engine/src/engine/index.ts @@ -16,6 +16,7 @@ import { Prisma, PrismaClient, PrismaClientOrTransaction, + PrismaReplicaClient, TaskRun, TaskRunExecutionSnapshot, Waitpoint, @@ -50,6 +51,7 @@ import { TtlSystem } from "./systems/ttlSystem.js"; import { WaitpointSystem } from "./systems/waitpointSystem.js"; import { EngineWorker, HeartbeatTimeouts, RunEngineOptions, TriggerParams } from "./types.js"; import { workerCatalog } from "./workerCatalog.js"; +import { getFinalRunStatuses, isFinalRunStatus } from "./statuses.js"; export class RunEngine { private runLockRedis: Redis; @@ -61,6 +63,7 @@ export class RunEngine { private heartbeatTimeouts: HeartbeatTimeouts; prisma: PrismaClient; + readOnlyPrisma: PrismaReplicaClient; runQueue: RunQueue; eventBus: EventBus = new EventEmitter(); executionSnapshotSystem: ExecutionSnapshotSystem; @@ -79,6 +82,7 @@ export class RunEngine { constructor(private readonly options: RunEngineOptions) { this.logger = options.logger ?? new Logger("RunEngine", this.options.logLevel ?? "info"); this.prisma = options.prisma; + this.readOnlyPrisma = options.readOnlyPrisma ?? this.prisma; this.runLockRedis = createRedisClient( { ...options.runLock.redis, @@ -123,7 +127,7 @@ export class RunEngine { defaultEnvConcurrencyLimit: options.queue?.defaultEnvConcurrency ?? 10, }), defaultEnvConcurrency: options.queue?.defaultEnvConcurrency ?? 10, - logger: new Logger("RunQueue", this.options.logLevel ?? "info"), + logger: new Logger("RunQueue", options.queue?.logLevel ?? "info"), redis: { ...options.queue.redis, keyPrefix: `${options.queue.redis.keyPrefix}runqueue:` }, retryOptions: options.queue?.retryOptions, workerOptions: { @@ -133,6 +137,13 @@ export class RunEngine { immediatePollIntervalMs: options.worker.immediatePollIntervalMs, shutdownTimeoutMs: options.worker.shutdownTimeoutMs, }, + concurrencySweeper: { + scanSchedule: options.queue?.concurrencySweeper?.scanSchedule, + processMarkedSchedule: options.queue?.concurrencySweeper?.processMarkedSchedule, + scanJitterInMs: options.queue?.concurrencySweeper?.scanJitterInMs, + processMarkedJitterInMs: options.queue?.concurrencySweeper?.processMarkedJitterInMs, + callback: this.#concurrencySweeperCallback.bind(this), + }, shardCount: options.queue?.shardCount, masterQueueConsumersDisabled: options.queue?.masterQueueConsumersDisabled, masterQueueConsumersIntervalMs: options.queue?.masterQueueConsumersIntervalMs, @@ -1329,4 +1340,44 @@ export class RunEngine { } }); } + + async #concurrencySweeperCallback( + runIds: string[] + ): Promise> { + const runs = await this.readOnlyPrisma.taskRun.findMany({ + where: { + id: { in: runIds }, + completedAt: { + lte: new Date(Date.now() - 1000 * 60 * 10), // This only finds runs that were completed more than 10 minutes ago + }, + organizationId: { + not: null, + }, + status: { + in: getFinalRunStatuses(), + }, + }, + select: { + id: true, + status: true, + organizationId: true, + }, + }); + + // Log the finished runs + for (const run of runs) { + this.logger.info("Concurrency sweeper callback found finished run", { + runId: run.id, + orgId: run.organizationId, + status: run.status, + }); + } + + return runs + .filter((run) => !!run.organizationId) + .map((run) => ({ + id: run.id, + orgId: run.organizationId!, + })); + } } diff --git a/internal-packages/run-engine/src/engine/statuses.ts b/internal-packages/run-engine/src/engine/statuses.ts index f8a66240d..36f082501 100644 --- a/internal-packages/run-engine/src/engine/statuses.ts +++ b/internal-packages/run-engine/src/engine/statuses.ts @@ -41,21 +41,25 @@ export function isInitialState(status: TaskRunExecutionStatus): boolean { return startedStatuses.includes(status); } -export function isFinalRunStatus(status: TaskRunStatus): boolean { - const finalStatuses: TaskRunStatus[] = [ - "CANCELED", - "INTERRUPTED", - "COMPLETED_SUCCESSFULLY", - "COMPLETED_WITH_ERRORS", - "SYSTEM_FAILURE", - "CRASHED", - "EXPIRED", - "TIMED_OUT", - ]; +const finalStatuses: TaskRunStatus[] = [ + "CANCELED", + "INTERRUPTED", + "COMPLETED_SUCCESSFULLY", + "COMPLETED_WITH_ERRORS", + "SYSTEM_FAILURE", + "CRASHED", + "EXPIRED", + "TIMED_OUT", +]; +export function isFinalRunStatus(status: TaskRunStatus): boolean { return finalStatuses.includes(status); } +export function getFinalRunStatuses(): TaskRunStatus[] { + return finalStatuses; +} + export function canReleaseConcurrency(status: TaskRunExecutionStatus): boolean { const releaseableStatuses: TaskRunExecutionStatus[] = ["SUSPENDED", "EXECUTING_WITH_WAITPOINTS"]; return releaseableStatuses.includes(status); diff --git a/internal-packages/run-engine/src/engine/types.ts b/internal-packages/run-engine/src/engine/types.ts index 286545c85..87612b2bf 100644 --- a/internal-packages/run-engine/src/engine/types.ts +++ b/internal-packages/run-engine/src/engine/types.ts @@ -8,7 +8,7 @@ import { RetryOptions, RunChainState, } from "@trigger.dev/core/v3"; -import { PrismaClient } from "@trigger.dev/database"; +import { PrismaClient, PrismaReplicaClient } from "@trigger.dev/database"; import { FairQueueSelectionStrategyOptions } from "../run-queue/fairQueueSelectionStrategy.js"; import { MinimalAuthenticatedEnvironment } from "../shared/index.js"; import { workerCatalog } from "./workerCatalog.js"; @@ -17,6 +17,7 @@ import { LockRetryConfig } from "./locking.js"; export type RunEngineOptions = { prisma: PrismaClient; + readOnlyPrisma?: PrismaReplicaClient; worker: { disabled?: boolean; redis: RedisOptions; @@ -38,11 +39,18 @@ export type RunEngineOptions = { workerOptions?: WorkerConcurrencyOptions; retryOptions?: RetryOptions; defaultEnvConcurrency?: number; + logLevel?: LogLevel; queueSelectionStrategyOptions?: Pick< FairQueueSelectionStrategyOptions, "parentQueueLimit" | "tracer" | "biases" | "reuseSnapshotCount" | "maximumEnvCount" >; dequeueBlockingTimeoutSeconds?: number; + concurrencySweeper?: { + scanSchedule?: string; + processMarkedSchedule?: string; + scanJitterInMs?: number; + processMarkedJitterInMs?: number; + }; }; runLock: { redis: RedisOptions; diff --git a/internal-packages/run-engine/src/run-queue/index.ts b/internal-packages/run-engine/src/run-queue/index.ts index f5772a9d6..0b6d5c1d1 100644 --- a/internal-packages/run-engine/src/run-queue/index.ts +++ b/internal-packages/run-engine/src/run-queue/index.ts @@ -36,11 +36,13 @@ import { type Result, } from "@internal/redis"; import { MessageNotFoundError } from "./errors.js"; -import { tryCatch } from "@trigger.dev/core"; +import { promiseWithResolvers, tryCatch } from "@trigger.dev/core"; import { setInterval } from "node:timers/promises"; import { nanoid } from "nanoid"; -import { Worker, type WorkerConcurrencyOptions } from "@trigger.dev/redis-worker"; +import { CronSchema, Worker, type WorkerConcurrencyOptions } from "@trigger.dev/redis-worker"; import { z } from "zod"; +import { Readable } from "node:stream"; +import { setTimeout } from "node:timers/promises"; const SemanticAttributes = { QUEUE: "runqueue.queue", @@ -78,14 +80,31 @@ export type RunQueueOptions = { }; meter?: Meter; dequeueBlockingTimeoutSeconds?: number; + concurrencySweeper?: { + scanSchedule?: string; + scanJitterInMs?: number; + processMarkedSchedule?: string; + processMarkedJitterInMs?: number; + callback: ConcurrencySweeperCallback; + }; }; +export interface ConcurrencySweeperCallback { + (runIds: string[]): Promise>; +} + type DequeuedMessage = { messageId: string; messageScore: string; message: OutputPayload; }; +type MarkedRun = { + orgId: string; + messageId: string; + score: number; +}; + const defaultRetrySettings = { maxAttempts: 12, factor: 2, @@ -102,6 +121,24 @@ const workerCatalog = { }), visibilityTimeoutMs: 30_000, }, + scanConcurrencySets: { + schema: CronSchema, + visibilityTimeoutMs: 60_000 * 5, + cron: "*/10 * * * *", + jitterInMs: 60_000, + retry: { + maxAttempts: 1, + }, + }, + processMarkedRuns: { + schema: CronSchema, + visibilityTimeoutMs: 60_000 * 5, + cron: "*/5 * * * *", + jitterInMs: 30_000, + retry: { + maxAttempts: 1, + }, + }, }; /** @@ -112,7 +149,7 @@ export class RunQueue { private subscriber: Redis; private luaDebugSubscriber: Redis; private logger: Logger; - private redis: Redis; + public redis: Redis; public keys: RunQueueKeyProducer; private queueSelectionStrategy: RunQueueSelectionStrategy; private shardCount: number; @@ -121,7 +158,7 @@ export class RunQueue { private _observableWorkerQueues: Set = new Set(); private _meter: Meter; - constructor(private readonly options: RunQueueOptions) { + constructor(public readonly options: RunQueueOptions) { this.shardCount = options.shardCount ?? 2; this.retryOptions = options.retryOptions ?? defaultRetrySettings; this.redis = createRedisClient(options.redis, { @@ -170,16 +207,40 @@ export class RunQueue { ...options.redis, keyPrefix: `${options.redis.keyPrefix}:worker`, }, - catalog: workerCatalog, + catalog: { + ...workerCatalog, + scanConcurrencySets: { + ...workerCatalog.scanConcurrencySets, + cron: options.concurrencySweeper?.scanSchedule ?? workerCatalog.scanConcurrencySets.cron, + jitter: + options.concurrencySweeper?.scanJitterInMs ?? + workerCatalog.scanConcurrencySets.jitterInMs, + }, + processMarkedRuns: { + ...workerCatalog.processMarkedRuns, + cron: + options.concurrencySweeper?.processMarkedSchedule ?? + workerCatalog.processMarkedRuns.cron, + jitterInMs: + options.concurrencySweeper?.processMarkedJitterInMs ?? + workerCatalog.processMarkedRuns.jitterInMs, + }, + }, concurrency: options.workerOptions?.concurrency, pollIntervalMs: options.workerOptions?.pollIntervalMs ?? 1000, immediatePollIntervalMs: options.workerOptions?.immediatePollIntervalMs ?? 100, shutdownTimeoutMs: options.workerOptions?.shutdownTimeoutMs ?? 10_000, - logger: new Logger("RunQueueWorker", options.logLevel ?? "log"), + logger: new Logger("RunQueueWorker", options.logLevel ?? "info"), jobs: { processQueueForWorkerQueue: async (job) => { await this.#processQueueForWorkerQueue(job.payload.queueKey, job.payload.environmentId); }, + scanConcurrencySets: async (job) => { + await this.scanConcurrencySets(); + }, + processMarkedRuns: async (job) => { + await this.processMarkedRuns(); + }, }, }); @@ -404,38 +465,7 @@ export class RunQueue { } public async readMessage(orgId: string, messageId: string) { - return this.#trace( - "readMessage", - async (span) => { - const rawMessage = await this.redis.get(this.keys.messageKey(orgId, messageId)); - - if (!rawMessage) { - return; - } - - const message = OutputPayload.safeParse(JSON.parse(rawMessage)); - - if (!message.success) { - this.logger.error(`[${this.name}] Failed to parse message`, { - messageId, - error: message.error, - service: this.name, - }); - - return; - } - - return message.data; - }, - { - attributes: { - [SEMATTRS_MESSAGING_OPERATION]: "receive", - [SEMATTRS_MESSAGE_ID]: messageId, - [SEMATTRS_MESSAGING_SYSTEM]: "marqs", - [SemanticAttributes.RUN_ID]: messageId, - }, - } - ); + return this.readMessageFromKey(this.keys.messageKey(orgId, messageId)); } public async readMessageFromKey(messageKey: string) { @@ -448,24 +478,34 @@ export class RunQueue { return; } - const message = OutputPayload.safeParse(JSON.parse(rawMessage)); + const deserializedMessage = safeJsonParse(rawMessage); + + const message = OutputPayload.safeParse(deserializedMessage); if (!message.success) { this.logger.error(`[${this.name}] Failed to parse message`, { messageKey, error: message.error, service: this.name, + deserializedMessage, }); - return; + return deserializedMessage as OutputPayload; } + span.setAttributes({ + [SemanticAttributes.QUEUE]: message.data.queue, + [SemanticAttributes.RUN_ID]: message.data.runId, + [SemanticAttributes.CONCURRENCY_KEY]: message.data.concurrencyKey, + [SemanticAttributes.WORKER_QUEUE]: this.#getWorkerQueueFromMessage(message.data), + }); + return message.data; }, { attributes: { [SEMATTRS_MESSAGING_OPERATION]: "receive", - [SEMATTRS_MESSAGING_SYSTEM]: "marqs", + [SEMATTRS_MESSAGING_SYSTEM]: "runqueue", }, } ); @@ -896,6 +936,35 @@ export class RunQueue { return await this.redis.lrange(workerQueueKey, 0, -1); } + /** + * Create a scan stream for queue current concurrency keys + */ + public currentConcurrencyScanStream( + count: number = 10, + onEnd?: () => void, + onError?: (error: Error) => void + ): { stream: Readable; redis: Redis } { + const pattern = this.keys.currentConcurrencySetKeyScanPattern(); + const stream = this.redis.scanStream({ + match: pattern, + count, + type: "set", + }); + + if (onEnd) { + stream.on("end", onEnd); + } + + if (onError) { + stream.on("error", onError); + } + + return { + stream, + redis: this.redis, + }; + } + private async handleRedriveMessage(channel: string, message: string) { try { const { runId, envId, projectId, orgId } = JSON.parse(message) as any; @@ -1651,6 +1720,249 @@ export class RunQueue { return blockingClient; } + // Call this every 10 minutes + private async scanConcurrencySets() { + if (this.abortController.signal.aborted) { + return; + } + + this.logger.debug("Scanning concurrency sets for completed runs"); + + const stats = { + streamCallbacks: 0, + processedKeys: 0, + }; + + const { promise, resolve, reject } = promiseWithResolvers(); + + const { stream, redis } = this.currentConcurrencyScanStream( + 10, + () => { + this.logger.debug("Concurrency scan stream closed", { stats }); + + resolve(stats); + }, + (error) => { + this.logger.error("Concurrency scan stream error", { + stats, + error: { + name: error.name, + message: error.message, + stack: error.stack, + }, + }); + + reject(error); + } + ); + + stream.on("data", async (keys: string[]) => { + if (!keys || keys.length === 0) { + return; + } + + stream.pause(); + + if (this.abortController.signal.aborted) { + stream.destroy(); + return; + } + + stats.streamCallbacks++; + + const uniqueKeys = Array.from(new Set(keys)).map((key) => + key.replace(redis.options.keyPrefix ?? "", "") + ); + + if (uniqueKeys.length === 0) { + stream.resume(); + return; + } + + this.logger.debug("Processing concurrency keys from stream", { + keys: uniqueKeys, + }); + + stats.processedKeys += uniqueKeys.length; + + await Promise.allSettled(uniqueKeys.map((key) => this.processConcurrencySet(key))).finally( + () => { + stream.resume(); + } + ); + }); + + return promise; + } + + private async processConcurrencySet(concurrencyKey: string) { + const stream = this.redis.sscanStream(concurrencyKey, { + count: 100, + }); + + const { promise, resolve, reject } = promiseWithResolvers(); + + stream.on("end", () => { + resolve(); + }); + + stream.on("error", (error) => { + this.logger.error("Error in sscanStream for concurrency set", { + concurrencyKey, + error, + }); + + reject(error); + }); + + stream.on("data", async (runIds: string[]) => { + stream.pause(); + + if (this.abortController.signal.aborted) { + stream.destroy(); + return; + } + + if (!runIds || runIds.length === 0) { + stream.resume(); + return; + } + + const deduplicatedRunIds = Array.from(new Set(runIds)); + + const [processError] = await tryCatch( + this.processCurrentConcurrencyRunIds(concurrencyKey, deduplicatedRunIds) + ); + + if (processError) { + this.logger.error("Error processing concurrency set", { + concurrencyKey, + runIds, + error: processError, + }); + } + + stream.resume(); + }); + + return promise; + } + + private async processCurrentConcurrencyRunIds(concurrencyKey: string, runIds: string[]) { + this.logger.debug(`Processing concurrency set with ${runIds.length} runs`, { + concurrencyKey, + runIds: runIds.slice(0, 5), // Log first 5 for debugging + }); + + // Call the callback to determine which runs are completed + const completedRuns = await this.options.concurrencySweeper?.callback(runIds); + + if (!completedRuns) { + this.logger.debug("No completed runs found in concurrency set", { concurrencyKey }); + return; + } + + if (completedRuns.length === 0) { + this.logger.debug("No completed runs found in concurrency set", { concurrencyKey }); + return; + } + + this.logger.debug(`Found ${completedRuns.length} completed runs to mark for ack`, { + concurrencyKey, + completedRunIds: completedRuns.map((r) => r.id).slice(0, 5), + }); + + // Mark the completed runs for acknowledgment + await this.markRunsForAck(completedRuns); + } + + private async markRunsForAck(completedRuns: Array<{ id: string; orgId: string }>) { + const markedForAckKey = this.keys.markedForAckKey(); + + // Prepare arguments: alternating orgId, messageId pairs + const args: Array = []; + for (const run of completedRuns) { + this.logger.info("Marking run for acknowledgment", { + orgId: run.orgId, + runId: run.id, + }); + + args.push(Date.now()); + args.push(`${run.orgId}:${run.id}`); + } + + const count = await this.redis.zadd(markedForAckKey, ...args); + + this.logger.debug(`Marked ${count} runs for acknowledgment`, { + markedForAckKey, + count, + }); + } + + // Call this every 5 minutes + private async processMarkedRuns() { + if (this.abortController.signal.aborted) { + return; + } + + try { + const markedForAckKey = this.keys.markedForAckKey(); + const results = await this.redis.getMarkedRunsForAck(markedForAckKey, "100"); + + if (results.length === 0) { + return; + } + + const markedRuns: MarkedRun[] = []; + + // Parse results: [orgId1, messageId1, score1, orgId2, messageId2, score2, ...] + for (let i = 0; i < results.length; i += 3) { + markedRuns.push({ + orgId: results[i], + messageId: results[i + 1], + score: Number(results[i + 2]), + }); + } + + this.logger.debug(`Processing ${markedRuns.length} marked runs for acknowledgment`, { + markedRuns: markedRuns, // Log first 3 for debugging + }); + + for (const run of markedRuns) { + const [processError] = await tryCatch(this.processMarkedRun(run)); + + if (processError) { + this.logger.error("Error processing marked run", { + error: processError, + orgId: run.orgId, + messageId: run.messageId, + }); + } + } + + const shouldProcessMoreRuns = (await this.redis.zcard(markedForAckKey)) > 0; + + if (shouldProcessMoreRuns) { + await setTimeout(1000); + await this.processMarkedRuns(); + } + } catch (error) { + this.logger.error("Error processing marked runs", { error }); + } + } + + async processMarkedRun(run: MarkedRun) { + this.logger.info("Acknowledging marked run", { + orgId: run.orgId, + messageId: run.messageId, + }); + + await this.acknowledgeMessage(run.orgId, run.messageId, { + skipDequeueProcessing: true, + removeFromWorkerQueue: false, + }); + } + #registerCommands() { this.redis.defineCommand("migrateLegacyMasterQueues", { numberOfKeys: 1, @@ -2020,6 +2332,77 @@ local envConcurrencyLimit = ARGV[1] redis.call('SET', envConcurrencyLimitKey, envConcurrencyLimit) `, }); + + this.redis.defineCommand("markCompletedRunsForAck", { + numberOfKeys: 1, + lua: ` +-- Keys: +local markedForAckKey = KEYS[1] + +-- Args: alternating orgId, messageId pairs +local currentTime = tonumber(redis.call('TIME')[1]) * 1000 + +for i = 1, #ARGV, 2 do + local orgId = ARGV[i] + local messageId = ARGV[i + 1] + local markedValue = orgId .. ':' .. messageId + + redis.call('ZADD', markedForAckKey, currentTime, markedValue) +end + +return #ARGV / 2 + `, + }); + + this.redis.defineCommand("getMarkedRunsForAck", { + numberOfKeys: 1, + lua: ` +-- Keys: +local markedForAckKey = KEYS[1] + +-- Args: +local maxCount = tonumber(ARGV[1] or '10') + +-- Get the oldest marked runs +local markedRuns = redis.call('ZRANGE', markedForAckKey, 0, maxCount - 1, 'WITHSCORES') + +local results = {} +for i = 1, #markedRuns, 2 do + local markedValue = markedRuns[i] + local score = markedRuns[i + 1] + + -- Parse orgId:messageId + local colonIndex = string.find(markedValue, ':') + if colonIndex then + local orgId = string.sub(markedValue, 1, colonIndex - 1) + local messageId = string.sub(markedValue, colonIndex + 1) + + table.insert(results, orgId) + table.insert(results, messageId) + table.insert(results, score) + end +end + +-- Remove the processed items +if #results > 0 then + local itemsToRemove = {} + for i = 1, #markedRuns, 2 do + table.insert(itemsToRemove, markedRuns[i]) + end + redis.call('ZREM', markedForAckKey, unpack(itemsToRemove)) +end + +return results + `, + }); + } +} + +function safeJsonParse(rawMessage: string): unknown { + try { + return JSON.parse(rawMessage); + } catch (e) { + return undefined; } } @@ -2145,5 +2528,11 @@ declare module "@internal/redis" { keyPrefix: string, ...queueNames: string[] ): Result; + + getMarkedRunsForAck( + markedForAckKey: string, + maxCount: string, + callback?: Callback + ): Result; } } diff --git a/internal-packages/run-engine/src/run-queue/keyProducer.ts b/internal-packages/run-engine/src/run-queue/keyProducer.ts index 49e165ad9..d2ac500f7 100644 --- a/internal-packages/run-engine/src/run-queue/keyProducer.ts +++ b/internal-packages/run-engine/src/run-queue/keyProducer.ts @@ -219,9 +219,16 @@ export class RunQueueFullKeyProducer implements RunQueueKeyProducer { } } deadLetterQueueKeyFromQueue(queue: string): string { - const descriptor = this.descriptorFromQueue(queue); + const { orgId, projectId, envId } = this.descriptorFromQueue(queue); + return this.deadLetterQueueKey({ orgId, projectId, envId }); + } - return this.deadLetterQueueKey(descriptor); + markedForAckKey(): string { + return "markedForAck"; + } + + currentConcurrencySetKeyScanPattern(): string { + return `*:${constants.ENV_PART}:*:queue:*:${constants.CURRENT_CONCURRENCY_PART}`; } descriptorFromQueue(queue: string): QueueDescriptor { diff --git a/internal-packages/run-engine/src/run-queue/tests/concurrencySweeper.test.ts b/internal-packages/run-engine/src/run-queue/tests/concurrencySweeper.test.ts new file mode 100644 index 000000000..97fef5303 --- /dev/null +++ b/internal-packages/run-engine/src/run-queue/tests/concurrencySweeper.test.ts @@ -0,0 +1,171 @@ +import { redisTest } from "@internal/testcontainers"; +import { trace } from "@internal/tracing"; +import { Logger } from "@trigger.dev/core/logger"; +import { describe } from "node:test"; +import { setTimeout } from "node:timers/promises"; +import { FairQueueSelectionStrategy } from "../fairQueueSelectionStrategy.js"; +import { RunQueue } from "../index.js"; +import { RunQueueFullKeyProducer } from "../keyProducer.js"; +import { InputPayload } from "../types.js"; + +const testOptions = { + name: "rq", + tracer: trace.getTracer("rq"), + workers: 1, + defaultEnvConcurrency: 25, + logger: new Logger("RunQueue", "warn"), + retryOptions: { + maxAttempts: 5, + factor: 1.1, + minTimeoutInMs: 100, + maxTimeoutInMs: 1_000, + randomize: true, + }, + keys: new RunQueueFullKeyProducer(), +}; + +const authenticatedEnvDev = { + id: "e1234", + type: "DEVELOPMENT" as const, + maximumConcurrencyLimit: 10, + project: { id: "p1234" }, + organization: { id: "o1234" }, +}; + +const messageDev: InputPayload = { + runId: "r4321", + taskIdentifier: "task/my-task", + orgId: "o1234", + projectId: "p1234", + environmentId: "e4321", + environmentType: "DEVELOPMENT", + queue: "task/my-task", + timestamp: Date.now(), + attempt: 0, +}; + +const messageDev2: InputPayload = { + ...messageDev, + runId: "r4322", +}; + +vi.setConfig({ testTimeout: 60_000 }); + +describe("RunQueue Concurrency Sweeper", () => { + redisTest( + "should process queue current concurrency sets and mark runs for ack if they are completed", + async ({ redisContainer }) => { + let enableConcurrencySweeper = false; + + const queue = new RunQueue({ + ...testOptions, + queueSelectionStrategy: new FairQueueSelectionStrategy({ + redis: { + keyPrefix: "runqueue:test:", + host: redisContainer.getHost(), + port: redisContainer.getPort(), + }, + keys: testOptions.keys, + }), + redis: { + keyPrefix: "runqueue:test:", + host: redisContainer.getHost(), + port: redisContainer.getPort(), + }, + concurrencySweeper: { + scanSchedule: "* * * * * *", // Every second + scanJitter: 5, + processMarkedSchedule: "* * * * * *", // Every second + processMarkedJitter: 5, + callback: async (runIds) => { + if (!enableConcurrencySweeper) { + return []; + } + + return [{ id: messageDev.runId, orgId: "o1234" }]; + }, + }, + }); + + try { + //enqueue message + const enqueueResult = await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageDev, + workerQueue: authenticatedEnvDev.id, + }); + + expect(enqueueResult).toBe(undefined); + + const enqueueResult2 = await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageDev2, + workerQueue: authenticatedEnvDev.id, + }); + + expect(enqueueResult2).toBe(undefined); + + //queue length + const result2 = await queue.lengthOfQueue(authenticatedEnvDev, messageDev.queue); + expect(result2).toBe(2); + + const envQueueLength2 = await queue.lengthOfEnvQueue(authenticatedEnvDev); + expect(envQueueLength2).toBe(2); + + //concurrencies + const queueConcurrency = await queue.currentConcurrencyOfQueue( + authenticatedEnvDev, + messageDev.queue + ); + expect(queueConcurrency).toBe(0); + + const envConcurrency = await queue.currentConcurrencyOfEnvironment(authenticatedEnvDev); + expect(envConcurrency).toBe(0); + + //dequeue message + const dequeued = await queue.dequeueMessageFromWorkerQueue( + "test_12345", + authenticatedEnvDev.id + ); + + expect(dequeued).toBeDefined(); + expect(dequeued?.messageId).toEqual(messageDev.runId); + + const dequeued2 = await queue.dequeueMessageFromWorkerQueue( + "test_12345", + authenticatedEnvDev.id + ); + expect(dequeued2).toBeDefined(); + expect(dequeued2?.messageId).toEqual(messageDev2.runId); + + // queue concurrency should be 2 + const queueConcurrency2 = await queue.currentConcurrencyOfQueue( + authenticatedEnvDev, + messageDev.queue + ); + expect(queueConcurrency2).toBe(2); + + // env concurrency should be 2 + const envConcurrency2 = await queue.currentConcurrencyOfEnvironment(authenticatedEnvDev); + expect(envConcurrency2).toBe(2); + + enableConcurrencySweeper = true; + + await setTimeout(3_000); // Now a run is "completed" and should be removed from the concurrency set + + // queue concurrency should be 0 + const queueConcurrency3 = await queue.currentConcurrencyOfQueue( + authenticatedEnvDev, + messageDev.queue + ); + expect(queueConcurrency3).toBe(1); + + // env concurrency should be 1 + const envConcurrency3 = await queue.currentConcurrencyOfEnvironment(authenticatedEnvDev); + expect(envConcurrency3).toBe(1); + } finally { + await queue.quit(); + } + } + ); +}); diff --git a/internal-packages/run-engine/src/run-queue/types.ts b/internal-packages/run-engine/src/run-queue/types.ts index 68431f4ed..38420ce89 100644 --- a/internal-packages/run-engine/src/run-queue/types.ts +++ b/internal-packages/run-engine/src/run-queue/types.ts @@ -96,6 +96,10 @@ export interface RunQueueKeyProducer { deadLetterQueueKey(env: MinimalAuthenticatedEnvironment): string; deadLetterQueueKey(env: EnvDescriptor): string; deadLetterQueueKeyFromQueue(queue: string): string; + + // Concurrency sweeper methods + markedForAckKey(): string; + currentConcurrencySetKeyScanPattern(): string; } export type EnvQueues = { diff --git a/packages/redis-worker/package.json b/packages/redis-worker/package.json index e572b095f..2f64432e1 100644 --- a/packages/redis-worker/package.json +++ b/packages/redis-worker/package.json @@ -27,7 +27,8 @@ "lodash.omit": "^4.5.0", "nanoid": "^5.0.7", "p-limit": "^6.2.0", - "zod": "3.23.8" + "zod": "3.23.8", + "cron-parser": "^4.9.0" }, "devDependencies": { "@internal/redis": "workspace:*", diff --git a/packages/redis-worker/src/cron.test.ts b/packages/redis-worker/src/cron.test.ts new file mode 100644 index 000000000..9a607234a --- /dev/null +++ b/packages/redis-worker/src/cron.test.ts @@ -0,0 +1,130 @@ +import { redisTest } from "@internal/testcontainers"; +import { Logger } from "@trigger.dev/core/logger"; +import { describe } from "node:test"; +import { expect } from "vitest"; +import { Worker, CronSchema } from "./worker.js"; +import { setTimeout } from "node:timers/promises"; + +describe("Worker with cron", () => { + redisTest( + "process items on the cron schedule", + { timeout: 180_000 }, + async ({ redisContainer }) => { + const processedItems: CronSchema[] = []; + const worker = new Worker({ + name: "test-worker", + redisOptions: { + host: redisContainer.getHost(), + port: redisContainer.getPort(), + password: redisContainer.getPassword(), + }, + catalog: { + cronJob: { + cron: "*/5 * * * * *", // Every 5 seconds + schema: CronSchema, + visibilityTimeoutMs: 5000, + retry: { maxAttempts: 3 }, + jitterInMs: 100, + }, + }, + jobs: { + cronJob: async ({ payload }) => { + await setTimeout(30); // Simulate work + processedItems.push(payload); + }, + }, + concurrency: { + workers: 2, + tasksPerWorker: 3, + }, + logger: new Logger("test", "debug"), + }).start(); + + await setTimeout(6_000); + + expect(processedItems.length).toBe(1); + + const firstItem = processedItems[0]; + + expect(firstItem?.timestamp).toBeGreaterThan(0); + expect(firstItem?.lastTimestamp).toBeUndefined(); + expect(firstItem?.cron).toBe("*/5 * * * * *"); + + await setTimeout(6_000); + + expect(processedItems.length).toBeGreaterThanOrEqual(2); + + const secondItem = processedItems[1]; + expect(secondItem?.timestamp).toBeGreaterThan(firstItem!.timestamp); + expect(secondItem?.lastTimestamp).toBe(firstItem?.timestamp); + expect(secondItem?.cron).toBe("*/5 * * * * *"); + + await worker.stop(); + } + ); + + redisTest( + "continues processing cron items even when job handler throws errors", + { timeout: 180_000 }, + async ({ redisContainer }) => { + const processedItems: CronSchema[] = []; + let executionCount = 0; + + const worker = new Worker({ + name: "test-worker-error", + redisOptions: { + host: redisContainer.getHost(), + port: redisContainer.getPort(), + password: redisContainer.getPassword(), + }, + catalog: { + cronJob: { + cron: "*/3 * * * * *", // Every 3 seconds + schema: CronSchema, + visibilityTimeoutMs: 5000, + retry: { maxAttempts: 1 }, // Only try once to fail faster + jitterInMs: 100, + }, + }, + jobs: { + cronJob: async ({ payload }) => { + executionCount++; + await setTimeout(30); // Simulate work + + // Throw error on first and third execution + if (executionCount === 1 || executionCount === 3) { + throw new Error(`Simulated error on execution ${executionCount}`); + } + + processedItems.push(payload); + }, + }, + concurrency: { + workers: 2, + tasksPerWorker: 3, + }, + logger: new Logger("test", "debug"), + }).start(); + + // Wait long enough for 4 executions (12 seconds + buffer) + await setTimeout(14_000); + + // Should have at least 4 executions total + expect(executionCount).toBeGreaterThanOrEqual(4); + + // Should have 2 successful items (executions 2 and 4) + expect(processedItems.length).toBeGreaterThanOrEqual(2); + + // Verify that some executions failed (execution count > successful count) + // This proves that errors occurred but cron scheduling continued + expect(executionCount).toBeGreaterThan(processedItems.length); + + // Verify that successful executions still have correct structure + const firstSuccessful = processedItems[0]; + expect(firstSuccessful?.timestamp).toBeGreaterThan(0); + expect(firstSuccessful?.cron).toBe("*/3 * * * * *"); + + await worker.stop(); + } + ); +}); diff --git a/packages/redis-worker/src/worker.ts b/packages/redis-worker/src/worker.ts index 69a1bbf8b..3542f305f 100644 --- a/packages/redis-worker/src/worker.ts +++ b/packages/redis-worker/src/worker.ts @@ -19,15 +19,28 @@ import { nanoid } from "nanoid"; import pLimit from "p-limit"; import { z } from "zod"; import { AnyQueueItem, SimpleQueue } from "./queue.js"; +import { parseExpression } from "cron-parser"; + +export const CronSchema = z.object({ + cron: z.string(), + lastTimestamp: z.number().optional(), + timestamp: z.number(), +}); + +export type CronSchema = z.infer; export type WorkerCatalog = { [key: string]: { schema: z.ZodFirstPartySchemaTypes | z.ZodDiscriminatedUnion; visibilityTimeoutMs: number; retry?: RetryOptions; + cron?: string; + jitterInMs?: number; }; }; +type WorkerCatalogItem = WorkerCatalog[keyof WorkerCatalog]; + type QueueCatalogFromWorkerCatalog = { [K in keyof Catalog]: Catalog[K]["schema"]; }; @@ -204,6 +217,12 @@ class Worker { public start() { const { workers, tasksPerWorker } = this.concurrency; + this.logger.info("Starting worker", { + workers, + tasksPerWorker, + concurrency: this.concurrency, + }); + // Launch a number of "worker loops" on the main thread. for (let i = 0; i < workers; i++) { this.workerLoops.push(this.runWorkerLoop(`worker-${nanoid(12)}`, tasksPerWorker, i, workers)); @@ -219,7 +238,9 @@ class Worker { }); }, }); + this.setupSubscriber(); + this.setupCron(); return this; } @@ -496,6 +517,11 @@ class Worker { return; } + if (!catalogItem) { + this.logger.error(`No catalog item found for job type: ${job}`); + return; + } + await startSpan( this.tracer, "processItem", @@ -513,6 +539,10 @@ class Worker { // On success, acknowledge the item. await this.queue.ack(id, deduplicationKey); + + if (catalogItem.cron) { + await this.rescheduleCronJob(job, catalogItem, item); + } }, { kind: SpanKind.CONSUMER, @@ -560,7 +590,13 @@ class Worker { attempt: newAttempt, errorMessage, }); + await this.queue.moveToDeadLetterQueue(id, errorMessage); + + if (catalogItem.cron) { + await this.rescheduleCronJob(job, catalogItem, item); + } + return; } @@ -622,6 +658,113 @@ class Worker { return new Promise((resolve) => setTimeout(resolve, ms)); } + private setupCron() { + const cronJobs = Object.entries(this.options.catalog).filter(([_, value]) => value.cron); + + if (cronJobs.length === 0) { + return; + } + + this.logger.info("Setting up cron jobs", { + cronJobs: cronJobs.map(([job, value]) => ({ + job, + cron: value.cron, + jitterInMs: value.jitterInMs, + })), + }); + + // For each cron job, we need to try and enqueue a job with the next timestamp of the cron job. + const enqueuePromises = cronJobs.map(([job, value]) => + this.enqueueCronJob(value.cron!, job, value.jitterInMs) + ); + + Promise.allSettled(enqueuePromises).then((results) => { + results.forEach((result) => { + if (result.status === "fulfilled") { + this.logger.info("Enqueued cron job", { result: result.value }); + } else { + this.logger.error("Failed to enqueue cron job", { reason: result.reason }); + } + }); + }); + } + + private async enqueueCronJob(cron: string, job: string, jitter?: number, lastTimestamp?: Date) { + const scheduledAt = this.calculateNextScheduledAt(cron, lastTimestamp); + const identifier = [job, this.timestampIdentifier(scheduledAt)].join(":"); + // Calculate the availableAt date by calculating a random number between -jitter/2 and jitter/2 and adding it to the scheduledAt + const availableAt = jitter + ? new Date(scheduledAt.getTime() + Math.random() * jitter - jitter / 2) + : scheduledAt; + + const enqueued = await this.enqueueOnce({ + id: identifier, + job, + payload: { + timestamp: scheduledAt.getTime(), + lastTimestamp: lastTimestamp?.getTime(), + cron, + }, + availableAt, + }); + + this.logger.info("Enqueued cron job", { + identifier, + cron, + job, + scheduledAt, + enqueued, + availableAt, + }); + + return { + identifier, + cron, + job, + scheduledAt, + enqueued, + }; + } + + private async rescheduleCronJob(job: string, catalogItem: WorkerCatalogItem, item: CronSchema) { + if (!catalogItem.cron) { + return; + } + + return this.enqueueCronJob( + catalogItem.cron, + job, + catalogItem.jitterInMs, + new Date(item.timestamp) + ); + } + + private calculateNextScheduledAt(cron: string, lastTimestamp?: Date): Date { + const scheduledAt = parseExpression(cron, { + currentDate: lastTimestamp, + }) + .next() + .toDate(); + + // If scheduledAt is in the past, we should just calculate the next one based on the current time + if (scheduledAt < new Date()) { + return this.calculateNextScheduledAt(cron); + } + + return scheduledAt; + } + + private timestampIdentifier(timestamp: Date) { + const year = timestamp.getUTCFullYear(); + const month = timestamp.getUTCMonth(); + const day = timestamp.getUTCDate(); + const hour = timestamp.getUTCHours(); + const minute = timestamp.getUTCMinutes(); + const second = timestamp.getUTCSeconds(); + + return `${year}-${month}-${day}-${hour}-${minute}-${second}`; + } + private setupSubscriber() { const channel = `${this.options.name}:redrive`; this.subscriber?.subscribe(channel, (err) => { diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index be39db842..096368205 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1660,6 +1660,9 @@ importers: '@trigger.dev/core': specifier: workspace:4.0.0-v4-beta.21 version: link:../core + cron-parser: + specifier: ^4.9.0 + version: 4.9.0 lodash.omit: specifier: ^4.5.0 version: 4.5.0 diff --git a/references/hello-world/trigger.config.ts b/references/hello-world/trigger.config.ts index fde795ccc..c3c6aea9e 100644 --- a/references/hello-world/trigger.config.ts +++ b/references/hello-world/trigger.config.ts @@ -5,7 +5,7 @@ export default defineConfig({ compatibilityFlags: ["run_engine_v2"], project: "proj_rrkpdguyagvsoktglnod", logLevel: "log", - maxDuration: 60, + maxDuration: 3600, retries: { enabledInDev: true, default: { From f38d35e9f6a907bc72c624035d6c7df0b1f1b5e7 Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Mon, 30 Jun 2025 16:32:02 +0100 Subject: [PATCH 2/3] ClickHouse replication improvements (retrying, strip bad unicode chars) (#2205) * Add retry logic for insert operations Add a generic retry mechanism for task run and payload inserts to handle transient connection errors. The new #insertWithRetry method retries up to three times with exponential backoff and jitter on retryable connection errors such as connection resets or timeouts. Errors are logged and recorded in tracing spans to improve observability and robustness of the replication service. * Replication settings are configurable * Log out the runIds for failed batches * Detecting bad JSON in run replication and ignoring it * Reproduced split unicode error * Move output file * Massively improved the performance * Minor performance improvements * Unskip tests * Remove unused test in CH package * Fix for the ClickHouse UI explorer * RunReplication keepAlive defaults to false * Add concurrency_key and bulk_action_group_ids to ClickHouse task runs * ClickHouse package doesn't need to be built anymore for the webapp * Set the concurrency_key from the run replication service --- apps/webapp/app/env.server.ts | 6 +- .../runsReplicationInstance.server.ts | 3 + .../services/runsReplicationService.server.ts | 124 +++++++++++- apps/webapp/app/utils/detectBadJsonStrings.ts | 68 +++++++ apps/webapp/test/bad-clickhouse-output.json | 3 + apps/webapp/test/detectbadJsonStrings.test.ts | 190 ++++++++++++++++++ .../test/runsReplicationService.part2.test.ts | 178 +++++++++++++++- docker/docker-compose.yml | 2 +- internal-packages/clickhouse/package.json | 12 +- ...d_task_runs_v2_concurrency_bulkactions.sql | 12 ++ .../clickhouse/src/taskRuns.test.ts | 6 + internal-packages/clickhouse/src/taskRuns.ts | 2 + 12 files changed, 598 insertions(+), 8 deletions(-) create mode 100644 apps/webapp/app/utils/detectBadJsonStrings.ts create mode 100644 apps/webapp/test/bad-clickhouse-output.json create mode 100644 apps/webapp/test/detectbadJsonStrings.test.ts create mode 100644 internal-packages/clickhouse/schema/005_add_task_runs_v2_concurrency_bulkactions.sql diff --git a/apps/webapp/app/env.server.ts b/apps/webapp/app/env.server.ts index 05aed6f3d..da99b394a 100644 --- a/apps/webapp/app/env.server.ts +++ b/apps/webapp/app/env.server.ts @@ -842,9 +842,13 @@ const EnvironmentSchema = z.object({ RUN_REPLICATION_LEADER_LOCK_ADDITIONAL_TIME_MS: z.coerce.number().int().default(10_000), RUN_REPLICATION_LEADER_LOCK_RETRY_INTERVAL_MS: z.coerce.number().int().default(500), RUN_REPLICATION_WAIT_FOR_ASYNC_INSERT: z.string().default("0"), - RUN_REPLICATION_KEEP_ALIVE_ENABLED: z.string().default("1"), + RUN_REPLICATION_KEEP_ALIVE_ENABLED: z.string().default("0"), RUN_REPLICATION_KEEP_ALIVE_IDLE_SOCKET_TTL_MS: z.coerce.number().int().optional(), RUN_REPLICATION_MAX_OPEN_CONNECTIONS: z.coerce.number().int().default(10), + // Retry configuration for insert operations + RUN_REPLICATION_INSERT_MAX_RETRIES: z.coerce.number().int().default(3), + RUN_REPLICATION_INSERT_BASE_DELAY_MS: z.coerce.number().int().default(100), + RUN_REPLICATION_INSERT_MAX_DELAY_MS: z.coerce.number().int().default(2000), // Clickhouse CLICKHOUSE_URL: z.string().optional(), diff --git a/apps/webapp/app/services/runsReplicationInstance.server.ts b/apps/webapp/app/services/runsReplicationInstance.server.ts index 8a8b0df65..86b17601a 100644 --- a/apps/webapp/app/services/runsReplicationInstance.server.ts +++ b/apps/webapp/app/services/runsReplicationInstance.server.ts @@ -62,6 +62,9 @@ function initializeRunsReplicationInstance() { logLevel: env.RUN_REPLICATION_LOG_LEVEL, waitForAsyncInsert: env.RUN_REPLICATION_WAIT_FOR_ASYNC_INSERT === "1", tracer: provider.getTracer("runs-replication-service"), + insertMaxRetries: env.RUN_REPLICATION_INSERT_MAX_RETRIES, + insertBaseDelayMs: env.RUN_REPLICATION_INSERT_BASE_DELAY_MS, + insertMaxDelayMs: env.RUN_REPLICATION_INSERT_MAX_DELAY_MS, }); if (env.RUN_REPLICATION_ENABLED === "1") { diff --git a/apps/webapp/app/services/runsReplicationService.server.ts b/apps/webapp/app/services/runsReplicationService.server.ts index eed1d3c9f..97ceb23ea 100644 --- a/apps/webapp/app/services/runsReplicationService.server.ts +++ b/apps/webapp/app/services/runsReplicationService.server.ts @@ -15,6 +15,8 @@ import { TaskRun } from "@trigger.dev/database"; import { nanoid } from "nanoid"; import EventEmitter from "node:events"; import pLimit from "p-limit"; +import { logger } from "./logger.server"; +import { detectBadJsonStrings } from "~/utils/detectBadJsonStrings"; interface TransactionEvent { tag: "insert" | "update" | "delete"; @@ -51,6 +53,10 @@ export type RunsReplicationServiceOptions = { logLevel?: LogLevel; tracer?: Tracer; waitForAsyncInsert?: boolean; + // Retry configuration for insert operations + insertMaxRetries?: number; + insertBaseDelayMs?: number; + insertMaxDelayMs?: number; }; type TaskRunInsert = { _version: bigint; run: TaskRun; event: "insert" | "update" | "delete" }; @@ -80,6 +86,10 @@ export class RunsReplicationService { private _latestCommitEndLsn: string | null = null; private _lastAcknowledgedLsn: string | null = null; private _acknowledgeInterval: NodeJS.Timeout | null = null; + // Retry configuration + private _insertMaxRetries: number; + private _insertBaseDelayMs: number; + private _insertMaxDelayMs: number; public readonly events: EventEmitter; @@ -151,6 +161,11 @@ export class RunsReplicationService { this._replicationClient.events.on("leaderElection", (isLeader) => { this.logger.info("Leader election", { isLeader }); }); + + // Initialize retry configuration + this._insertMaxRetries = options.insertMaxRetries ?? 3; + this._insertBaseDelayMs = options.insertBaseDelayMs ?? 100; + this._insertMaxDelayMs = options.insertMaxDelayMs ?? 2000; } public async shutdown() { @@ -445,8 +460,37 @@ export class RunsReplicationService { payloadInserts: payloadInserts.length, }); - await this.#insertTaskRunInserts(taskRunInserts); - await this.#insertPayloadInserts(payloadInserts); + // Insert task runs and payloads with retry logic for connection errors + const [taskRunError, taskRunResult] = await this.#insertWithRetry( + () => this.#insertTaskRunInserts(taskRunInserts), + "task run inserts", + flushId + ); + + const [payloadError, payloadResult] = await this.#insertWithRetry( + () => this.#insertPayloadInserts(payloadInserts), + "payload inserts", + flushId + ); + + // Log any errors that occurred + if (taskRunError) { + this.logger.error("Error inserting task run inserts", { + error: taskRunError, + flushId, + runIds: taskRunInserts.map((r) => r.run_id), + }); + recordSpanError(span, taskRunError); + } + + if (payloadError) { + this.logger.error("Error inserting payload inserts", { + error: payloadError, + flushId, + runIds: payloadInserts.map((r) => r.run_id), + }); + recordSpanError(span, payloadError); + } this.logger.debug("Flushed inserts", { flushId, @@ -456,6 +500,73 @@ export class RunsReplicationService { }); } + // New method to handle inserts with retry logic for connection errors + async #insertWithRetry( + insertFn: () => Promise, + operationName: string, + flushId: string + ): Promise<[Error | null, T | null]> { + let lastError: Error | null = null; + + for (let attempt = 1; attempt <= this._insertMaxRetries; attempt++) { + try { + const result = await insertFn(); + return [null, result]; + } catch (error) { + lastError = error instanceof Error ? error : new Error(String(error)); + + // Check if this is a retryable connection error + if (this.#isRetryableConnectionError(lastError) && attempt < this._insertMaxRetries) { + const delay = this.#calculateConnectionRetryDelay(attempt); + + this.logger.warn(`Retrying ${operationName} due to connection error`, { + flushId, + attempt, + maxRetries: this._insertMaxRetries, + error: lastError.message, + delay, + }); + + await new Promise((resolve) => setTimeout(resolve, delay)); + continue; + } + break; + } + } + + return [lastError, null]; + } + + // New method to check if an error is a retryable connection error + #isRetryableConnectionError(error: Error): boolean { + const errorMessage = error.message.toLowerCase(); + const retryableConnectionPatterns = [ + "socket hang up", + "econnreset", + "connection reset", + "connection refused", + "connection timeout", + "network error", + "read econnreset", + "write econnreset", + ]; + + return retryableConnectionPatterns.some((pattern) => errorMessage.includes(pattern)); + } + + // New method to calculate retry delay for connection errors + #calculateConnectionRetryDelay(attempt: number): number { + // Exponential backoff: baseDelay, baseDelay*2, baseDelay*4, etc. + const delay = Math.min( + this._insertBaseDelayMs * Math.pow(2, attempt - 1), + this._insertMaxDelayMs + ); + + // Add some jitter to prevent thundering herd + const jitter = Math.random() * 100; + return delay + jitter; + } + async #insertTaskRunInserts(taskRunInserts: TaskRunV2[]) { return await startSpan(this._tracer, "insertTaskRunsInserts", async (span) => { const [insertError, insertResult] = await this.options.clickhouse.taskRuns.insert( @@ -604,6 +715,7 @@ export class RunsReplicationService { idempotency_key: run.idempotencyKey ?? "", expiration_ttl: run.ttl ?? "", output, + concurrency_key: run.concurrencyKey ?? "", _version: _version.toString(), _is_deleted: event === "delete" ? 1 : 0, }; @@ -631,6 +743,14 @@ export class RunsReplicationService { return { data: undefined }; } + if (detectBadJsonStrings(data)) { + this.logger.warn("Detected bad JSON strings", { + data, + dataType, + }); + return { data: undefined }; + } + const packet = { data, dataType, diff --git a/apps/webapp/app/utils/detectBadJsonStrings.ts b/apps/webapp/app/utils/detectBadJsonStrings.ts new file mode 100644 index 000000000..4a000b542 --- /dev/null +++ b/apps/webapp/app/utils/detectBadJsonStrings.ts @@ -0,0 +1,68 @@ +export function detectBadJsonStrings(jsonString: string): boolean { + // Fast path: skip everything if no \u + let idx = jsonString.indexOf("\\u"); + if (idx === -1) return false; + + // Use a more efficient scanning strategy + const length = jsonString.length; + + while (idx !== -1 && idx < length - 5) { + // Only check if we have enough characters left + if (idx + 6 > length) break; + + if (jsonString[idx + 1] === "u" && jsonString[idx + 2] === "d") { + const third = jsonString[idx + 3]; + + // High surrogate check + if ( + /[89ab]/.test(third) && + /[0-9a-f]/.test(jsonString[idx + 4]) && + /[0-9a-f]/.test(jsonString[idx + 5]) + ) { + // Check for low surrogate after (need at least 6 more chars) + if (idx + 12 > length) { + return true; // Incomplete high surrogate (not enough chars left) + } + + if ( + jsonString[idx + 6] !== "\\" || + jsonString[idx + 7] !== "u" || + jsonString[idx + 8] !== "d" || + !/[cd]/.test(jsonString[idx + 9]) || + !/[0-9a-f]/.test(jsonString[idx + 10]) || + !/[0-9a-f]/.test(jsonString[idx + 11]) + ) { + return true; // Incomplete high surrogate + } + } + + // Low surrogate check + if ( + (third === "c" || third === "d") && + /[0-9a-f]/.test(jsonString[idx + 4]) && + /[0-9a-f]/.test(jsonString[idx + 5]) + ) { + // Check for high surrogate before (need at least 6 chars before) + if (idx < 6) { + return true; // Incomplete low surrogate (not enough chars before) + } + + if ( + jsonString[idx - 6] !== "\\" || + jsonString[idx - 5] !== "u" || + jsonString[idx - 4] !== "d" || + !/[89ab]/.test(jsonString[idx - 3]) || + !/[0-9a-f]/.test(jsonString[idx - 2]) || + !/[0-9a-f]/.test(jsonString[idx - 1]) + ) { + return true; // Incomplete low surrogate + } + } + } + + // More efficient next search - skip ahead by 2 to avoid overlapping matches + idx = jsonString.indexOf("\\u", idx + 2); + } + + return false; +} diff --git a/apps/webapp/test/bad-clickhouse-output.json b/apps/webapp/test/bad-clickhouse-output.json new file mode 100644 index 000000000..8900b56fe --- /dev/null +++ b/apps/webapp/test/bad-clickhouse-output.json @@ -0,0 +1,3 @@ +{ + "title": "❜ 𝐒 𝐏𝗈𝗌𝗍 . . . 𝐍ð–ū𝗐 𝐂𝗈𝗇𝗍ð–ū𝗇𝗍 ꒰ ⚔ïļ ꒱ 𝐒𝐋 ❜ 𝐔𝐋\n\n꒰ âĪïļ ꒱ 𓃊 𝐋ð—ēð—Ū𝘃ð—ē 𝖚 ð—đð—ķð—ļð—ē 𝖚𝗇\ud835" +} diff --git a/apps/webapp/test/detectbadJsonStrings.test.ts b/apps/webapp/test/detectbadJsonStrings.test.ts new file mode 100644 index 000000000..7d14bf4ae --- /dev/null +++ b/apps/webapp/test/detectbadJsonStrings.test.ts @@ -0,0 +1,190 @@ +import { describe, expect, it } from "vitest"; +import { detectBadJsonStrings } from "~/utils/detectBadJsonStrings"; + +describe("detectBadJsonStrings", () => { + it("should not detect valid JSON string", () => { + const goodJson = `{"title": "hello"}`; + const result = detectBadJsonStrings(goodJson); + expect(result).toBe(false); + }); + + it("should detect incomplete Unicode escape sequences", () => { + const badJson = `{"title": "hello\\ud835"}`; + const result = detectBadJsonStrings(badJson); + expect(result).toBe(true); + }); + + it("should not detect complete Unicode escape sequences", () => { + const goodJson = `{"title": "hello\\ud835\\udc00"}`; + const result = detectBadJsonStrings(goodJson); + expect(result).toBe(false); + }); + + it("should detect incomplete low surrogate", () => { + const badJson = `{"title": "hello\\udc00"}`; + const result = detectBadJsonStrings(badJson); + expect(result).toBe(true); + }); + + it("should handle multiple Unicode sequences correctly", () => { + const goodJson = `{"title": "hello\\ud835\\udc00\\ud835\\udc01"}`; + const result = detectBadJsonStrings(goodJson); + expect(result).toBe(false); + }); + + it("should detect mixed complete and incomplete sequences", () => { + const badJson = `{"title": "hello\\ud835\\udc00\\ud835"}`; + const result = detectBadJsonStrings(badJson); + expect(result).toBe(true); + }); + + it("should have acceptable performance overhead", () => { + const longText = `hello world `.repeat(1_000); + const goodJson = `{"title": "hello", "text": "${longText}"}`; + const badJson = `{"title": "hello\\ud835", "text": "${longText}"}`; + + const iterations = 100_000; + + // Warm up + for (let i = 0; i < 1000; i++) { + detectBadJsonStrings(goodJson); + detectBadJsonStrings(badJson); + } + + // Measure good JSON (most common case) + const goodStart = performance.now(); + for (let i = 0; i < iterations; i++) { + detectBadJsonStrings(goodJson); + } + const goodTime = performance.now() - goodStart; + + // Measure bad JSON (edge case) + const badStart = performance.now(); + for (let i = 0; i < iterations; i++) { + detectBadJsonStrings(badJson); + } + const badTime = performance.now() - badStart; + + // Measure baseline (just function call overhead) + const baselineStart = performance.now(); + for (let i = 0; i < iterations; i++) { + // Empty function call to measure baseline + } + const baselineTime = performance.now() - baselineStart; + + const goodOverhead = goodTime - baselineTime; + const badOverhead = badTime - baselineTime; + + console.log(`Baseline (${iterations} iterations): ${baselineTime.toFixed(2)}ms`); + console.log( + `Good JSON (${iterations} iterations): ${goodTime.toFixed( + 2 + )}ms (overhead: ${goodOverhead.toFixed(2)}ms)` + ); + console.log( + `Bad JSON (${iterations} iterations): ${badTime.toFixed( + 2 + )}ms (overhead: ${badOverhead.toFixed(2)}ms)` + ); + console.log( + `Average per call - Good: ${(goodOverhead / iterations).toFixed(4)}ms, Bad: ${( + badOverhead / iterations + ).toFixed(4)}ms` + ); + + // Assertions for performance expectations + // Good JSON should be reasonably fast (most common case) + expect(goodOverhead / iterations).toBeLessThan(0.01); // Less than 10 microseconds per call + + // Bad JSON can be slower due to regex matching, but still reasonable + expect(badOverhead / iterations).toBeLessThan(0.01); // Less than 20 microseconds per call + + // Total overhead for 100k calls should be reasonable + expect(goodOverhead).toBeLessThan(1000); // Less than 1 second for 100k calls + }); + + it("should handle various JSON sizes efficiently", () => { + const sizes = [100, 1000, 10000, 100000]; + const iterations = 10_000; + + for (const size of sizes) { + const text = `hello world `.repeat(size / 11); // Approximate size + const goodJson = `{"title": "hello", "text": "${text}"}`; + + const start = performance.now(); + for (let i = 0; i < iterations; i++) { + detectBadJsonStrings(goodJson); + } + const time = performance.now() - start; + + console.log( + `Size ${size} chars (${iterations} iterations): ${time.toFixed(2)}ms (${( + time / iterations + ).toFixed(4)}ms per call)` + ); + + // Performance should scale reasonably with size + expect(time / iterations).toBeLessThan(size / 1000); // Roughly linear scaling + } + }); + + it("should show significant performance improvement with quick rejection", () => { + const longText = `hello world `.repeat(1_000); + const goodJson = `{"title": "hello", "text": "${longText}"}`; + const badJson = `{"title": "hello\\ud835", "text": "${longText}"}`; + const noUnicodeJson = `{"title": "hello", "text": "${longText}"}`; + + const iterations = 100_000; + + // Warm up + for (let i = 0; i < 1000; i++) { + detectBadJsonStrings(goodJson); + detectBadJsonStrings(badJson); + detectBadJsonStrings(noUnicodeJson); + } + + // Test strings with no Unicode escapes (99.9% case) + const noUnicodeStart = performance.now(); + for (let i = 0; i < iterations; i++) { + detectBadJsonStrings(noUnicodeJson); + } + const noUnicodeTime = performance.now() - noUnicodeStart; + + // Test strings with Unicode escapes (0.1% case) + const withUnicodeStart = performance.now(); + for (let i = 0; i < iterations; i++) { + detectBadJsonStrings(badJson); + } + const withUnicodeTime = performance.now() - withUnicodeStart; + + console.log( + `No Unicode escapes (${iterations} iterations): ${noUnicodeTime.toFixed(2)}ms (${( + noUnicodeTime / iterations + ).toFixed(4)}ms per call)` + ); + console.log( + `With Unicode escapes (${iterations} iterations): ${withUnicodeTime.toFixed(2)}ms (${( + withUnicodeTime / iterations + ).toFixed(4)}ms per call)` + ); + console.log( + `Performance ratio: ${(withUnicodeTime / noUnicodeTime).toFixed( + 2 + )}x slower for Unicode strings` + ); + + // Both cases should be extremely fast (under 1 microsecond per call) + expect(noUnicodeTime / iterations).toBeLessThan(0.001); // Less than 1 microsecond + expect(withUnicodeTime / iterations).toBeLessThan(0.001); // Less than 1 microsecond + + // The difference should be reasonable (not more than 5x) + expect(noUnicodeTime / withUnicodeTime).toBeLessThan(5); + }); +}); + +function processPacket(data: string): { data?: string; dataType?: string } { + if (detectBadJsonStrings(data)) { + return { data: undefined }; + } + return { data, dataType: "application/json" }; +} diff --git a/apps/webapp/test/runsReplicationService.part2.test.ts b/apps/webapp/test/runsReplicationService.part2.test.ts index 01627e4ec..cb04867f9 100644 --- a/apps/webapp/test/runsReplicationService.part2.test.ts +++ b/apps/webapp/test/runsReplicationService.part2.test.ts @@ -1,12 +1,11 @@ import { ClickHouse } from "@internal/clickhouse"; import { containerTest } from "@internal/testcontainers"; import { Logger } from "@trigger.dev/core/logger"; +import { readFile } from "node:fs/promises"; import { setTimeout } from "node:timers/promises"; import { z } from "zod"; -import { TaskRunStatus } from "~/database-types"; import { RunsReplicationService } from "~/services/runsReplicationService.server"; -import { createInMemoryTracing } from "./utils/tracing"; -import superjson from "superjson"; +import { detectBadJsonStrings } from "~/utils/detectBadJsonStrings"; vi.setConfig({ testTimeout: 60_000 }); @@ -611,4 +610,177 @@ describe("RunsReplicationService (part 2/2)", () => { }, { timeout: 60_000 * 5 } ); + + containerTest( + "should insert TaskRuns even if there are incomplete Unicode escape sequences in the JSON", + 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-stress-bulk-insert", + }); + + const runsReplicationService = new RunsReplicationService({ + clickhouse, + pgConnectionUrl: postgresContainer.getConnectionUri(), + serviceName: "runs-replication-stress-bulk-insert", + slotName: "task_runs_to_clickhouse_v1", + publicationName: "task_runs_to_clickhouse_v1_publication", + redisOptions, + maxFlushConcurrency: 10, + flushIntervalMs: 100, + flushBatchSize: 50, + leaderLockTimeoutMs: 5000, + leaderLockExtendIntervalMs: 1000, + ackIntervalSeconds: 5, + logger: new Logger("runs-replication-stress-bulk-insert", "info"), + }); + + await runsReplicationService.start(); + + const organization = await prisma.organization.create({ + data: { + title: "test-stress-bulk-insert", + slug: "test-stress-bulk-insert", + }, + }); + + const project = await prisma.project.create({ + data: { + name: "test-stress-bulk-insert", + slug: "test-stress-bulk-insert", + organizationId: organization.id, + externalRef: "test-stress-bulk-insert", + }, + }); + + const runtimeEnvironment = await prisma.runtimeEnvironment.create({ + data: { + slug: "test-stress-bulk-insert", + type: "DEVELOPMENT", + projectId: project.id, + organizationId: organization.id, + apiKey: "test-stress-bulk-insert", + pkApiKey: "test-stress-bulk-insert", + shortcode: "test-stress-bulk-insert", + }, + }); + + // Prepare 9 unique TaskRuns + const now = Date.now(); + const runsData = Array.from({ length: 9 }, (_, i) => ({ + friendlyId: `run_bulk_${now}_${i}`, + taskIdentifier: `my-task-bulk`, + payload: `{"title": "hello"}`, + payloadType: "application/json", + traceId: `bulk-${i}`, + spanId: `bulk-${i}`, + queue: "test-stress-bulk-insert", + runtimeEnvironmentId: runtimeEnvironment.id, + projectId: project.id, + organizationId: organization.id, + environmentType: "DEVELOPMENT" as const, + engine: "V2" as const, + status: "PENDING" as const, + attemptNumber: 1, + createdAt: new Date(now + i), + updatedAt: new Date(now + i), + })); + + //add a run with incomplete Unicode escape sequences + const badPayload = await readFile(`${__dirname}/bad-clickhouse-output.json`, "utf-8"); + const hasProblems = detectBadJsonStrings(badPayload); + expect(hasProblems).toBe(true); + + runsData.push({ + friendlyId: `run_bulk_${now}_10`, + taskIdentifier: `my-task-bulk`, + payload: badPayload, + payloadType: "application/json", + traceId: `bulk-10`, + spanId: `bulk-10`, + queue: "test-stress-bulk-insert", + runtimeEnvironmentId: runtimeEnvironment.id, + projectId: project.id, + organizationId: organization.id, + environmentType: "DEVELOPMENT" as const, + engine: "V2" as const, + status: "PENDING" as const, + attemptNumber: 1, + createdAt: new Date(now + 10), + updatedAt: new Date(now + 10), + }); + + // Bulk insert + const created = await prisma.taskRun.createMany({ data: runsData }); + expect(created.count).toBe(10); + + // Update the runs (not the 10th one) + await prisma.taskRun.updateMany({ + where: { + spanId: { not: "bulk-10" }, + }, + data: { + status: "COMPLETED_SUCCESSFULLY", + output: `{"foo":"bar"}`, + outputType: "application/json", + }, + }); + + // Give the 10th one a bad payload + await prisma.taskRun.updateMany({ + where: { + spanId: "bulk-10", + }, + data: { + status: "COMPLETED_SUCCESSFULLY", + output: badPayload, + outputType: "application/json", + }, + }); + + // Wait for replication + await setTimeout(5000); + + // Query ClickHouse for all runs using FINAL + const queryRuns = clickhouse.reader.query({ + name: "runs-replication-stress-bulk-insert", + query: `SELECT * FROM trigger_dev.task_runs_v2 FINAL`, + schema: z.any(), + }); + + const [queryError, result] = await queryRuns({}); + expect(queryError).toBeNull(); + expect(result?.length).toBe(10); + + console.log("Data", { + runsData, + result, + }); + + // Check a few random runs for correctness + for (let i = 0; i < 9; i++) { + const expected = runsData[i]; + const found = result?.find((r: any) => r.friendly_id === expected.friendlyId); + expect(found).toBeDefined(); + expect(found).toEqual( + expect.objectContaining({ + friendly_id: expected.friendlyId, + trace_id: expected.traceId, + task_identifier: expected.taskIdentifier, + status: "COMPLETED_SUCCESSFULLY", + }) + ); + expect(found?.output).toBeDefined(); + } + + // Check the run with the bad JSON + const foundBad = result?.find((r: any) => r.span_id === "bulk-10"); + expect(foundBad).toBeDefined(); + expect(foundBad?.output).toStrictEqual({}); + + await runsReplicationService.stop(); + } + ); }); diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml index 02b00de7b..933881664 100644 --- a/docker/docker-compose.yml +++ b/docker/docker-compose.yml @@ -111,7 +111,7 @@ services: ports: - "5521:5521" environment: - VITE_CLICKHOUSE_URL: "http://clickhouse:8123" + VITE_CLICKHOUSE_URL: "http://localhost:8123" VITE_CLICKHOUSE_USER: "default" VITE_CLICKHOUSE_PASS: "password" networks: diff --git a/internal-packages/clickhouse/package.json b/internal-packages/clickhouse/package.json index e2767708e..d85051b50 100644 --- a/internal-packages/clickhouse/package.json +++ b/internal-packages/clickhouse/package.json @@ -25,5 +25,15 @@ "db:migrate:down": "GOOSE_COMMAND=down pnpm run db:migrate", "test": "vitest --sequence.concurrent=false --no-file-parallelism", "test:coverage": "vitest --sequence.concurrent=false --no-file-parallelism --coverage.enabled" + }, + "exports": { + "./package.json": "./package.json", + ".": { + "import": { + "@triggerdotdev/source": "./src/index.ts", + "types": "./dist/src/index.d.ts", + "default": "./dist/src/index.js" + } + } } -} \ No newline at end of file +} diff --git a/internal-packages/clickhouse/schema/005_add_task_runs_v2_concurrency_bulkactions.sql b/internal-packages/clickhouse/schema/005_add_task_runs_v2_concurrency_bulkactions.sql new file mode 100644 index 000000000..3dd624746 --- /dev/null +++ b/internal-packages/clickhouse/schema/005_add_task_runs_v2_concurrency_bulkactions.sql @@ -0,0 +1,12 @@ +-- +goose Up +/* +Add concurrency_key and bulk_action_group_ids columns with defaults. + */ +ALTER TABLE trigger_dev.task_runs_v2 +ADD COLUMN concurrency_key String DEFAULT '', +ADD COLUMN bulk_action_group_ids Array(String) DEFAULT []; + +-- +goose Down +ALTER TABLE trigger_dev.task_runs_v2 +DROP COLUMN concurrency_key, +DROP COLUMN bulk_action_group_ids; \ No newline at end of file diff --git a/internal-packages/clickhouse/src/taskRuns.test.ts b/internal-packages/clickhouse/src/taskRuns.test.ts index 30ea0270b..b51c9f38c 100644 --- a/internal-packages/clickhouse/src/taskRuns.test.ts +++ b/internal-packages/clickhouse/src/taskRuns.test.ts @@ -61,6 +61,8 @@ describe("Task Runs V2", () => { root_run_id: "root_run_1234", parent_run_id: "parent_run_1234", depth: 1, + concurrency_key: "concurrency_key_1234", + bulk_action_group_ids: ["bulk_action_group_id_1234", "bulk_action_group_id_1235"], _version: "1", }, ]); @@ -75,6 +77,8 @@ describe("Task Runs V2", () => { schema: z.object({ environment_id: z.string(), run_id: z.string(), + concurrency_key: z.string(), + bulk_action_group_ids: z.array(z.string()), }), params: z.object({ run_id: z.string(), @@ -89,6 +93,8 @@ describe("Task Runs V2", () => { expect.objectContaining({ environment_id: "env_1234", run_id: "run_1234", + concurrency_key: "concurrency_key_1234", + bulk_action_group_ids: ["bulk_action_group_id_1234", "bulk_action_group_id_1235"], }), ]) ); diff --git a/internal-packages/clickhouse/src/taskRuns.ts b/internal-packages/clickhouse/src/taskRuns.ts index aff8af829..86830b5bd 100644 --- a/internal-packages/clickhouse/src/taskRuns.ts +++ b/internal-packages/clickhouse/src/taskRuns.ts @@ -42,6 +42,8 @@ export const TaskRunV2 = z.object({ idempotency_key: z.string(), expiration_ttl: z.string(), is_test: z.boolean().default(false), + concurrency_key: z.string().default(""), + bulk_action_group_ids: z.array(z.string()).default([]), _version: z.string(), _is_deleted: z.number().int().default(0), }); From 08d84eb5960c690078cc4dbefa93a2954b22127e Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Mon, 30 Jun 2025 16:37:52 +0100 Subject: [PATCH 3/3] Fix v4 concurrency sweeper flaky test (#2213) --- .../src/run-queue/tests/concurrencySweeper.test.ts | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/internal-packages/run-engine/src/run-queue/tests/concurrencySweeper.test.ts b/internal-packages/run-engine/src/run-queue/tests/concurrencySweeper.test.ts index 97fef5303..8a1399911 100644 --- a/internal-packages/run-engine/src/run-queue/tests/concurrencySweeper.test.ts +++ b/internal-packages/run-engine/src/run-queue/tests/concurrencySweeper.test.ts @@ -74,9 +74,9 @@ describe("RunQueue Concurrency Sweeper", () => { }, concurrencySweeper: { scanSchedule: "* * * * * *", // Every second - scanJitter: 5, + scanJitterInMs: 5, processMarkedSchedule: "* * * * * *", // Every second - processMarkedJitter: 5, + processMarkedJitterInMs: 5, callback: async (runIds) => { if (!enableConcurrencySweeper) { return []; @@ -151,9 +151,9 @@ describe("RunQueue Concurrency Sweeper", () => { enableConcurrencySweeper = true; - await setTimeout(3_000); // Now a run is "completed" and should be removed from the concurrency set + await setTimeout(5_000); // Now a run is "completed" and should be removed from the concurrency set - // queue concurrency should be 0 + // queue concurrency should be 1 const queueConcurrency3 = await queue.currentConcurrencyOfQueue( authenticatedEnvDev, messageDev.queue