diff --git a/.changeset/clean-pianos-listen.md b/.changeset/clean-pianos-listen.md new file mode 100644 index 000000000..fd6177147 --- /dev/null +++ b/.changeset/clean-pianos-listen.md @@ -0,0 +1,8 @@ +--- +"@trigger.dev/sdk": patch +"@trigger.dev/core-apps": patch +"trigger.dev": patch +"@trigger.dev/core": patch +--- + +add machine config and secure zod connection diff --git a/apps/coordinator/.env.example b/apps/coordinator/.env.example index 472939071..77377ab3c 100644 --- a/apps/coordinator/.env.example +++ b/apps/coordinator/.env.example @@ -1,3 +1,4 @@ HTTP_SERVER_PORT=8020 PLATFORM_ENABLED=true -PLATFORM_WS_PORT=3030 \ No newline at end of file +PLATFORM_WS_PORT=3030 +SECURE_CONNECTION=false \ No newline at end of file diff --git a/apps/coordinator/src/index.ts b/apps/coordinator/src/index.ts index feb2722ba..176bc5087 100644 --- a/apps/coordinator/src/index.ts +++ b/apps/coordinator/src/index.ts @@ -26,6 +26,7 @@ const PLATFORM_ENABLED = ["1", "true"].includes(process.env.PLATFORM_ENABLED ?? const PLATFORM_HOST = process.env.PLATFORM_HOST || "127.0.0.1"; const PLATFORM_WS_PORT = process.env.PLATFORM_WS_PORT || 3030; const PLATFORM_SECRET = process.env.PLATFORM_SECRET || "coordinator-secret"; +const SECURE_CONNECTION = ["1", "true"].includes(process.env.SECURE_CONNECTION ?? "false"); const logger = new SimpleLogger(`[${NODE_NAME}]`); @@ -365,6 +366,7 @@ class TaskCoordinator { namespace: "coordinator", host: PLATFORM_HOST, port: Number(PLATFORM_WS_PORT), + secure: SECURE_CONNECTION, clientMessages: CoordinatorToPlatformMessages, serverMessages: PlatformToCoordinatorMessages, authToken: PLATFORM_SECRET, diff --git a/apps/docker-provider/.env.example b/apps/docker-provider/.env.example index 075b6823a..2d24f79c8 100644 --- a/apps/docker-provider/.env.example +++ b/apps/docker-provider/.env.example @@ -2,6 +2,7 @@ HTTP_SERVER_PORT=8050 PLATFORM_WS_PORT=3030 PLATFORM_SECRET=provider-secret +SECURE_CONNECTION=false # Use this if you are on macOS # COORDINATOR_HOST="host.docker.internal" diff --git a/apps/docker-provider/src/index.ts b/apps/docker-provider/src/index.ts index 7fed78137..7728449da 100644 --- a/apps/docker-provider/src/index.ts +++ b/apps/docker-provider/src/index.ts @@ -8,6 +8,7 @@ import { TaskOperationsIndexOptions, } from "@trigger.dev/core-apps"; import { setTimeout } from "node:timers/promises"; +import { PostStartCauses, PreStopCauses } from "@trigger.dev/core/v3"; const MACHINE_NAME = process.env.MACHINE_NAME || "local"; const COORDINATOR_PORT = process.env.COORDINATOR_PORT || 8020; @@ -190,6 +191,9 @@ class DockerTaskOperations implements TaskOperations { async delete(opts: { runId: string }) { await this.#initialize(); + const containerName = this.#getRunContainerName(opts.runId); + await this.#sendPreStop(containerName); + logger.log("noop: delete"); } @@ -208,6 +212,26 @@ class DockerTaskOperations implements TaskOperations { } async #sendPostStart(containerName: string): Promise { + try { + const port = await this.#getHttpServerPort(containerName); + logger.debug(await this.#runLifecycleCommand(containerName, port, "postStart", "restore")); + } catch (error) { + logger.error("postStart error", { error }); + throw new Error("postStart command failed"); + } + } + + async #sendPreStop(containerName: string): Promise { + try { + const port = await this.#getHttpServerPort(containerName); + logger.debug(await this.#runLifecycleCommand(containerName, port, "preStop", "terminate")); + } catch (error) { + logger.error("preStop error", { error }); + throw new Error("preStop command failed"); + } + } + + async #getHttpServerPort(containerName: string): Promise { // We first get the correct port, which is random during dev as we run with host networking and need to avoid clashes // FIXME: Skip this in prod const logs = logger.debug(await $`docker logs ${containerName}`); @@ -219,19 +243,14 @@ class DockerTaskOperations implements TaskOperations { throw new Error("failed to extract port from logs"); } - try { - logger.debug(await this.#runLifecycleCommand(containerName, port, "postStart", "restore")); - } catch (error) { - logger.error("postStart error", { error }); - throw new Error("postStart command failed"); - } + return port; } - async #runLifecycleCommand( + async #runLifecycleCommand( containerName: string, port: number, - type: "postStart" | "preStop", - cause: "index" | "create" | "restore", + type: THookType, + cause: THookType extends "postStart" ? PostStartCauses : PreStopCauses, retryCount = 0 ): Promise { try { @@ -244,15 +263,15 @@ class DockerTaskOperations implements TaskOperations { `127.0.0.1:${port}/${type}?cause=${cause}`, ]); } catch (error: any) { - if (retryCount < 6) { - logger.debug("retriable postStart error", { retryCount, message: error?.message }); + if (type === "postStart" && retryCount < 6) { + logger.debug(`retriable ${type} error`, { retryCount, message: error?.message }); await setTimeout(exponentialBackoff(retryCount + 1, 2, 50, 1150, 50)); return this.#runLifecycleCommand(containerName, port, type, cause, retryCount + 1); } - logger.error("final postStart error", { message: error?.message }); - throw new Error(`postStart command failed after ${retryCount - 1} retries`); + logger.error(`final ${type} error`, { message: error?.message }); + throw new Error(`${type} command failed after ${retryCount - 1} retries`); } } } diff --git a/apps/kubernetes-provider/.env.example b/apps/kubernetes-provider/.env.example index 10949a807..f21ee29ba 100644 --- a/apps/kubernetes-provider/.env.example +++ b/apps/kubernetes-provider/.env.example @@ -2,6 +2,7 @@ HTTP_SERVER_PORT=8060 PLATFORM_WS_PORT=3030 PLATFORM_SECRET=provider-secret +SECURE_CONNECTION=false # Use this if you are on macOS # COORDINATOR_HOST="host.docker.internal" diff --git a/apps/kubernetes-provider/src/index.ts b/apps/kubernetes-provider/src/index.ts index 5b4d290b6..5442e0677 100644 --- a/apps/kubernetes-provider/src/index.ts +++ b/apps/kubernetes-provider/src/index.ts @@ -7,6 +7,7 @@ import { TaskOperationsIndexOptions, TaskOperationsRestoreOptions, } from "@trigger.dev/core-apps"; +import { Machine, PostStartCauses, PreStopCauses, EnvironmentType } from "@trigger.dev/core/v3"; import { randomUUID } from "crypto"; const RUNTIME_ENV = process.env.KUBERNETES_PORT ? "kubernetes" : "local"; @@ -55,6 +56,12 @@ class KubernetesTaskOperations implements TaskOperations { metadata: { labels: { app: "task-index", + "app.kubernetes.io/part-of": "trigger-worker", + "app.kubernetes.io/component": "index", + env: opts.envId, + envtype: this.#envTypeToLabelValue(opts.envType), + org: opts.orgId, + project: opts.projectId, }, }, spec: { @@ -64,6 +71,9 @@ class KubernetesTaskOperations implements TaskOperations { name: "registry-trigger", }, ], + nodeSelector: { + nodetype: "worker", + }, containers: [ { name: this.#getIndexContainerName(opts.shortCode), @@ -79,6 +89,13 @@ class KubernetesTaskOperations implements TaskOperations { // memory: "50Mi", // }, // }, + lifecycle: { + preStop: { + exec: { + command: this.#getLifecycleCommand("preStop", "terminate"), + }, + }, + }, env: [ { name: "DEBUG", @@ -151,6 +168,13 @@ class KubernetesTaskOperations implements TaskOperations { namespace: this.#namespace.metadata.name, labels: { app: "task-run", + "app.kubernetes.io/part-of": "trigger-worker", + "app.kubernetes.io/component": "create", + env: opts.envId, + envtype: this.#envTypeToLabelValue(opts.envType), + org: opts.orgId, + project: opts.projectId, + run: opts.runId, }, }, spec: { @@ -160,6 +184,9 @@ class KubernetesTaskOperations implements TaskOperations { name: "registry-trigger", }, ], + nodeSelector: { + nodetype: "worker", + }, containers: [ { name: this.#getRunContainerName(opts.runId), @@ -169,9 +196,9 @@ class KubernetesTaskOperations implements TaskOperations { containerPort: 8000, }, ], - // resources: { - // limits: opts.machine, - // }, + resources: { + limits: this.#getResourcesFromMachineConfig(opts.machine), + }, lifecycle: { postStart: { exec: { @@ -180,7 +207,7 @@ class KubernetesTaskOperations implements TaskOperations { }, preStop: { exec: { - command: this.#getLifecycleCommand("preStop", "create"), + command: this.#getLifecycleCommand("preStop", "terminate"), }, }, }, @@ -262,6 +289,14 @@ class KubernetesTaskOperations implements TaskOperations { namespace: this.#namespace.metadata.name, labels: { app: "task-run", + "app.kubernetes.io/part-of": "trigger-worker", + "app.kubernetes.io/component": "restore", + env: opts.envId, + envtype: this.#envTypeToLabelValue(opts.envType), + org: opts.orgId, + project: opts.projectId, + run: opts.runId, + checkpoint: opts.checkpointId, }, }, spec: { @@ -271,6 +306,9 @@ class KubernetesTaskOperations implements TaskOperations { name: "registry-trigger", }, ], + nodeSelector: { + nodetype: "worker", + }, initContainers: [ { name: "pull-base-image", @@ -309,9 +347,9 @@ class KubernetesTaskOperations implements TaskOperations { containerPort: 8000, }, ], - // resources: { - // limits: opts.machine, - // }, + resources: { + limits: this.#getResourcesFromMachineConfig(opts.machine), + }, lifecycle: { postStart: { exec: { @@ -320,7 +358,7 @@ class KubernetesTaskOperations implements TaskOperations { }, preStop: { exec: { - command: this.#getLifecycleCommand("preStop", "restore"), + command: this.#getLifecycleCommand("preStop", "terminate"), }, }, }, @@ -355,7 +393,30 @@ class KubernetesTaskOperations implements TaskOperations { await this.#getPod(opts.runId, this.#namespace); } - #getLifecycleCommand(type: "postStart" | "preStop", cause: "index" | "create" | "restore") { + #envTypeToLabelValue(type: EnvironmentType) { + switch (type) { + case "PRODUCTION": + return "prod"; + case "STAGING": + return "stg"; + case "DEVELOPMENT": + return "dev"; + case "PREVIEW": + return "preview"; + } + } + + #getResourcesFromMachineConfig(config: Machine) { + return { + cpu: `${config.cpu}`, + memory: `${config.memory}G`, + }; + } + + #getLifecycleCommand( + type: THookType, + cause: THookType extends "postStart" ? PostStartCauses : PreStopCauses + ) { return ["/bin/sh", "-c", `sleep 1; wget -q -O- 127.0.0.1:8000/${type}?cause=${cause}`]; } diff --git a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts index 3b3fbd6b1..116a2320d 100644 --- a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts +++ b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts @@ -1,5 +1,6 @@ import { Context, ROOT_CONTEXT, Span, SpanKind, context, trace } from "@opentelemetry/api"; import { + Machine, ProdTaskRunExecution, ProdTaskRunExecutionPayload, TaskRunError, @@ -438,10 +439,27 @@ export class SharedQueueConsumer { queueId: queue.id, runtimeEnvironmentId: environment.id, }, + include: { + backgroundWorkerTask: true, + }, }); const isRetry = taskRunAttempt.number > 1; + const { machineConfig } = taskRunAttempt.backgroundWorkerTask; + const machine = Machine.safeParse(machineConfig ?? {}); + + if (!machine.success) { + logger.error("Failed to parse machine config", { + queueMessage: message.data, + messageId: message.messageId, + attemptId: taskRunAttempt.id, + machineConfig, + }); + + await this.#ackAndDoMoreWork(message.messageId); + return; + } try { if (messageBody.data.checkpointEventId) { const restoreService = new RestoreCheckpointService(); @@ -470,11 +488,16 @@ export class SharedQueueConsumer { backgroundWorkerId: deployment.worker.friendlyId, data: { type: "SCHEDULE_ATTEMPT", - id: taskRunAttempt.id, image: deployment.imageReference, - envId: environment.id, - runId: taskRunAttempt.taskRunId, version: deployment.version, + machine: machine.data, + // identifiers + id: taskRunAttempt.id, + envId: environment.id, + envType: environment.type, + orgId: environment.organizationId, + projectId: environment.projectId, + runId: taskRunAttempt.taskRunId, }, }); } diff --git a/apps/webapp/app/v3/services/createBackgroundWorker.server.ts b/apps/webapp/app/v3/services/createBackgroundWorker.server.ts index 00765d34a..7b7d4eec8 100644 --- a/apps/webapp/app/v3/services/createBackgroundWorker.server.ts +++ b/apps/webapp/app/v3/services/createBackgroundWorker.server.ts @@ -101,6 +101,7 @@ export async function createBackgroundTasks( exportName: task.exportName, retryConfig: task.retry, queueConfig: task.queue, + machineConfig: task.machine, }, }); diff --git a/apps/webapp/app/v3/services/indexDeployment.server.ts b/apps/webapp/app/v3/services/indexDeployment.server.ts index 854dcb034..8487a2845 100644 --- a/apps/webapp/app/v3/services/indexDeployment.server.ts +++ b/apps/webapp/app/v3/services/indexDeployment.server.ts @@ -49,9 +49,13 @@ export class IndexDeploymentService extends BaseService { version: "v1", shortCode: deployment.shortCode, imageTag: deployment.imageReference, - envId: deployment.environmentId, apiKey: deployment.environment.apiKey, apiUrl: env.APP_ORIGIN, + // identifiers + envId: deployment.environmentId, + envType: deployment.environment.type, + projectId: deployment.projectId, + orgId: deployment.environment.organizationId, }); logger.debug("Index ACK received", { responses }); diff --git a/apps/webapp/app/v3/services/restoreCheckpoint.server.ts b/apps/webapp/app/v3/services/restoreCheckpoint.server.ts index a61f95603..9e362f4d9 100644 --- a/apps/webapp/app/v3/services/restoreCheckpoint.server.ts +++ b/apps/webapp/app/v3/services/restoreCheckpoint.server.ts @@ -3,6 +3,7 @@ import { logger } from "~/services/logger.server"; import { socketIo } from "../handleSocketIo.server"; import { CreateCheckpointRestoreEventService } from "./createCheckpointRestoreEvent.server"; import { BaseService } from "./baseService.server"; +import { Machine } from "@trigger.dev/core/v3"; const RESTORABLE_RUN_STATUSES: TaskRunStatus[] = ["WAITING_TO_RESUME"]; const RESTORABLE_ATTEMPT_STATUSES: TaskRunAttemptStatus[] = ["PAUSED"]; @@ -30,8 +31,14 @@ export class RestoreCheckpointService extends BaseService { attempt: { select: { status: true, + backgroundWorkerTask: { + select: { + machineConfig: true, + }, + }, }, }, + runtimeEnvironment: true, }, }, }, @@ -63,17 +70,34 @@ export class RestoreCheckpointService extends BaseService { return; } + const { machineConfig } = checkpoint.attempt.backgroundWorkerTask; + const machine = Machine.safeParse(machineConfig ?? {}); + + if (!machine.success) { + logger.error("Failed to parse machine config", { + attemptId: checkpoint.attemptId, + machineConfig: checkpoint.attempt.backgroundWorkerTask.machineConfig, + }); + return; + } + const eventService = new CreateCheckpointRestoreEventService(this._prisma); await eventService.restore({ checkpointId: checkpoint.id }); socketIo.providerNamespace.emit("RESTORE", { version: "v1", - checkpointId: checkpoint.id, - runId: checkpoint.runId, type: checkpoint.type, location: checkpoint.location, reason: checkpoint.reason ?? undefined, imageRef: checkpoint.imageRef, + machine: machine.data, + // identifiers + checkpointId: checkpoint.id, + envId: checkpoint.runtimeEnvironment.id, + envType: checkpoint.runtimeEnvironment.type, + orgId: checkpoint.runtimeEnvironment.organizationId, + projectId: checkpoint.runtimeEnvironment.projectId, + runId: checkpoint.runId, }); return checkpoint; diff --git a/packages/cli-v3/src/workers/dev/worker-facade.ts b/packages/cli-v3/src/workers/dev/worker-facade.ts index 37fe0c6e4..b7585240b 100644 --- a/packages/cli-v3/src/workers/dev/worker-facade.ts +++ b/packages/cli-v3/src/workers/dev/worker-facade.ts @@ -84,6 +84,7 @@ function getTasks(): Array { filePath: (taskFile as any).filePath, queue: (task as any).__trigger.queue, retry: (task as any).__trigger.retry, + machine: (task as any).__trigger.machine, fns: (task as any).__trigger.fns, }); } diff --git a/packages/cli-v3/src/workers/prod/entry-point.ts b/packages/cli-v3/src/workers/prod/entry-point.ts index 39a853798..02162aa06 100644 --- a/packages/cli-v3/src/workers/prod/entry-point.ts +++ b/packages/cli-v3/src/workers/prod/entry-point.ts @@ -1,6 +1,8 @@ import { Config, CoordinatorToProdWorkerMessages, + PostStartCauses, + PreStopCauses, ProdWorkerToCoordinatorMessages, TaskResource, WaitReason, @@ -567,23 +569,15 @@ class ProdWorker { } case "/preStop": { - const schema = z.enum(["index", "create", "restore"]); - - const cause = schema.safeParse(url.searchParams.get("cause")); + const cause = PreStopCauses.safeParse(url.searchParams.get("cause")); if (!cause.success) { logger.error("Failed to parse cause", { cause }); - return; + return reply.text("Failed to parse cause", 400); } switch (cause.data) { - case "index": { - break; - } - case "create": { - break; - } - case "restore": { + case "terminate": { break; } default: { @@ -597,13 +591,11 @@ class ProdWorker { } case "/postStart": { - const schema = z.enum(["index", "create", "restore"]); - - const cause = schema.safeParse(url.searchParams.get("cause")); + const cause = PostStartCauses.safeParse(url.searchParams.get("cause")); if (!cause.success) { logger.error("Failed to parse cause", { cause }); - return; + return reply.text("Failed to parse cause", 400); } switch (cause.data) { @@ -678,11 +670,7 @@ class ProdWorker { } for (const task of this.#backgroundWorker.tasks) { - taskResources.push({ - id: task.id, - filePath: task.filePath, - exportName: task.exportName, - }); + taskResources.push(task); packageVersion = task.packageVersion; } diff --git a/packages/cli-v3/src/workers/prod/worker-facade.ts b/packages/cli-v3/src/workers/prod/worker-facade.ts index 21185c3b8..f25269fd3 100644 --- a/packages/cli-v3/src/workers/prod/worker-facade.ts +++ b/packages/cli-v3/src/workers/prod/worker-facade.ts @@ -78,6 +78,7 @@ function getTasks(): Array { filePath: (taskFile as any).filePath, queue: (task as any).__trigger.queue, retry: (task as any).__trigger.retry, + machine: (task as any).__trigger.machine, fns: (task as any).__trigger.fns, }); } diff --git a/packages/core-apps/src/provider.ts b/packages/core-apps/src/provider.ts index d304b4fa5..84e449aad 100644 --- a/packages/core-apps/src/provider.ts +++ b/packages/core-apps/src/provider.ts @@ -2,6 +2,7 @@ import { createServer } from "node:http"; import { ClientToSharedQueueMessages, clientWebsocketMessages, + EnvironmentType, Machine, PlatformToProviderMessages, ProviderToPlatformMessages, @@ -18,30 +19,46 @@ const MACHINE_NAME = process.env.MACHINE_NAME || "local"; const PLATFORM_HOST = process.env.PLATFORM_HOST || "127.0.0.1"; const PLATFORM_WS_PORT = process.env.PLATFORM_WS_PORT || 3030; const PLATFORM_SECRET = process.env.PLATFORM_SECRET || "provider-secret"; +const SECURE_CONNECTION = ["1", "true"].includes(process.env.SECURE_CONNECTION ?? "false"); const logger = new SimpleLogger(`[${MACHINE_NAME}]`); export interface TaskOperationsIndexOptions { shortCode: string; imageRef: string; - envId: string; apiKey: string; apiUrl: string; + // identifiers + envId: string; + envType: EnvironmentType; + orgId: string; + projectId: string; } export interface TaskOperationsCreateOptions { - runId: string; image: string; machine: Machine; - envId: string; version: string; + // identifiers + envId: string; + envType: EnvironmentType; + orgId: string; + projectId: string; + runId: string; + attemptId: string; } export interface TaskOperationsRestoreOptions { - runId: string; imageRef: string; checkpointRef: string; machine: Machine; + // identifiers + envId: string; + envType: EnvironmentType; + orgId: string; + projectId: string; + runId: string; + checkpointId: string; } export interface TaskOperations { @@ -87,6 +104,7 @@ export class ProviderShell implements Provider { namespace: "shared-queue", host: PLATFORM_HOST, port: Number(PLATFORM_WS_PORT), + secure: SECURE_CONNECTION, clientMessages: ClientToSharedQueueMessages, serverMessages: SharedQueueToClientMessages, authToken: PLATFORM_SECRET, @@ -101,11 +119,16 @@ export class ProviderShell implements Provider { if (message.data.type === "SCHEDULE_ATTEMPT") { try { this.tasks.create({ - envId: message.data.envId, - runId: message.data.runId, image: message.data.image, - machine: {}, - version: message.version, + machine: message.data.machine, + version: message.data.version, + // identifiers + envId: message.data.envId, + envType: message.data.envType, + orgId: message.data.orgId, + projectId: message.data.projectId, + runId: message.data.runId, + attemptId: message.data.id, }); } catch (error) { logger.error("create failed", error); @@ -138,6 +161,7 @@ export class ProviderShell implements Provider { namespace: "provider", host: PLATFORM_HOST, port: Number(PLATFORM_WS_PORT), + secure: SECURE_CONNECTION, clientMessages: ProviderToPlatformMessages, serverMessages: PlatformToProviderMessages, authToken: PLATFORM_SECRET, @@ -165,9 +189,13 @@ export class ProviderShell implements Provider { await this.tasks.index({ shortCode: message.shortCode, imageRef: message.imageTag, - envId: message.envId, apiKey: message.apiKey, apiUrl: message.apiUrl, + // identifiers + envId: message.envId, + envType: message.envType, + orgId: message.orgId, + projectId: message.projectId, }); } catch (error) { logger.error("index failed", error); @@ -206,13 +234,16 @@ export class ProviderShell implements Provider { try { await this.tasks.restore({ - runId: message.runId, checkpointRef: message.location, - machine: { - cpu: "1", - memory: "100Mi", - }, + machine: message.machine, imageRef: message.imageRef, + // identifiers + envId: message.envId, + envType: message.envType, + orgId: message.orgId, + projectId: message.projectId, + runId: message.runId, + checkpointId: message.checkpointId, }); } catch (error) { logger.error("restore failed", error); @@ -230,54 +261,34 @@ export class ProviderShell implements Provider { const reply = new HttpReply(res); - switch (req.url) { - case "/health": { - return reply.text("ok"); - } - case "/whoami": { - return reply.text(`${MACHINE_NAME}`); - } - case "/close": { - this.#platformSocket.close(); - return reply.text("platform socket closed"); - } - case "/delete": { - const body = await getTextBody(req); + try { + const url = new URL(req.url ?? "", `http://${req.headers.host}`); - await this.tasks.delete({ runId: body }); + switch (url.pathname) { + case "/health": { + return reply.text("ok"); + } + case "/whoami": { + return reply.text(`${MACHINE_NAME}`); + } + case "/close": { + this.#platformSocket.close(); + return reply.text("platform socket closed"); + } + case "/delete": { + const body = await getTextBody(req); - return reply.text(`sent delete request: ${body}`); - } - case "/invoke": { - const body = await getTextBody(req); - - await this.tasks.create({ - envId: "placeholder", - image: body, - machine: { - cpu: "1", - memory: "100Mi", - }, - runId: "", - version: "", - }); - - return reply.text(`sent restore request: ${body}`); - } - case "/restore": { - const body = await getTextBody(req); - - const items = body.split("&"); - const image = items[0]; - const baseImageTag = items[1] ?? image; - - // await this.tasks.restore({}); - - return reply.text(`sent restore request: ${body}`); - } - default: { - return reply.empty(404); + await this.tasks.delete({ runId: body }); + + return reply.text(`sent delete request: ${body}`); + } + default: { + return reply.empty(404); + } } + } catch (error) { + logger.error("HTTP server error", { error }); + reply.empty(500); } }); diff --git a/packages/core/src/v3/schemas/messages.ts b/packages/core/src/v3/schemas/messages.ts index f37ad1c27..8c0d9f0a6 100644 --- a/packages/core/src/v3/schemas/messages.ts +++ b/packages/core/src/v3/schemas/messages.ts @@ -1,6 +1,29 @@ import { z } from "zod"; import { TaskRunExecution, TaskRunExecutionResult } from "./common"; +export const EnvironmentType = z.enum(["PRODUCTION", "STAGING", "DEVELOPMENT", "PREVIEW"]) +export type EnvironmentType = z.infer; + +export const MachineCpu = z + .union([z.literal(0.25), z.literal(0.5), z.literal(1), z.literal(2), z.literal(4)]) + .default(0.5); + +export type MachineCpu = z.infer; + +export const MachineMemory = z + .union([z.literal(0.25), z.literal(0.5), z.literal(1), z.literal(2), z.literal(4), z.literal(8)]) + .default(1); + +export type MachineMemory = z.infer; + +export const Machine = z.object({ + version: z.literal("v1").default("v1"), + cpu: MachineCpu, + memory: MachineMemory, +}); + +export type Machine = z.infer; + export const TaskRunExecutionPayload = z.object({ execution: TaskRunExecution, traceContext: z.record(z.unknown()), @@ -39,11 +62,16 @@ export const BackgroundWorkerServerMessages = z.discriminatedUnion("type", [ }), z.object({ type: z.literal("SCHEDULE_ATTEMPT"), - id: z.string(), image: z.string(), - envId: z.string(), - runId: z.string(), version: z.string(), + machine: Machine, + // identifiers + id: z.string(), // attempt + envId: z.string(), + envType: EnvironmentType, + orgId: z.string(), + projectId: z.string(), + runId: z.string(), }), ]); @@ -229,6 +257,7 @@ export const TaskMetadata = z.object({ packageVersion: z.string(), queue: QueueOptions.optional(), retry: RetryOptions.optional(), + machine: Machine.partial().optional(), }); export type TaskMetadata = z.infer; diff --git a/packages/core/src/v3/schemas/resources.ts b/packages/core/src/v3/schemas/resources.ts index ae7e5c7f2..b105920f8 100644 --- a/packages/core/src/v3/schemas/resources.ts +++ b/packages/core/src/v3/schemas/resources.ts @@ -1,5 +1,5 @@ import { z } from "zod"; -import { QueueOptions, RetryOptions } from "./messages"; +import { Machine, QueueOptions, RetryOptions } from "./messages"; export const TaskResource = z.object({ id: z.string(), @@ -7,6 +7,7 @@ export const TaskResource = z.object({ exportName: z.string(), queue: QueueOptions.optional(), retry: RetryOptions.optional(), + machine: Machine.partial().optional(), }); export type TaskResource = z.infer; diff --git a/packages/core/src/v3/schemas/schemas.ts b/packages/core/src/v3/schemas/schemas.ts index 92fbf950b..ffdbcf135 100644 --- a/packages/core/src/v3/schemas/schemas.ts +++ b/packages/core/src/v3/schemas/schemas.ts @@ -7,9 +7,17 @@ import { ProdTaskRunExecution, ProdTaskRunExecutionPayload, RetryOptions, + Machine, + EnvironmentType, } from "./messages"; import { TaskResource } from "./resources"; +export const PostStartCauses = z.enum(["index", "create", "restore"]); +export type PostStartCauses = z.infer; + +export const PreStopCauses = z.enum(["terminate"]); +export type PreStopCauses = z.infer; + const RegexSchema = z.custom((val) => { try { // Check to see if val is a regex @@ -42,13 +50,6 @@ export type ResolvedConfig = RequireKeys< "triggerDirectories" | "triggerUrl" | "projectDir" | "tsconfigPath" >; -export const Machine = z.object({ - cpu: z.string().default("1").optional(), - memory: z.string().default("500Mi").optional(), -}); - -export type Machine = z.infer; - export const WaitReason = z.enum(["WAIT_FOR_DURATION", "WAIT_FOR_TASK", "WAIT_FOR_BATCH"]); export type WaitReason = z.infer; @@ -85,9 +86,13 @@ export const PlatformToProviderMessages = { version: z.literal("v1").default("v1"), imageTag: z.string(), shortCode: z.string(), - envId: z.string(), apiKey: z.string(), apiUrl: z.string(), + // identifiers + envId: z.string(), + envType: EnvironmentType, + orgId: z.string(), + projectId: z.string(), }), callback: z.discriminatedUnion("success", [ z.object({ @@ -103,22 +108,22 @@ export const PlatformToProviderMessages = { }), ]), }, - INVOKE: { - message: z.object({ - version: z.literal("v1").default("v1"), - name: z.string(), - machine: Machine, - }), - }, + // TODO: this should be a shared queue message instead RESTORE: { message: z.object({ version: z.literal("v1").default("v1"), - checkpointId: z.string(), - runId: z.string(), type: z.enum(["DOCKER", "KUBERNETES"]), location: z.string(), reason: z.string().optional(), imageRef: z.string(), + machine: Machine, + // identifiers + checkpointId: z.string(), + envId: z.string(), + envType: EnvironmentType, + orgId: z.string(), + projectId: z.string(), + runId: z.string(), }), }, DELETE: { diff --git a/packages/core/src/v3/zodSocket.ts b/packages/core/src/v3/zodSocket.ts index 294cdc89f..0898ce210 100644 --- a/packages/core/src/v3/zodSocket.ts +++ b/packages/core/src/v3/zodSocket.ts @@ -263,7 +263,8 @@ interface ZodSocketConnectionOptions< TServerMessages extends ZodSocketMessageCatalogSchema, > { host: string; - port: number; + port?: number; + secure?: boolean; namespace: string; clientMessages: TClientMessages; serverMessages: TServerMessages; @@ -302,15 +303,24 @@ export class ZodSocketConnection< #logger: StructuredLogger; constructor(opts: ZodSocketConnectionOptions) { - this.socket = io(`ws://${opts.host}:${opts.port}/${opts.namespace}`, { + const uri = `${opts.secure ? "wss" : "ws"}://${opts.host}:${ + opts.port ?? (opts.secure ? "443" : "80") + }/${opts.namespace}`; + + const logger = new SimpleStructuredLogger(opts.namespace, LogLevel.info); + logger.log("new zod socket", { uri }); + + this.socket = io(uri, { transports: ["websocket"], auth: { token: opts.authToken, }, extraHeaders: opts.extraHeaders, + reconnectionDelay: 500, + reconnectionDelayMax: 1000, }); - this.#logger = new SimpleStructuredLogger(opts.namespace, LogLevel.info, { + this.#logger = logger.child({ socketId: this.socket.id, }); diff --git a/packages/database/prisma/migrations/20240327121557_add_machine_config_to_worker_task/migration.sql b/packages/database/prisma/migrations/20240327121557_add_machine_config_to_worker_task/migration.sql new file mode 100644 index 000000000..ba5809af5 --- /dev/null +++ b/packages/database/prisma/migrations/20240327121557_add_machine_config_to_worker_task/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "BackgroundWorkerTask" ADD COLUMN "machineConfig" JSONB; diff --git a/packages/database/prisma/schema.prisma b/packages/database/prisma/schema.prisma index 6b7483d46..8e4b54bbe 100644 --- a/packages/database/prisma/schema.prisma +++ b/packages/database/prisma/schema.prisma @@ -1559,6 +1559,7 @@ model BackgroundWorkerTask { queueConfig Json? retryConfig Json? + machineConfig Json? @@unique([workerId, slug]) } diff --git a/packages/trigger-sdk/src/v3/shared.ts b/packages/trigger-sdk/src/v3/shared.ts index 3ac436f57..500a2f7fc 100644 --- a/packages/trigger-sdk/src/v3/shared.ts +++ b/packages/trigger-sdk/src/v3/shared.ts @@ -9,6 +9,8 @@ import { HandleErrorResult, InitFnParams, InitOutput, + MachineCpu, + MachineMemory, MiddlewareFnParams, QueueOptions, RetryOptions, @@ -114,8 +116,8 @@ export type TaskOptions( packageVersion: packageJson.version, queue: params.queue, retry: params.retry ? { ...defaultRetryOptions, ...params.retry } : undefined, + machine: params.machine, fns: { run: params.run, init: params.init,