Publish redis-worker and add graceful shutdown manager (#1810)

* add shutdown manager

* update ai test instructions

* add shutdown timeout to redis-worker

* move redis worker to packages

* add unregister method

* prep for publishing package

* fix types

* update ai files

* fix cursor terminal links

* prevent overly friendly ids

* use structured logger

* use unique shutdown handler names

* rework suspend completion

* add trycatch util

* rework suspend restore

* add http server metrics

* add missing prom-client to core

* add prom metrics to redis worker

* bundle redis-worker

* fix esm/cjs interop

* remove proxy from changeset ignore and add supervisor

* add pause to prerelease script for any manual edits

* unregister the correct handler and add early detection

* small change to http handler return

* fix worker tests

* fix shutdown manager tests
This commit is contained in:
nicktrn
2025-03-21 14:53:27 +00:00
committed by GitHub
parent adca1997af
commit 49a3f72e13
58 changed files with 1637 additions and 355 deletions
+1 -1
View File
@@ -12,7 +12,7 @@
"access": "public",
"baseBranch": "main",
"updateInternalDependencies": "patch",
"ignore": ["webapp", "proxy", "coordinator", "docker-provider", "kubernetes-provider"],
"ignore": ["webapp", "supervisor", "coordinator", "docker-provider", "kubernetes-provider"],
"___experimentalUnsafeOptions_WILL_CHANGE_IN_PATCH": {
"onlyUpdatePeerDependentsWhenOutOfRange": true
}
+1 -1
View File
@@ -9,7 +9,7 @@ The main trigger.dev webapp, which powers it's API and dashboard and makes up th
- `@trigger.dev/database` exports a Prisma 5.4.1 client that is used extensively in the webapp to access a PostgreSQL instance. The schema file is [schema.prisma](mdc:internal-packages/database/prisma/schema.prisma)
- `@trigger.dev/core` is a published package and is used to share code between the `@trigger.dev/sdk` and the webapp. It includes functionality but also a load of Zod schemas for data validation. When importing from `@trigger.dev/core` in the webapp, we never import the root `@trigger.dev/core` path, instead we favor one of the subpath exports that you can find in [package.json](mdc:packages/core/package.json)
- `@internal/run-engine` has all the code needed to trigger a run and take it through it's lifecycle to completion.
- `@internal/redis-worker` is a custom redis based background job/worker system that's used in the webapp and also used inside the run engine.
- `@trigger.dev/redis-worker` is a custom redis based background job/worker system that's used in the webapp and also used inside the run engine.
## Environment variables and testing
+1 -1
View File
@@ -14,13 +14,13 @@ This is a pnpm 8.15.5 monorepo that uses turborepo @turbo.json. The following wo
- <root>/packages/core is the `@trigger.dev/core` package that is shared across the SDK and other packages
- <root>/packages/build defines the types and prebuilt build extensions for trigger.dev. See our [build extensions docs](https://trigger.dev/docs/config/extensions/overview.md) for more information.
- <root>/packages/react-hooks defines some useful react hooks like our realtime hooks. See our [Realtime hooks](https://trigger.dev/docs/frontend/react-hooks/realtime.md) and our [Trigger hooks](https://trigger.dev/docs/frontend/react-hooks/triggering.md) for more information.
- <root>/packages/redis-worker is the `@trigger.dev/redis-worker` package that implements a custom background job/worker sytem powered by redis for offloading work to the background, used in the webapp and also in the Run Engine 2.0.
## Internal Packages
- <root>/internal-packages/\* are packages that are used internally only, not published, and usually they have a tsc build step and are used in the webapp
- <root>/internal-packages/database is the `@trigger.dev/database` package that exports a prisma client, has the schema file, and exports a few other helpers.
- <root>/internal-packages/run-engine is the `@internal/run-engine` package that is "Run Engine 2.0" and handles moving a run all the way through it's lifecycle
- <root>/internal-packages/redis-worker is the `@internal/redis-worker` package that implements a custom background job/worker sytem powered by redis for offloading work to the background, used in the webapp and also in the Run Engine 2.0.
- <root>/internal-packages/redis is the `@internal/redis` package that exports Redis types and the `createRedisClient` function to unify how we create redis clients in the repo. It's not used everywhere yet, but it's the preferred way to create redis clients from now on.
- <root>/internal-packages/testcontainers is the `@internal/testcontainers` package that exports a few useful functions for spinning up local testcontainers when writing vitest tests. See our [tests.md](./tests.md) file for more information.
- <root>/internal-packages/zodworker is the `@internal/zodworker` package that implements a wrapper around graphile-worker that allows us to use zod to validate our background jobs. We are moving away from using graphile-worker as our background job system, replacing it with our own redis-worker package.
+5
View File
@@ -25,6 +25,11 @@ pnpm run test ./src/components/Button.test.ts
We use vitest for testing. We almost NEVER mock anything. Start with a top-level "describe", and have multiple "it" statements inside of it.
New test files should be placed right next to the file being tested. For example:
- Source file: `./src/services/MyService.ts`
- Test file: `./src/services/MyService.test.ts`
When writing anything that needs redis or postgresql, we have some internal "testcontainers" that are used to spin up a local instance, redis, or both.
redisTest:
+1 -1
View File
@@ -6,7 +6,7 @@
"type": "module",
"scripts": {
"build": "tsc",
"dev": "tsx --experimental-sqlite --require dotenv/config --watch src/index.ts",
"dev": "tsx --experimental-sqlite --require dotenv/config --watch src/index.ts || (echo '!! Remember to run: nvm use'; exit 1)",
"start": "node --experimental-sqlite dist/index.js",
"typecheck": "tsc --noEmit"
},
+6
View File
@@ -45,6 +45,12 @@ const Env = z.object({
// Used by the resource monitor
OVERRIDE_CPU_TOTAL: z.coerce.number().optional(),
OVERRIDE_MEMORY_TOTAL_GB: z.coerce.number().optional(),
// Kubernetes specific settings
KUBERNETES_FORCE_ENABLED: BoolEnv.default(false),
KUBERNETES_NAMESPACE: z.string().default("default"),
EPHEMERAL_STORAGE_SIZE_LIMIT: z.string().default("10Gi"),
EPHEMERAL_STORAGE_SIZE_REQUEST: z.string().default("2Gi"),
});
export const env = Env.parse(stdEnv);
+14 -4
View File
@@ -13,7 +13,11 @@ import {
} from "./resourceMonitor.js";
import { KubernetesWorkloadManager } from "./workloadManager/kubernetes.js";
import { DockerWorkloadManager } from "./workloadManager/docker.js";
import { HttpServer, CheckpointClient } from "@trigger.dev/core/v3/serverOnly";
import {
HttpServer,
CheckpointClient,
isKubernetesEnvironment,
} from "@trigger.dev/core/v3/serverOnly";
import { createK8sApi, RUNTIME_ENV } from "./clients/kubernetes.js";
class ManagedSupervisor {
@@ -25,7 +29,7 @@ class ManagedSupervisor {
private readonly resourceMonitor: ResourceMonitor;
private readonly checkpointClient?: CheckpointClient;
private readonly isKubernetes = RUNTIME_ENV === "kubernetes";
private readonly isKubernetes = isKubernetesEnvironment(env.KUBERNETES_FORCE_ENABLED);
private readonly warmStartUrl = env.TRIGGER_WARM_START_URL;
constructor() {
@@ -94,6 +98,7 @@ class ManagedSupervisor {
this.checkpointClient = new CheckpointClient({
apiUrl: new URL(env.TRIGGER_CHECKPOINT_URL),
workerClient: this.workerSession.httpClient,
orchestrator: this.isKubernetes ? "KUBERNETES" : "DOCKER",
});
}
@@ -127,7 +132,9 @@ class ManagedSupervisor {
return;
}
if (message.checkpoint) {
const { checkpoint, ...rest } = message;
if (checkpoint) {
this.logger.log("[ManagedWorker] Restoring run", { runId: message.run.id });
if (!this.checkpointClient) {
@@ -139,7 +146,10 @@ class ManagedSupervisor {
const didRestore = await this.checkpointClient.restoreRun({
runFriendlyId: message.run.friendlyId,
snapshotFriendlyId: message.snapshot.friendlyId,
checkpoint: message.checkpoint,
body: {
...rest,
checkpoint,
},
});
if (didRestore) {
@@ -9,9 +9,6 @@ import type { EnvironmentType, MachinePreset } from "@trigger.dev/core/v3";
import { env } from "../env.js";
import { type K8sApi, createK8sApi, type k8s } from "../clients/kubernetes.js";
const POD_EPHEMERAL_STORAGE_SIZE_LIMIT = process.env.POD_EPHEMERAL_STORAGE_SIZE_LIMIT || "10Gi";
const POD_EPHEMERAL_STORAGE_SIZE_REQUEST = process.env.POD_EPHEMERAL_STORAGE_SIZE_REQUEST || "2Gi";
type ResourceQuantities = {
[K in "cpu" | "memory" | "ephemeral-storage"]?: string;
};
@@ -19,7 +16,7 @@ type ResourceQuantities = {
export class KubernetesWorkloadManager implements WorkloadManager {
private readonly logger = new SimpleStructuredLogger("kubernetes-workload-provider");
private k8s: K8sApi;
private namespace = "default";
private namespace = env.KUBERNETES_NAMESPACE;
constructor(private opts: WorkloadManagerOptions) {
this.k8s = createK8sApi();
@@ -205,13 +202,13 @@ export class KubernetesWorkloadManager implements WorkloadManager {
get #defaultResourceRequests(): ResourceQuantities {
return {
"ephemeral-storage": POD_EPHEMERAL_STORAGE_SIZE_REQUEST,
"ephemeral-storage": env.EPHEMERAL_STORAGE_SIZE_REQUEST,
};
}
get #defaultResourceLimits(): ResourceQuantities {
return {
"ephemeral-storage": POD_EPHEMERAL_STORAGE_SIZE_LIMIT,
"ephemeral-storage": env.EPHEMERAL_STORAGE_SIZE_LIMIT,
};
}
+28 -7
View File
@@ -94,8 +94,8 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
this.websocketServer = this.createWebsocketServer();
}
private runnerIdFromRequest(req: IncomingMessage): string | undefined {
const value = req.headers[WORKLOAD_HEADERS.RUNNER_ID];
private headerValueFromRequest(req: IncomingMessage, headerName: string): string | undefined {
const value = req.headers[headerName];
if (Array.isArray(value)) {
return value[0];
@@ -104,6 +104,22 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
return value;
}
private runnerIdFromRequest(req: IncomingMessage): string | undefined {
return this.headerValueFromRequest(req, WORKLOAD_HEADERS.RUNNER_ID);
}
private deploymentIdFromRequest(req: IncomingMessage): string | undefined {
return this.headerValueFromRequest(req, WORKLOAD_HEADERS.DEPLOYMENT_ID);
}
private deploymentVersionFromRequest(req: IncomingMessage): string | undefined {
return this.headerValueFromRequest(req, WORKLOAD_HEADERS.DEPLOYMENT_VERSION);
}
private projectRefFromRequest(req: IncomingMessage): string | undefined {
return this.headerValueFromRequest(req, WORKLOAD_HEADERS.PROJECT_REF);
}
private createHttpServer({ host, port }: { host: string; port: number }) {
return new HttpServer({ port, host })
.route(
@@ -213,8 +229,10 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
}
const runnerId = this.runnerIdFromRequest(req);
const deploymentVersion = this.deploymentVersionFromRequest(req);
const projectRef = this.projectRefFromRequest(req);
if (!runnerId) {
if (!runnerId || !deploymentVersion || !projectRef) {
console.error("Invalid headers for suspend request", {
...params,
headers: req.headers,
@@ -241,16 +259,19 @@ export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
const suspendResult = await this.checkpointClient.suspendRun({
runFriendlyId: params.runFriendlyId,
snapshotFriendlyId: params.snapshotFriendlyId,
containerId: runnerId,
runnerId,
body: {
runnerId,
runId: params.runFriendlyId,
snapshotId: params.snapshotFriendlyId,
projectRef,
deploymentVersion,
},
});
if (!suspendResult) {
console.error("Failed to suspend run", { params });
return;
}
console.log("Suspended run", { params });
},
}
)
+3
View File
@@ -424,6 +424,7 @@ const EnvironmentSchema = z.object({
RUN_ENGINE_QUEUE_AGE_RANDOMIZATION_BIAS: z.coerce.number().default(0.25),
RUN_ENGINE_REUSE_SNAPSHOT_COUNT: z.coerce.number().int().default(0),
RUN_ENGINE_MAXIMUM_ENV_COUNT: z.coerce.number().int().optional(),
RUN_ENGINE_WORKER_SHUTDOWN_TIMEOUT_MS: z.coerce.number().int().default(60_000),
RUN_ENGINE_WORKER_REDIS_HOST: z
.string()
@@ -591,6 +592,7 @@ const EnvironmentSchema = z.object({
LEGACY_RUN_ENGINE_WORKER_POLL_INTERVAL: z.coerce.number().int().default(1000),
LEGACY_RUN_ENGINE_WORKER_IMMEDIATE_POLL_INTERVAL: z.coerce.number().int().default(50),
LEGACY_RUN_ENGINE_WORKER_CONCURRENCY_LIMIT: z.coerce.number().int().default(100),
LEGACY_RUN_ENGINE_WORKER_SHUTDOWN_TIMEOUT_MS: z.coerce.number().int().default(60_000),
LEGACY_RUN_ENGINE_WORKER_REDIS_HOST: z
.string()
@@ -633,6 +635,7 @@ const EnvironmentSchema = z.object({
COMMON_WORKER_POLL_INTERVAL: z.coerce.number().int().default(1000),
COMMON_WORKER_IMMEDIATE_POLL_INTERVAL: z.coerce.number().int().default(50),
COMMON_WORKER_CONCURRENCY_LIMIT: z.coerce.number().int().default(100),
COMMON_WORKER_SHUTDOWN_TIMEOUT_MS: z.coerce.number().int().default(60_000),
COMMON_WORKER_REDIS_HOST: z
.string()
+2 -1
View File
@@ -1,4 +1,4 @@
import { Worker as RedisWorker } from "@internal/redis-worker";
import { Worker as RedisWorker } from "@trigger.dev/redis-worker";
import { Logger } from "@trigger.dev/core/logger";
import { z } from "zod";
import { env } from "~/env.server";
@@ -80,6 +80,7 @@ function initializeWorker() {
},
pollIntervalMs: env.COMMON_WORKER_POLL_INTERVAL,
immediatePollIntervalMs: env.COMMON_WORKER_IMMEDIATE_POLL_INTERVAL,
shutdownTimeoutMs: env.COMMON_WORKER_SHUTDOWN_TIMEOUT_MS,
logger: new Logger("CommonWorker", "debug"),
jobs: {
"v3.deliverAlert": async ({ payload }) => {
@@ -1,4 +1,4 @@
import { Worker as RedisWorker } from "@internal/redis-worker";
import { Worker as RedisWorker } from "@trigger.dev/redis-worker";
import { Logger } from "@trigger.dev/core/logger";
import { z } from "zod";
import { env } from "~/env.server";
@@ -67,6 +67,7 @@ function initializeWorker() {
},
pollIntervalMs: env.LEGACY_RUN_ENGINE_WORKER_POLL_INTERVAL,
immediatePollIntervalMs: env.LEGACY_RUN_ENGINE_WORKER_IMMEDIATE_POLL_INTERVAL,
shutdownTimeoutMs: env.LEGACY_RUN_ENGINE_WORKER_SHUTDOWN_TIMEOUT_MS,
logger: new Logger("LegacyRunEngineWorker", "debug"),
jobs: {
runHeartbeat: async ({ payload }) => {
+1
View File
@@ -17,6 +17,7 @@ function createRunEngine() {
workers: env.RUN_ENGINE_WORKER_COUNT,
tasksPerWorker: env.RUN_ENGINE_TASKS_PER_WORKER,
pollIntervalMs: env.RUN_ENGINE_WORKER_POLL_INTERVAL,
shutdownTimeoutMs: env.RUN_ENGINE_WORKER_SHUTDOWN_TIMEOUT_MS,
redis: {
keyPrefix: "engine:",
port: env.RUN_ENGINE_WORKER_REDIS_PORT ?? undefined,
+1 -1
View File
@@ -53,7 +53,7 @@
"@internal/run-engine": "workspace:*",
"@internal/zod-worker": "workspace:*",
"@internal/redis": "workspace:*",
"@internal/redis-worker": "workspace:*",
"@trigger.dev/redis-worker": "workspace:*",
"@internationalized/date": "^3.5.1",
"@lezer/highlight": "^1.1.6",
"@opentelemetry/api": "1.9.0",
@@ -1,38 +0,0 @@
{
"name": "@internal/redis-worker",
"private": true,
"version": "0.0.1",
"main": "./dist/src/index.js",
"types": "./dist/src/index.d.ts",
"type": "module",
"exports": {
".": {
"@triggerdotdev/source": "./src/index.ts",
"import": "./dist/src/index.js",
"types": "./dist/src/index.d.ts",
"default": "./dist/src/index.js"
}
},
"dependencies": {
"@internal/tracing": "workspace:*",
"@internal/redis": "workspace:*",
"@trigger.dev/core": "workspace:*",
"lodash.omit": "^4.5.0",
"nanoid": "^5.0.7",
"p-limit": "^6.2.0",
"zod": "3.23.8"
},
"devDependencies": {
"@internal/testcontainers": "workspace:*",
"@types/lodash.omit": "^4.5.7",
"vitest": "^1.4.0",
"rimraf": "6.0.1"
},
"scripts": {
"clean": "rimraf dist",
"typecheck": "tsc --noEmit -p tsconfig.build.json",
"test": "vitest --sequence.concurrent=false --no-file-parallelism",
"build": "pnpm run clean && tsc -p tsconfig.build.json",
"dev": "tsc --watch -p tsconfig.build.json"
}
}
@@ -1,8 +0,0 @@
{
"references": [{ "path": "./tsconfig.src.json" }, { "path": "./tsconfig.test.json" }],
"compilerOptions": {
"moduleResolution": "Node16",
"module": "Node16",
"customConditions": ["@triggerdotdev/source"]
}
}
@@ -1,19 +0,0 @@
{
"include": ["src/**/*.ts"],
"exclude": ["node_modules", "src/**/*.test.ts"],
"compilerOptions": {
"composite": true,
"target": "ES2019",
"lib": ["ES2019", "DOM", "DOM.Iterable", "DOM.AsyncIterable"],
"module": "Node16",
"moduleResolution": "Node16",
"moduleDetection": "force",
"verbatimModuleSyntax": false,
"esModuleInterop": true,
"forceConsistentCasingInFileNames": true,
"isolatedModules": true,
"preserveWatchOutput": true,
"skipLibCheck": true,
"strict": true
}
}
@@ -1,20 +0,0 @@
{
"include": ["src/**/*.test.ts"],
"references": [{ "path": "./tsconfig.src.json" }],
"compilerOptions": {
"composite": true,
"target": "ES2019",
"lib": ["ES2019", "DOM", "DOM.Iterable", "DOM.AsyncIterable"],
"module": "Node16",
"moduleResolution": "Node16",
"moduleDetection": "force",
"verbatimModuleSyntax": false,
"types": ["vitest/globals"],
"esModuleInterop": true,
"forceConsistentCasingInFileNames": true,
"isolatedModules": true,
"preserveWatchOutput": true,
"skipLibCheck": true,
"strict": true
}
}
+1 -1
View File
@@ -1,4 +1,4 @@
import { Redis, RedisOptions } from "ioredis";
import { Redis, type RedisOptions } from "ioredis";
import { Logger } from "@trigger.dev/core/logger";
export { Redis, type Callback, type RedisOptions, type Result, type RedisCommander } from "ioredis";
+2 -2
View File
@@ -15,7 +15,7 @@
},
"dependencies": {
"@internal/redis": "workspace:*",
"@internal/redis-worker": "workspace:*",
"@trigger.dev/redis-worker": "workspace:*",
"@internal/tracing": "workspace:*",
"@trigger.dev/core": "workspace:*",
"@trigger.dev/database": "workspace:*",
@@ -41,4 +41,4 @@
"build": "pnpm run clean && tsc -p tsconfig.build.json",
"dev": "tsc --watch -p tsconfig.build.json"
}
}
}
@@ -1,5 +1,5 @@
import { createRedisClient, Redis } from "@internal/redis";
import { Worker } from "@internal/redis-worker";
import { Worker } from "@trigger.dev/redis-worker";
import { startSpan, trace, Tracer } from "@internal/tracing";
import { Logger } from "@trigger.dev/core/logger";
import {
@@ -120,6 +120,7 @@ export class RunEngine {
concurrency: options.worker,
pollIntervalMs: options.worker.pollIntervalMs,
immediatePollIntervalMs: options.worker.immediatePollIntervalMs,
shutdownTimeoutMs: options.worker.shutdownTimeoutMs,
logger: new Logger("RunEngineWorker", "debug"),
jobs: {
finishWaitpoint: async ({ payload }) => {
@@ -1,5 +1,5 @@
import { type RedisOptions } from "@internal/redis";
import { Worker, type WorkerConcurrencyOptions } from "@internal/redis-worker";
import { Worker, type WorkerConcurrencyOptions } from "@trigger.dev/redis-worker";
import { Tracer } from "@internal/tracing";
import { MachinePreset, MachinePresetName, QueueOptions, RetryOptions } from "@trigger.dev/core/v3";
import { PrismaClient } from "@trigger.dev/database";
@@ -13,6 +13,7 @@ export type RunEngineOptions = {
redis: RedisOptions;
pollIntervalMs?: number;
immediatePollIntervalMs?: number;
shutdownTimeoutMs?: number;
} & WorkerConcurrencyOptions;
machines: {
defaultMachine: MachinePresetName;
+2 -2
View File
@@ -1,5 +1,5 @@
import { Span, SpanOptions, SpanStatusCode, Tracer } from "@opentelemetry/api";
import { Logger, SeverityNumber } from "@opentelemetry/api-logs";
import { type Span, type SpanOptions, SpanStatusCode, type Tracer } from "@opentelemetry/api";
import { type Logger, SeverityNumber } from "@opentelemetry/api-logs";
import { flattenAttributes } from "@trigger.dev/core/v3/utils/flattenAttributes";
export * from "@opentelemetry/semantic-conventions";
+2 -1
View File
@@ -74,7 +74,8 @@
"@changesets/assemble-release-plan@5.2.4": "patches/@changesets__assemble-release-plan@5.2.4.patch",
"engine.io-parser@5.2.2": "patches/engine.io-parser@5.2.2.patch",
"graphile-worker@0.16.6": "patches/graphile-worker@0.16.6.patch",
"redlock@5.0.0-beta.2": "patches/redlock@5.0.0-beta.2.patch"
"redlock@5.0.0-beta.2": "patches/redlock@5.0.0-beta.2.patch",
"supports-hyperlinks@2.3.0": "patches/supports-hyperlinks@2.3.0.patch"
}
}
}
@@ -149,8 +149,10 @@ class ManagedRunController {
this.httpClient = new WorkloadHttpClient({
workerApiUrl: this.workerApiUrl,
deploymentId: env.TRIGGER_DEPLOYMENT_ID,
runnerId: env.TRIGGER_RUNNER_ID,
deploymentId: env.TRIGGER_DEPLOYMENT_ID,
deploymentVersion: env.TRIGGER_DEPLOYMENT_VERSION,
projectRef: env.TRIGGER_PROJECT_REF,
});
if (env.TRIGGER_WARM_START_URL) {
+2
View File
@@ -189,8 +189,10 @@
"humanize-duration": "^3.27.3",
"jose": "^5.4.0",
"nanoid": "^3.3.4",
"prom-client": "^15.1.0",
"socket.io": "4.7.4",
"socket.io-client": "4.7.5",
"std-env": "^3.8.1",
"superjson": "^2.2.1",
"tinyexec": "^0.3.2",
"zod": "3.23.8",
+9
View File
@@ -1,3 +1,12 @@
export function assertExhaustive(x: never): never {
throw new Error("Unexpected object: " + x);
}
export async function tryCatch<T, E = Error>(promise: Promise<T>): Promise<[null, T] | [E, null]> {
try {
const data = await promise;
return [null, data];
} catch (error) {
return [error as E, null];
}
}
@@ -21,6 +21,10 @@ export function toFriendlyId(entityName: string, internalId: string): string {
throw new Error("Internal ID cannot be empty");
}
if (internalId.startsWith(`${entityName}_`)) {
return internalId;
}
return `${entityName}_${internalId}`;
}
@@ -8,4 +8,6 @@ export const WORKER_HEADERS = {
export const WORKLOAD_HEADERS = {
DEPLOYMENT_ID: "x-trigger-workload-deployment-id",
RUNNER_ID: "x-trigger-workload-runner-id",
DEPLOYMENT_VERSION: "x-trigger-workload-deployment-version",
PROJECT_REF: "x-trigger-workload-project-ref",
};
@@ -15,6 +15,8 @@ import {
WorkerApiRunHeartbeatResponseBody,
WorkerApiRunLatestSnapshotResponseBody,
WorkerApiDebugLogBody,
WorkerApiSuspendRunRequestBody,
WorkerApiSuspendRunResponseBody,
} from "./schemas.js";
import { SupervisorClientCommonOptions } from "./types.js";
import { getDefaultWorkerHeaders } from "./util.js";
@@ -220,14 +222,30 @@ export class SupervisorHttpClient {
);
}
getSuspendCompletionUrl(runId: string, snapshotId: string, runnerId?: string) {
return {
url: `${this.apiUrl}/engine/v1/worker-actions/runs/${runId}/snapshots/${snapshotId}/suspend`,
headers: {
...this.defaultHeaders,
...this.runnerIdHeader(runnerId),
},
};
async submitSuspendCompletion({
runId,
snapshotId,
runnerId,
body,
}: {
runId: string;
snapshotId: string;
runnerId?: string;
body: WorkerApiSuspendRunRequestBody;
}) {
return wrapZodFetch(
WorkerApiSuspendRunResponseBody,
`${this.apiUrl}/engine/v1/worker-actions/runs/${runId}/snapshots/${snapshotId}/suspend`,
{
method: "POST",
headers: {
...this.defaultHeaders,
...this.runnerIdHeader(runnerId),
"Content-Type": "application/json",
},
body: JSON.stringify(body),
}
);
}
private runnerIdHeader(runnerId?: string): Record<string, string> {
@@ -142,3 +142,11 @@ export const WorkerApiDebugLogBody = z.object({
properties: Attributes.optional(),
});
export type WorkerApiDebugLogBody = z.infer<typeof WorkerApiDebugLogBody>;
export const WorkerApiSuspendCompletionResponseBody = z.object({
success: z.boolean(),
error: z.string().optional(),
});
export type WorkerApiSuspendCompletionResponseBody = z.infer<
typeof WorkerApiSuspendCompletionResponseBody
>;
@@ -1,5 +1,7 @@
export type WorkloadClientCommonOptions = {
workerApiUrl: string;
deploymentId: string;
runnerId: string;
deploymentId: string;
deploymentVersion: string;
projectRef: string;
};
@@ -8,5 +8,7 @@ export function getDefaultWorkloadHeaders(
return createHeaders({
[WORKLOAD_HEADERS.DEPLOYMENT_ID]: options.deploymentId,
[WORKLOAD_HEADERS.RUNNER_ID]: options.runnerId,
[WORKLOAD_HEADERS.DEPLOYMENT_VERSION]: options.deploymentVersion,
[WORKLOAD_HEADERS.PROJECT_REF]: options.projectRef,
});
}
+7 -24
View File
@@ -1,4 +1,4 @@
import { CheckpointType } from "./runEngine.js";
import { CheckpointType, DequeuedMessage } from "./runEngine.js";
import z from "zod";
const CallbackUrl = z
@@ -8,20 +8,12 @@ const CallbackUrl = z
export const CheckpointServiceSuspendRequestBody = z.object({
type: CheckpointType,
containerId: z.string(),
simulate: z.boolean().optional(),
leaveRunning: z.boolean().optional(),
runId: z.string(),
snapshotId: z.string(),
runnerId: z.string(),
projectRef: z.string(),
deploymentVersion: z.string(),
reason: z.string().optional(),
callbacks: z
.object({
/** These headers will sent to all callbacks */
headers: z.record(z.string()).optional(),
/** This will be hit before suspending the container. Suspension will proceed unless we receive an error response. */
preSuspend: CallbackUrl.optional(),
/** This will be hit after suspending or failure to suspend the container */
completion: CallbackUrl.optional(),
})
.optional(),
});
export type CheckpointServiceSuspendRequestBody = z.infer<
@@ -39,16 +31,7 @@ export type CheckpointServiceSuspendResponseBody = z.infer<
typeof CheckpointServiceSuspendResponseBody
>;
export const CheckpointServiceRestoreRequestBody = z.discriminatedUnion("type", [
z.object({
type: z.literal(CheckpointType.Enum.DOCKER),
containerId: z.string(),
}),
z.object({
type: z.literal(CheckpointType.Enum.KUBERNETES),
containerId: z.string(),
}),
]);
export const CheckpointServiceRestoreRequestBody = DequeuedMessage.required({ checkpoint: true });
export type CheckpointServiceRestoreRequestBody = z.infer<
typeof CheckpointServiceRestoreRequestBody
+48 -46
View File
@@ -124,52 +124,6 @@ export const ExecutionResult = z.object({
export type ExecutionResult = z.infer<typeof ExecutionResult>;
/** This is sent to a Worker when a run is dequeued (a new run or continuing run) */
export const DequeuedMessage = z.object({
version: z.literal("1"),
snapshot: ExecutionSnapshot,
dequeuedAt: z.coerce.date(),
image: z.string().optional(),
checkpoint: z
.object({
id: z.string(),
type: z.string(),
location: z.string(),
reason: z.string().nullish(),
})
.optional(),
completedWaitpoints: z.array(CompletedWaitpoint),
backgroundWorker: z.object({
id: z.string(),
friendlyId: z.string(),
version: z.string(),
}),
deployment: z.object({
id: z.string().optional(),
friendlyId: z.string().optional(),
}),
run: z.object({
id: z.string(),
friendlyId: z.string(),
isTest: z.boolean(),
machine: MachinePreset,
attemptNumber: z.number(),
masterQueue: z.string(),
traceContext: z.record(z.unknown()),
}),
environment: z.object({
id: z.string(),
type: EnvironmentType,
}),
organization: z.object({
id: z.string(),
}),
project: z.object({
id: z.string(),
}),
});
export type DequeuedMessage = z.infer<typeof DequeuedMessage>;
/** The response to the Worker when starting an attempt */
export const StartRunAttemptResult = ExecutionResult.and(
z.object({
@@ -257,3 +211,51 @@ export const MachineResources = z.object({
memory: z.number(),
});
export type MachineResources = z.infer<typeof MachineResources>;
export const DequeueMessageCheckpoint = z.object({
id: z.string(),
type: CheckpointType,
location: z.string(),
imageRef: z.string(),
reason: z.string().nullish(),
});
export type DequeueMessageCheckpoint = z.infer<typeof DequeueMessageCheckpoint>;
/** This is sent to a Worker when a run is dequeued (a new run or continuing run) */
export const DequeuedMessage = z.object({
version: z.literal("1"),
snapshot: ExecutionSnapshot,
dequeuedAt: z.coerce.date(),
image: z.string().optional(),
checkpoint: DequeueMessageCheckpoint.optional(),
completedWaitpoints: z.array(CompletedWaitpoint),
backgroundWorker: z.object({
id: z.string(),
friendlyId: z.string(),
version: z.string(),
}),
deployment: z.object({
id: z.string().optional(),
friendlyId: z.string().optional(),
}),
run: z.object({
id: z.string(),
friendlyId: z.string(),
isTest: z.boolean(),
machine: MachinePreset,
attemptNumber: z.number(),
masterQueue: z.string(),
traceContext: z.record(z.unknown()),
}),
environment: z.object({
id: z.string(),
type: EnvironmentType,
}),
organization: z.object({
id: z.string(),
}),
project: z.object({
id: z.string(),
}),
});
export type DequeuedMessage = z.infer<typeof DequeuedMessage>;
@@ -4,69 +4,51 @@ import {
CheckpointServiceSuspendResponseBody,
CheckpointServiceRestoreRequestBodyInput,
} from "../schemas/checkpoints.js";
import { DequeuedMessage } from "../schemas/runEngine.js";
import { CheckpointType, DequeuedMessage } from "../schemas/runEngine.js";
import { SimpleStructuredLogger } from "../utils/structuredLogger.js";
export type CheckpointClientOptions = {
apiUrl: URL;
workerClient: SupervisorHttpClient;
orchestrator: CheckpointType;
};
export class CheckpointClient {
private readonly logger = new SimpleStructuredLogger("checkpoint-client");
private readonly apiUrl: URL;
private readonly workerClient: SupervisorHttpClient;
private get restoreUrl() {
return new URL("/api/v1/restore", this.apiUrl);
}
private get suspendUrl() {
return new URL("/api/v1/suspend", this.apiUrl);
}
constructor(opts: CheckpointClientOptions) {
this.apiUrl = opts.apiUrl;
this.workerClient = opts.workerClient;
}
constructor(private readonly opts: CheckpointClientOptions) {}
async suspendRun({
runFriendlyId,
snapshotFriendlyId,
containerId,
runnerId,
body,
}: {
runFriendlyId: string;
snapshotFriendlyId: string;
containerId: string;
runnerId: string;
body: Omit<CheckpointServiceSuspendRequestBodyInput, "type">;
}): Promise<boolean> {
const completionUrl = this.workerClient.getSuspendCompletionUrl(
runFriendlyId,
snapshotFriendlyId,
runnerId
);
const res = await fetch(this.suspendUrl, {
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify({
type: "DOCKER",
containerId,
callbacks: {
completion: completionUrl.url,
headers: completionUrl.headers,
const res = await fetch(
new URL(
`/api/v1/runs/${runFriendlyId}/snapshots/${snapshotFriendlyId}/suspend`,
this.opts.apiUrl
),
{
method: "POST",
headers: {
"Content-Type": "application/json",
},
} satisfies CheckpointServiceSuspendRequestBodyInput),
});
body: JSON.stringify({
type: this.opts.orchestrator,
...body,
} satisfies CheckpointServiceSuspendRequestBodyInput),
}
);
if (!res.ok) {
this.logger.error("[CheckpointClient] Suspend request failed", {
runFriendlyId,
snapshotFriendlyId,
containerId,
body,
});
return false;
}
@@ -74,7 +56,7 @@ export class CheckpointClient {
this.logger.debug("[CheckpointClient] Suspend request success", {
runFriendlyId,
snapshotFriendlyId,
containerId,
body,
status: res.status,
contentType: res.headers.get("content-type"),
});
@@ -87,7 +69,7 @@ export class CheckpointClient {
this.logger.error("[CheckpointClient] Suspend response invalid", {
runFriendlyId,
snapshotFriendlyId,
containerId,
body,
data,
});
return false;
@@ -106,28 +88,31 @@ export class CheckpointClient {
async restoreRun({
runFriendlyId,
snapshotFriendlyId,
checkpoint,
body,
}: {
runFriendlyId: string;
snapshotFriendlyId: string;
checkpoint: NonNullable<DequeuedMessage["checkpoint"]>;
body: CheckpointServiceRestoreRequestBodyInput;
}): Promise<boolean> {
const res = await fetch(this.restoreUrl, {
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify({
type: "DOCKER",
containerId: checkpoint.location,
} satisfies CheckpointServiceRestoreRequestBodyInput),
});
const res = await fetch(
new URL(
`/api/v1/runs/${runFriendlyId}/snapshots/${snapshotFriendlyId}/restore`,
this.opts.apiUrl
),
{
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify(body),
}
);
if (!res.ok) {
this.logger.error("[CheckpointClient] Restore request failed", {
runFriendlyId,
snapshotFriendlyId,
checkpoint,
body,
});
return false;
}
+83 -14
View File
@@ -2,6 +2,8 @@ import { createServer, type IncomingMessage, type ServerResponse } from "node:ht
import { z } from "zod";
import { SimpleStructuredLogger } from "../utils/structuredLogger.js";
import { HttpReply, getJsonBody } from "../apps/http.js";
import { Registry, Histogram, Counter } from "prom-client";
import { tryCatch } from "../../utils.js";
const logger = new SimpleStructuredLogger("worker-http");
@@ -51,21 +53,72 @@ type RouteMap = Partial<{
type HttpServerOptions = {
port: number;
host: string;
metrics?: {
register?: Registry;
expose?: boolean;
collect?: boolean;
};
};
export class HttpServer {
private static httpRequestDuration?: Histogram;
private static httpRequestTotal?: Counter;
private readonly metricsRegister?: Registry;
private readonly port: number;
private readonly host: string;
private routes: RouteMap = {};
public readonly server: ReturnType<typeof createServer>;
constructor(options: HttpServerOptions) {
this.port = options.port;
this.host = options.host;
this.metricsRegister = options.metrics?.register;
const collectMetrics = options.metrics?.collect ?? true;
const exposeMetrics = options.metrics?.expose ?? false;
// Initialize metrics only if registry is provided and not already initialized
if (this.metricsRegister && collectMetrics) {
if (!HttpServer.httpRequestDuration) {
HttpServer.httpRequestDuration = new Histogram({
name: "http_request_duration_seconds",
help: "Duration of HTTP requests in seconds",
labelNames: ["method", "route", "status", "port", "host"],
registers: [this.metricsRegister],
});
}
if (!HttpServer.httpRequestTotal) {
HttpServer.httpRequestTotal = new Counter({
name: "http_requests_total",
help: "Total number of HTTP requests",
labelNames: ["method", "route", "status", "port", "host"],
registers: [this.metricsRegister],
});
}
}
if (exposeMetrics) {
// Register metrics route
this.route("/metrics", "GET", {
handler: async ({ reply }) => {
if (!this.metricsRegister) {
return reply.text("Metrics registry not found", 500);
}
return reply.text(
await this.metricsRegister.metrics(),
200,
this.metricsRegister.contentType
);
},
});
}
this.server = createServer(async (req, res) => {
const reply = new HttpReply(res);
const startTime = process.hrtime();
try {
const { url, method } = req;
@@ -98,13 +151,6 @@ export class HttpServer {
const routeDefinition = this.routes[route]?.[httpMethod.data];
// logger.debug("Matched route", {
// url,
// method,
// route,
// routeDefinition,
// });
if (!routeDefinition) {
logger.error("Invalid method", { url, method, parsedMethod: httpMethod.data });
return reply.empty(405);
@@ -141,25 +187,29 @@ export class HttpServer {
return reply.json({ ok: false, error: "Invalid body" }, false, 400);
}
try {
await handler({
const [error] = await tryCatch(
handler({
reply,
req,
res,
params: parsedParams.data,
queryParams: parsedQueryParams.data,
body: parsedBody.data,
});
} catch (handlerError) {
logger.error("Route handler error", { error: handlerError });
})
);
if (error) {
logger.error("Route handler error", { error });
return reply.empty(500);
}
} catch (error) {
logger.error("Failed to handle request", { error });
return reply.empty(500);
} finally {
this.collectMetrics(req, res, startTime);
}
return;
return reply.empty(501);
});
this.server.on("clientError", (_, socket) => {
@@ -197,6 +247,25 @@ export class HttpServer {
});
}
private collectMetrics(req: IncomingMessage, res: ServerResponse, startTime: [number, number]) {
if (!this.metricsRegister || !HttpServer.httpRequestDuration || !HttpServer.httpRequestTotal) {
return;
}
const [seconds, nanoseconds] = process.hrtime(startTime);
const duration = seconds + nanoseconds / 1e9;
const route = this.findRoute(req.url ?? "") ?? "unknown";
const method = req.method ?? "unknown";
const status = res.statusCode.toString();
HttpServer.httpRequestDuration.observe(
{ method, route, status, port: this.port, host: this.host },
duration
);
HttpServer.httpRequestTotal.inc({ method, route, status, port: this.port, host: this.host });
}
private optionalSchema<
TSchema extends z.ZodFirstPartySchemaTypes | undefined,
TData extends TSchema extends z.ZodFirstPartySchemaTypes ? z.TypeOf<TSchema> : TData,
+3
View File
@@ -1,3 +1,6 @@
export * from "./checkpointClient.js";
export * from "./checkpointTest.js";
export * from "./httpServer.js";
export * from "./singleton.js";
export * from "./shutdownManager.js";
export * from "./k8s.js";
+16
View File
@@ -0,0 +1,16 @@
import { env } from "std-env";
export function isKubernetesEnvironment(override?: boolean): boolean {
if (override !== undefined) {
return override;
}
// Then check for common Kubernetes environment variables
const k8sIndicators = [
env.KUBERNETES_PORT,
env.KUBERNETES_SERVICE_HOST,
env.KUBERNETES_SERVICE_PORT,
];
return k8sIndicators.some(Boolean);
}
@@ -0,0 +1,183 @@
import { describe, test, expect, vi, beforeEach } from "vitest";
import { ShutdownManager } from "./shutdownManager.js";
describe("ShutdownManager", { concurrent: false }, () => {
// Mock process.exit to prevent actual exit
const mockExit = vi.spyOn(process, "exit").mockImplementation(() => undefined as never);
beforeEach(() => {
vi.clearAllMocks();
});
test("should successfully register a new handler", () => {
const manager = new ShutdownManager(false);
const handler = vi.fn();
manager.register("test-handler", handler);
expect(manager._getHandlersForTesting().has("test-handler")).toBe(true);
const registeredHandler = manager._getHandlersForTesting().get("test-handler");
expect(registeredHandler?.handler).toBe(handler);
expect(registeredHandler?.signals).toEqual(["SIGTERM", "SIGINT"]);
});
test("should throw error when registering duplicate handler name", () => {
const manager = new ShutdownManager(false);
const handler = vi.fn();
manager.register("duplicate-handler", handler);
expect(() => {
manager.register("duplicate-handler", handler);
}).toThrow('Shutdown handler "duplicate-handler" already registered');
});
test("should register handler with custom signals", () => {
const manager = new ShutdownManager(false);
const handler = vi.fn();
manager.register("custom-signals", handler, ["SIGTERM"]);
const registeredHandler = manager._getHandlersForTesting().get("custom-signals");
expect(registeredHandler?.signals).toEqual(["SIGTERM"]);
});
test("should call registered handlers when shutdown is triggered", async () => {
const manager = new ShutdownManager(false);
const handler1 = vi.fn();
const handler2 = vi.fn();
manager.register("handler1", handler1);
manager.register("handler2", handler2);
await manager.shutdown("SIGTERM");
expect(handler1).toHaveBeenCalledWith("SIGTERM");
expect(handler2).toHaveBeenCalledWith("SIGTERM");
expect(mockExit).toHaveBeenCalledWith(128 + 15); // SIGTERM number
});
test("should only call handlers registered for specific signal", async () => {
const manager = new ShutdownManager(false);
const handler1 = vi.fn();
const handler2 = vi.fn();
manager.register("handler1", handler1, ["SIGTERM"]);
manager.register("handler2", handler2, ["SIGINT"]);
await manager.shutdown("SIGTERM");
expect(handler1).toHaveBeenCalledWith("SIGTERM");
expect(handler2).not.toHaveBeenCalled();
expect(mockExit).toHaveBeenCalledWith(128 + 15);
});
test("should handle errors in shutdown handlers gracefully", async () => {
const manager = new ShutdownManager(false);
const successHandler = vi.fn();
const errorHandler = vi.fn().mockRejectedValue(new Error("Handler failed"));
manager.register("success-handler", successHandler);
manager.register("error-handler", errorHandler);
await manager.shutdown("SIGTERM");
expect(successHandler).toHaveBeenCalledWith("SIGTERM");
expect(errorHandler).toHaveBeenCalledWith("SIGTERM");
expect(mockExit).toHaveBeenCalledWith(128 + 15);
});
test("should only run shutdown sequence once even if called multiple times", async () => {
const manager = new ShutdownManager(false);
const handler = vi.fn();
manager.register("test-handler", handler);
await Promise.all([manager.shutdown("SIGTERM"), manager.shutdown("SIGTERM")]);
expect(handler).toHaveBeenCalledTimes(1);
expect(mockExit).toHaveBeenCalledTimes(1);
expect(mockExit).toHaveBeenCalledWith(128 + 15);
});
test("should exit with correct signal number on SIGINT", async () => {
const manager = new ShutdownManager(false);
manager.register("test-handler", vi.fn());
await manager.shutdown("SIGINT");
expect(mockExit).toHaveBeenCalledWith(128 + 2); // SIGINT number
});
test("should exit with correct signal number on SIGTERM", async () => {
const manager = new ShutdownManager(false);
manager.register("test-handler", vi.fn());
await manager.shutdown("SIGTERM");
expect(mockExit).toHaveBeenCalledWith(128 + 15); // SIGTERM number
});
test("should only exit after all handlers have finished", async () => {
const sequence: string[] = [];
const manager = new ShutdownManager(false);
const handler1 = vi.fn().mockImplementation(async () => {
sequence.push("handler1 start");
await new Promise((resolve) => setTimeout(resolve, 10));
sequence.push("handler1 end");
});
const handler2 = vi.fn().mockImplementation(async () => {
sequence.push("handler2 start");
await new Promise((resolve) => setTimeout(resolve, 20));
sequence.push("handler2 end");
});
const handler3 = vi.fn().mockImplementation(async () => {
sequence.push("handler3 start");
await new Promise((resolve) => setTimeout(resolve, 5));
sequence.push("handler3 end");
});
// Store the current mock implementation
const currentExit = mockExit.getMockImplementation();
// Override with our sequence-tracking implementation
mockExit.mockImplementation((code?: number | string | null) => {
sequence.push("exit");
return undefined as never;
});
manager.register("handler1", handler1);
manager.register("handler2", handler2);
manager.register("handler3", handler3);
await manager.shutdown("SIGTERM");
// Verify the execution order
expect(sequence).toEqual([
"handler1 start",
"handler2 start",
"handler3 start",
"handler3 end",
"handler1 end",
"handler2 end",
"exit",
]);
// Verify the handlers were called with correct arguments
expect(handler1).toHaveBeenCalledWith("SIGTERM");
expect(handler2).toHaveBeenCalledWith("SIGTERM");
expect(handler3).toHaveBeenCalledWith("SIGTERM");
expect(mockExit).toHaveBeenCalledWith(128 + 15);
// Restore original mock implementation
if (currentExit) {
mockExit.mockImplementation(currentExit);
}
});
});
@@ -0,0 +1,112 @@
import { isTest } from "std-env";
import { SimpleStructuredLogger } from "../utils/structuredLogger.js";
import { singleton } from "./singleton.js";
type ShutdownHandler = NodeJS.SignalsListener;
// We intentionally keep these limited to avoid unexpected issues with signal handling
type ShutdownSignal = Extract<NodeJS.Signals, "SIGTERM" | "SIGINT">;
export class ShutdownManager {
private isShuttingDown = false;
private signalNumbers: Record<ShutdownSignal, number> = {
SIGINT: 2,
SIGTERM: 15,
};
private logger = new SimpleStructuredLogger("shutdownManager");
private handlers: Map<string, { handler: ShutdownHandler; signals: ShutdownSignal[] }> =
new Map();
constructor(private disableForTesting = true) {
if (disableForTesting) return;
process.on("SIGTERM", () => this.shutdown("SIGTERM"));
process.on("SIGINT", () => this.shutdown("SIGINT"));
}
register(
name: string,
handler: ShutdownHandler,
signals: ShutdownSignal[] = ["SIGTERM", "SIGINT"]
) {
if (!this.isEnabled()) return;
if (this.handlers.has(name)) {
throw new Error(`Shutdown handler "${name}" already registered`);
}
this.handlers.set(name, { handler, signals });
}
unregister(name: string) {
if (!this.isEnabled()) return;
if (!this.handlers.has(name)) {
throw new Error(`Shutdown handler "${name}" not registered`);
}
this.handlers.delete(name);
}
async shutdown(signal: ShutdownSignal) {
if (!this.isEnabled()) return;
if (this.isShuttingDown) return;
this.isShuttingDown = true;
this.logger.info(`Received ${signal}. Starting graceful shutdown...`);
// Get handlers that are registered for this signal
const handlersToRun = Array.from(this.handlers.entries()).filter(([_, { signals }]) =>
signals.includes(signal)
);
try {
const results = await Promise.allSettled(
handlersToRun.map(async ([name, { handler }]) => {
try {
this.logger.info(`Running shutdown handler: ${name}`);
await handler(signal);
this.logger.info(`Shutdown handler completed: ${name}`);
} catch (error) {
this.logger.error(`Shutdown handler failed: ${name}`, { error });
throw error;
}
})
);
// Log any failures
results.forEach((result, index) => {
if (result.status === "rejected") {
const handlerEntry = handlersToRun[index];
if (handlerEntry) {
const [name] = handlerEntry;
this.logger.error(`Shutdown handler "${name}" failed:`, { reason: result.reason });
}
}
});
} catch (error) {
this.logger.error("Error during shutdown:", { error });
} finally {
// Exit with the correct signal number
process.exit(128 + this.signalNumbers[signal]);
}
}
private isEnabled() {
if (!this.disableForTesting) {
return true;
}
return !isTest;
}
// Only for testing
public _getHandlersForTesting(): ReadonlyMap<
string,
{ handler: ShutdownHandler; signals: ShutdownSignal[] }
> {
return new Map(this.handlers);
}
}
export const shutdownManager = singleton("shutdownManager", () => new ShutdownManager());
@@ -0,0 +1,8 @@
export function singleton<T>(name: string, getValue: () => T): T {
const thusly = globalThis as unknown as {
__trigger_singletons: Record<string, T>;
};
thusly.__trigger_singletons ??= {};
thusly.__trigger_singletons[name] ??= getValue();
return thusly.__trigger_singletons[name];
}
+56
View File
@@ -0,0 +1,56 @@
{
"name": "@trigger.dev/redis-worker",
"version": "3.3.17",
"description": "Redis worker for trigger.dev",
"license": "MIT",
"publishConfig": {
"access": "public"
},
"repository": {
"type": "git",
"url": "https://github.com/triggerdotdev/trigger.dev",
"directory": "packages/redis-worker"
},
"type": "module",
"files": [
"dist"
],
"scripts": {
"clean": "rimraf dist .turbo",
"build": "tsup",
"dev": "tsup --watch",
"typecheck": "tsc --noEmit -p tsconfig.src.json",
"test": "vitest --sequence.concurrent=false --no-file-parallelism"
},
"dependencies": {
"@trigger.dev/core": "workspace:*",
"lodash.omit": "^4.5.0",
"nanoid": "^5.0.7",
"p-limit": "^6.2.0",
"prom-client": "^15.1.0",
"zod": "3.23.8"
},
"devDependencies": {
"@internal/redis": "workspace:*",
"@internal/testcontainers": "workspace:*",
"@internal/tracing": "workspace:*",
"@types/lodash.omit": "^4.5.7",
"rimraf": "6.0.1",
"tsup": "^8.4.0",
"tsx": "4.17.0",
"vitest": "^1.4.0"
},
"engines": {
"node": ">=18.20.0"
},
"main": "./dist/index.cjs",
"module": "./dist/index.js",
"types": "./dist/index.d.ts",
"exports": {
".": {
"types": "./dist/index.d.ts",
"import": "./dist/index.js",
"require": "./dist/index.cjs"
}
}
}
@@ -30,6 +30,11 @@ describe("SimpleQueue", () => {
expect(await queue.size()).toBe(2);
const [first] = await queue.dequeue(1);
if (!first) {
throw new Error("No item dequeued");
}
expect(first).toEqual(
expect.objectContaining({
id: "1",
@@ -47,6 +52,11 @@ describe("SimpleQueue", () => {
expect(await queue.size({ includeFuture: true })).toBe(1);
const [second] = await queue.dequeue(1);
if (!second) {
throw new Error("No item dequeued");
}
expect(second).toEqual(
expect.objectContaining({
id: "2",
@@ -253,6 +263,10 @@ describe("SimpleQueue", () => {
expect(await queue.size()).toBe(1);
expect(await queue.size({ includeFuture: true })).toBe(3);
if (!dequeued[0] || !dequeued[1]) {
throw new Error("No items dequeued");
}
await queue.ack(dequeued[0].id);
await queue.ack(dequeued[1].id);
@@ -270,6 +284,10 @@ describe("SimpleQueue", () => {
})
);
if (!last) {
throw new Error("No item dequeued");
}
await queue.ack(last.id);
expect(await queue.size({ includeFuture: true })).toBe(0);
} finally {
@@ -315,6 +333,11 @@ describe("SimpleQueue", () => {
// Dequeue the redriven item
const [redrivenItem] = await queue.dequeue(1);
if (!redrivenItem) {
throw new Error("No item dequeued");
}
expect(redrivenItem).toEqual(
expect.objectContaining({
id: "1",
@@ -8,6 +8,8 @@ import { AnyQueueItem, SimpleQueue } from "./queue.js";
import { nanoid } from "nanoid";
import pLimit from "p-limit";
import { createRedisClient } from "@internal/redis";
import { shutdownManager } from "@trigger.dev/core/v3/serverOnly";
import { Registry, Histogram } from "prom-client";
export type WorkerCatalog = {
[key: string]: {
@@ -44,8 +46,12 @@ type WorkerOptions<TCatalog extends WorkerCatalog> = {
concurrency?: WorkerConcurrencyOptions;
pollIntervalMs?: number;
immediatePollIntervalMs?: number;
shutdownTimeoutMs?: number;
logger?: Logger;
tracer?: Tracer;
metrics?: {
register: Registry;
};
};
// This results in attempt 12 being a delay of 1 hour
@@ -63,12 +69,23 @@ class Worker<TCatalog extends WorkerCatalog> {
private subscriber: Redis | undefined;
private tracer: Tracer;
private metrics: {
register?: Registry;
enqueueDuration?: Histogram;
dequeueDuration?: Histogram;
jobDuration?: Histogram;
ackDuration?: Histogram;
redriveDuration?: Histogram;
rescheduleDuration?: Histogram;
} = {};
queue: SimpleQueue<QueueCatalogFromWorkerCatalog<TCatalog>>;
private jobs: WorkerOptions<TCatalog>["jobs"];
private logger: Logger;
private workerLoops: Promise<void>[] = [];
private isShuttingDown = false;
private concurrency: Required<NonNullable<WorkerOptions<TCatalog>["concurrency"]>>;
private shutdownTimeoutMs: number;
// The p-limit limiter to control overall concurrency.
private limiter: ReturnType<typeof pLimit>;
@@ -77,6 +94,8 @@ class Worker<TCatalog extends WorkerCatalog> {
this.logger = options.logger ?? new Logger("Worker", "debug");
this.tracer = options.tracer ?? trace.getTracer(options.name);
this.shutdownTimeoutMs = options.shutdownTimeoutMs ?? 60_000;
const schema: QueueCatalogFromWorkerCatalog<TCatalog> = Object.fromEntries(
Object.entries(this.options.catalog).map(([key, value]) => [key, value.schema])
) as QueueCatalogFromWorkerCatalog<TCatalog>;
@@ -95,6 +114,61 @@ class Worker<TCatalog extends WorkerCatalog> {
// Create a p-limit instance using this limit.
this.limiter = pLimit(this.concurrency.limit);
this.metrics.register = options.metrics?.register;
if (!this.metrics.register) {
return;
}
this.metrics.enqueueDuration = new Histogram({
name: "redis_worker_enqueue_duration_seconds",
help: "The duration of enqueue operations",
labelNames: ["worker_name", "job_type", "has_available_at"],
buckets: [0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1],
registers: [this.metrics.register],
});
this.metrics.dequeueDuration = new Histogram({
name: "redis_worker_dequeue_duration_seconds",
help: "The duration of dequeue operations",
labelNames: ["worker_name", "worker_id", "task_count"],
buckets: [0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1],
registers: [this.metrics.register],
});
this.metrics.jobDuration = new Histogram({
name: "redis_worker_job_duration_seconds",
help: "The duration of job operations",
labelNames: ["worker_name", "worker_id", "batch_size", "job_type", "attempt"],
// use different buckets here as jobs can take a while to run
buckets: [0.1, 0.25, 0.5, 1, 2.5, 5, 10, 20, 30, 45, 60],
registers: [this.metrics.register],
});
this.metrics.ackDuration = new Histogram({
name: "redis_worker_ack_duration_seconds",
help: "The duration of ack operations",
labelNames: ["worker_name"],
buckets: [0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1],
registers: [this.metrics.register],
});
this.metrics.redriveDuration = new Histogram({
name: "redis_worker_redrive_duration_seconds",
help: "The duration of redrive operations",
labelNames: ["worker_name"],
buckets: [0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1],
registers: [this.metrics.register],
});
this.metrics.rescheduleDuration = new Histogram({
name: "redis_worker_reschedule_duration_seconds",
help: "The duration of reschedule operations",
labelNames: ["worker_name"],
buckets: [0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1],
registers: [this.metrics.register],
});
}
public start() {
@@ -147,22 +221,33 @@ class Worker<TCatalog extends WorkerCatalog> {
this.tracer,
"enqueue",
async (span) => {
const timeout = visibilityTimeoutMs ?? this.options.catalog[job].visibilityTimeoutMs;
const timeout = visibilityTimeoutMs ?? this.options.catalog[job]?.visibilityTimeoutMs;
if (!timeout) {
throw new Error(`No visibility timeout found for job ${String(job)} with id ${id}`);
}
span.setAttribute("job_visibility_timeout_ms", timeout);
return this.queue.enqueue({
id,
job,
item: payload,
visibilityTimeoutMs: timeout,
availableAt,
});
return this.withHistogram(
this.metrics.enqueueDuration,
this.queue.enqueue({
id,
job,
item: payload,
visibilityTimeoutMs: timeout,
availableAt,
}),
{
job_type: String(job),
has_available_at: availableAt ? "true" : "false",
}
);
},
{
kind: SpanKind.PRODUCER,
attributes: {
job_type: job as string,
job_type: String(job),
job_id: id,
},
}
@@ -178,7 +263,10 @@ class Worker<TCatalog extends WorkerCatalog> {
this.tracer,
"reschedule",
async (span) => {
return this.queue.reschedule(id, availableAt);
return this.withHistogram(
this.metrics.rescheduleDuration,
this.queue.reschedule(id, availableAt)
);
},
{
kind: SpanKind.PRODUCER,
@@ -194,7 +282,7 @@ class Worker<TCatalog extends WorkerCatalog> {
this.tracer,
"ack",
() => {
return this.queue.ack(id);
return this.withHistogram(this.metrics.ackDuration, this.queue.ack(id));
},
{
attributes: {
@@ -220,7 +308,14 @@ class Worker<TCatalog extends WorkerCatalog> {
}
try {
const items = await this.queue.dequeue(taskCount);
const items = await this.withHistogram(
this.metrics.dequeueDuration,
this.queue.dequeue(taskCount),
{
worker_id: workerId,
task_count: taskCount,
}
);
if (items.length === 0) {
await Worker.delay(pollIntervalMs);
@@ -265,7 +360,17 @@ class Worker<TCatalog extends WorkerCatalog> {
this.tracer,
"processItem",
async () => {
await handler({ id, payload: item, visibilityTimeoutMs, attempt });
await this.withHistogram(
this.metrics.jobDuration,
handler({ id, payload: item, visibilityTimeoutMs, attempt }),
{
worker_id: workerId,
batch_size: batchSize,
job_type: job,
attempt,
}
);
// On success, acknowledge the item.
await this.queue.ack(id);
},
@@ -301,7 +406,7 @@ class Worker<TCatalog extends WorkerCatalog> {
const newAttempt = attempt + 1;
const retrySettings = {
...defaultRetrySettings,
...catalogItem.retry,
...catalogItem?.retry,
};
const retryDelay = calculateNextRetryDelay(retrySettings, newAttempt);
@@ -354,6 +459,23 @@ class Worker<TCatalog extends WorkerCatalog> {
});
}
private async withHistogram<T>(
histogram: Histogram<string> | undefined,
promise: Promise<T>,
labels?: Record<string, string | number>
): Promise<T> {
if (!histogram || !this.metrics.register) {
return promise;
}
const end = histogram.startTimer({ worker_name: this.options.name, ...labels });
try {
return await promise;
} finally {
end();
}
}
// A simple helper to delay for a given number of milliseconds.
private static delay(ms: number): Promise<void> {
return new Promise((resolve) => setTimeout(resolve, ms));
@@ -378,7 +500,10 @@ class Worker<TCatalog extends WorkerCatalog> {
if (typeof id !== "string") {
throw new Error("Invalid message format: id must be a string");
}
await this.queue.redriveFromDeadLetterQueue(id);
await this.withHistogram(
this.metrics.redriveDuration,
this.queue.redriveFromDeadLetterQueue(id)
);
this.logger.log(`Redrived item ${id} from Dead Letter Queue`);
} catch (error) {
this.logger.error("Error processing redrive message", { error, message });
@@ -386,25 +511,37 @@ class Worker<TCatalog extends WorkerCatalog> {
}
private setupShutdownHandlers() {
process.on("SIGTERM", this.shutdown.bind(this));
process.on("SIGINT", this.shutdown.bind(this));
shutdownManager.register(`redis-worker:${this.options.name}`, this.shutdown.bind(this));
}
private async shutdown() {
if (this.isShuttingDown) return;
private async shutdown(signal?: NodeJS.Signals) {
if (this.isShuttingDown) {
this.logger.log("Worker already shutting down", { signal });
return;
}
this.isShuttingDown = true;
this.logger.log("Shutting down worker loops...");
this.logger.log("Shutting down worker loops...", { signal });
// Wait for all worker loops to finish.
await Promise.all(this.workerLoops);
await Promise.race([
Promise.all(this.workerLoops),
Worker.delay(this.shutdownTimeoutMs).then(() => {
this.logger.error("Worker shutdown timed out", {
signal,
shutdownTimeoutMs: this.shutdownTimeoutMs,
});
}),
]);
await this.subscriber?.unsubscribe();
await this.subscriber?.quit();
await this.queue.close();
this.logger.log("All workers and subscribers shut down.");
this.logger.log("All workers and subscribers shut down.", { signal });
}
public async stop() {
shutdownManager.unregister(`redis-worker:${this.options.name}`);
await this.shutdown();
}
}
+11
View File
@@ -0,0 +1,11 @@
{
"extends": "../../.configs/tsconfig.base.json",
"references": [
{
"path": "./tsconfig.src.json"
},
{
"path": "./tsconfig.test.json"
}
]
}
+10
View File
@@ -0,0 +1,10 @@
{
"extends": "./tsconfig.json",
"include": ["./src/**/*.ts"],
"compilerOptions": {
"isolatedDeclarations": false,
"composite": true,
"sourceMap": true,
"customConditions": ["@triggerdotdev/source"]
}
}
+11
View File
@@ -0,0 +1,11 @@
{
"extends": "./tsconfig.json",
"include": ["./test/**/*.ts"],
"references": [{ "path": "./tsconfig.src.json" }],
"compilerOptions": {
"isolatedDeclarations": false,
"composite": true,
"sourceMap": true,
"types": ["vitest/globals"]
}
}
+27
View File
@@ -0,0 +1,27 @@
import { defineConfig } from "tsup";
export default defineConfig({
entry: ["src/index.ts"],
format: ["cjs", "esm"],
dts: true,
splitting: false,
sourcemap: true,
clean: true,
treeshake: true,
bundle: true,
minify: false,
noExternal: [
// Always bundle internal packages
/^@internal/,
// Always bundle ESM-only packages
"nanoid",
"p-limit",
],
banner: ({ format }) => {
if (format !== "esm") return;
return {
js: `import { createRequire } from 'module'; const require = createRequire(import.meta.url || process.cwd() + '/index.js');`,
};
},
});
+16
View File
@@ -0,0 +1,16 @@
diff --git a/index.js b/index.js
index 89f1b2cb86fad204b0493da3b8a3d5ed28937260..945f4cff27ed501fca75e269dfd7172e74c9c955 100644
--- a/index.js
+++ b/index.js
@@ -62,6 +62,11 @@ function supportsHyperlink(stream) {
return false;
}
+ // Cursor supports hyperlinks
+ if ("CURSOR_TRACE_ID" in env) {
+ return true;
+ }
+
if ('TERM_PROGRAM' in env) {
const version = parseVersion(env.TERM_PROGRAM_VERSION);
+671 -64
View File
File diff suppressed because it is too large Load Diff
+15 -4
View File
@@ -35,6 +35,17 @@ else
echo "Git status is clean. Proceeding with the script.";
fi
# From here on, if the user aborts the script, we will clean up the git stage
git_reset() {
git reset --hard HEAD
}
abort() {
echo "Aborted. Cleaning up..."
git_reset
exit 1
}
trap abort INT
# Run your commands
# Run changeset version command and capture its output
echo "Running: pnpm exec changeset version --snapshot $version"
@@ -48,6 +59,8 @@ else
exit 1
fi
read -e -p "Pausing for manual changes, press Enter when ready to continue..."
echo "Running: pnpm run clean --filter \"@trigger.dev/*\" --filter \"trigger.dev\""
pnpm run clean --filter "@trigger.dev/*" --filter "trigger.dev"
@@ -59,11 +72,9 @@ read -p "Do you wish to continue? (y/N): " prompt
if [[ $prompt =~ [yY](es)* ]]; then
pnpm exec changeset publish --no-git-tag --snapshot --tag $version
else
echo "Publish command aborted by the user."
git reset --hard HEAD
exit 1;
abort
fi
# If there were no errors, clear the git stage
echo "Commands ran successfully. Clearing the git stage."
git reset --hard HEAD
git_reset