chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,883 @@
|
||||
import { SnapshotCallbackPayloadSchema } from "@internal/compute";
|
||||
import { type CheckpointClient, HttpServer } from "@trigger.dev/core/v3/serverOnly";
|
||||
import { SimpleStructuredLogger } from "@trigger.dev/core/v3/utils/structuredLogger";
|
||||
import {
|
||||
type WorkloadRunSnapshotsSinceResponseBody,
|
||||
type SupervisorHttpClient,
|
||||
WORKLOAD_HEADERS,
|
||||
type WorkloadClientSocketData,
|
||||
type WorkloadClientToServerEvents,
|
||||
type WorkloadContinueRunExecutionResponseBody,
|
||||
WorkloadDebugLogRequestBody,
|
||||
type WorkloadDequeueFromVersionResponseBody,
|
||||
WorkloadHeartbeatRequestBody,
|
||||
type WorkloadHeartbeatResponseBody,
|
||||
WorkloadRunAttemptCompleteRequestBody,
|
||||
type WorkloadRunAttemptCompleteResponseBody,
|
||||
WorkloadRunAttemptStartRequestBody,
|
||||
type WorkloadRunAttemptStartResponseBody,
|
||||
type WorkloadServerToClientEvents,
|
||||
type WorkloadSuspendRunResponseBody,
|
||||
} from "@trigger.dev/core/v3/workers";
|
||||
import EventEmitter from "node:events";
|
||||
import type { IncomingMessage, ServerResponse } from "node:http";
|
||||
import { type Namespace, Server, type Socket } from "socket.io";
|
||||
import { z } from "zod";
|
||||
import { env } from "../env.js";
|
||||
import { register } from "../metrics.js";
|
||||
import {
|
||||
ComputeSnapshotService,
|
||||
type RunTraceContext,
|
||||
} from "../services/computeSnapshotService.js";
|
||||
import type { OtlpTraceService } from "../services/otlpTraceService.js";
|
||||
import {
|
||||
emitOneShot,
|
||||
runWideEvent,
|
||||
setMeta,
|
||||
type State,
|
||||
type WideEventOptions,
|
||||
} from "../wideEvents/index.js";
|
||||
import type { ComputeWorkloadManager } from "../workloadManager/compute.js";
|
||||
|
||||
// Use the official export when upgrading to socket.io@4.8.0
|
||||
interface DefaultEventsMap {
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
[event: string]: (...args: any[]) => void;
|
||||
}
|
||||
|
||||
const WorkloadActionParams = z.object({
|
||||
runFriendlyId: z.string(),
|
||||
snapshotFriendlyId: z.string(),
|
||||
});
|
||||
|
||||
// Workloads bundled into customer task images before CLI v4.4.4 use a strict
|
||||
// zod enum for checkpoint type that only allows DOCKER and KUBERNETES. The
|
||||
// workload never reads this field - it only validates the response shape - so
|
||||
// rewriting it to a known value keeps older runners working without affecting
|
||||
// the value stored in the database or seen by internal services.
|
||||
function legacifyCheckpointType<T extends { checkpoint?: { type: string } | null }>(item: T): T {
|
||||
if (item.checkpoint?.type === "COMPUTE") {
|
||||
return { ...item, checkpoint: { ...item.checkpoint, type: "KUBERNETES" } } as T;
|
||||
}
|
||||
return item;
|
||||
}
|
||||
|
||||
type WorkloadServerEvents = {
|
||||
runConnected: [
|
||||
{
|
||||
run: {
|
||||
friendlyId: string;
|
||||
};
|
||||
},
|
||||
];
|
||||
runDisconnected: [
|
||||
{
|
||||
run: {
|
||||
friendlyId: string;
|
||||
};
|
||||
},
|
||||
];
|
||||
};
|
||||
|
||||
type WorkloadServerOptions = {
|
||||
port: number;
|
||||
host?: string;
|
||||
workerClient: SupervisorHttpClient;
|
||||
checkpointClient?: CheckpointClient;
|
||||
computeManager?: ComputeWorkloadManager;
|
||||
tracing?: OtlpTraceService;
|
||||
wideEventOpts: WideEventOptions;
|
||||
/** When true, high-frequency HTTP routes also emit wide events. */
|
||||
wideEventsNoisyRoutes: boolean;
|
||||
};
|
||||
|
||||
export class WorkloadServer extends EventEmitter<WorkloadServerEvents> {
|
||||
private checkpointClient?: CheckpointClient;
|
||||
private readonly snapshotService?: ComputeSnapshotService;
|
||||
|
||||
private readonly logger = new SimpleStructuredLogger("workload-server");
|
||||
private readonly wideEventOpts: WideEventOptions;
|
||||
private readonly wideEventsNoisyRoutes: boolean;
|
||||
|
||||
private readonly httpServer: HttpServer;
|
||||
private readonly websocketServer: Namespace<
|
||||
WorkloadClientToServerEvents,
|
||||
WorkloadServerToClientEvents,
|
||||
DefaultEventsMap,
|
||||
WorkloadClientSocketData
|
||||
>;
|
||||
|
||||
private readonly runSockets = new Map<
|
||||
string,
|
||||
Socket<
|
||||
WorkloadClientToServerEvents,
|
||||
WorkloadServerToClientEvents,
|
||||
DefaultEventsMap,
|
||||
WorkloadClientSocketData
|
||||
>
|
||||
>();
|
||||
|
||||
private readonly workerClient: SupervisorHttpClient;
|
||||
|
||||
constructor(opts: WorkloadServerOptions) {
|
||||
super();
|
||||
|
||||
const host = opts.host ?? "0.0.0.0";
|
||||
const port = opts.port;
|
||||
|
||||
this.workerClient = opts.workerClient;
|
||||
this.checkpointClient = opts.checkpointClient;
|
||||
this.wideEventOpts = opts.wideEventOpts;
|
||||
this.wideEventsNoisyRoutes = opts.wideEventsNoisyRoutes;
|
||||
|
||||
if (opts.computeManager?.snapshotsEnabled) {
|
||||
this.snapshotService = new ComputeSnapshotService({
|
||||
computeManager: opts.computeManager,
|
||||
workerClient: opts.workerClient,
|
||||
tracing: opts.tracing,
|
||||
wideEventOpts: this.wideEventOpts,
|
||||
});
|
||||
}
|
||||
|
||||
this.httpServer = this.createHttpServer({ host, port });
|
||||
this.websocketServer = this.createWebsocketServer();
|
||||
}
|
||||
|
||||
private headerValueFromRequest(req: IncomingMessage, headerName: string): string | undefined {
|
||||
const value = req.headers[headerName];
|
||||
|
||||
if (Array.isArray(value)) {
|
||||
return value[0];
|
||||
}
|
||||
|
||||
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);
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets common route meta on the wide-event state from URL params.
|
||||
*/
|
||||
private attachRouteMeta(state: State, params: unknown): void {
|
||||
if (!params || typeof params !== "object") return;
|
||||
const p = params as Record<string, unknown>;
|
||||
if (typeof p.runFriendlyId === "string") setMeta(state, "run_id", p.runFriendlyId);
|
||||
if (typeof p.snapshotFriendlyId === "string") {
|
||||
setMeta(state, "snapshot_id", p.snapshotFriendlyId);
|
||||
}
|
||||
if (typeof p.deploymentId === "string") setMeta(state, "deployment_id", p.deploymentId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Wraps an HTTP route handler body with the wide-event lifecycle. Reads
|
||||
* `traceparent` and `x-request-id` from `req.headers`, attaches `run_id` /
|
||||
* `snapshot_id` / `deployment_id` meta from `params` when present, and
|
||||
* captures the response status from `res.statusCode` after `fn` returns.
|
||||
*
|
||||
* Pass `highFrequency: true` for noisy routes (heartbeat, polling). Those
|
||||
* still go through the wrapper but only emit when
|
||||
* `TRIGGER_WIDE_EVENTS_NOISY_ROUTES` is on, so prod can keep them dark
|
||||
* while test envs capture full-fidelity traffic for debugging.
|
||||
*/
|
||||
private wideRoute<T>(
|
||||
ctx: { req: IncomingMessage; res: ServerResponse; params?: unknown },
|
||||
op: string,
|
||||
route: string,
|
||||
method: string,
|
||||
fn: () => Promise<T> | T,
|
||||
routeOpts: { highFrequency?: boolean } = {}
|
||||
): Promise<T> {
|
||||
const enabled =
|
||||
this.wideEventOpts.enabled && (!routeOpts.highFrequency || this.wideEventsNoisyRoutes);
|
||||
return runWideEvent(
|
||||
{
|
||||
...this.wideEventOpts,
|
||||
enabled,
|
||||
op,
|
||||
kind: "inbound",
|
||||
route,
|
||||
method,
|
||||
traceparent: this.headerValueFromRequest(ctx.req, "traceparent"),
|
||||
inboundRequestId: this.headerValueFromRequest(ctx.req, "x-request-id"),
|
||||
setup: (state) => this.attachRouteMeta(state, ctx.params),
|
||||
},
|
||||
fn,
|
||||
(state) => {
|
||||
state.statusCode = ctx.res.statusCode;
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
private createHttpServer({ host, port }: { host: string; port: number }) {
|
||||
const httpServer = new HttpServer({
|
||||
port,
|
||||
host,
|
||||
metrics: {
|
||||
register,
|
||||
expose: false,
|
||||
},
|
||||
})
|
||||
.route("/health", "GET", {
|
||||
handler: async ({ reply }) => {
|
||||
reply.text("OK");
|
||||
},
|
||||
})
|
||||
.route(
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/attempts/start",
|
||||
"POST",
|
||||
{
|
||||
paramsSchema: WorkloadActionParams,
|
||||
bodySchema: WorkloadRunAttemptStartRequestBody,
|
||||
handler: async (ctx) =>
|
||||
this.wideRoute(
|
||||
ctx,
|
||||
"attempt.start",
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/attempts/start",
|
||||
"POST",
|
||||
async () => {
|
||||
const { req, reply, params, body } = ctx;
|
||||
const startResponse = await this.workerClient.startRunAttempt(
|
||||
params.runFriendlyId,
|
||||
params.snapshotFriendlyId,
|
||||
body,
|
||||
this.runnerIdFromRequest(req)
|
||||
);
|
||||
|
||||
if (!startResponse.success) {
|
||||
this.logger.error("Failed to start run", {
|
||||
params,
|
||||
error: startResponse.error,
|
||||
});
|
||||
reply.empty(500);
|
||||
return;
|
||||
}
|
||||
|
||||
reply.json(startResponse.data satisfies WorkloadRunAttemptStartResponseBody);
|
||||
return;
|
||||
}
|
||||
),
|
||||
}
|
||||
)
|
||||
.route(
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/attempts/complete",
|
||||
"POST",
|
||||
{
|
||||
paramsSchema: WorkloadActionParams,
|
||||
bodySchema: WorkloadRunAttemptCompleteRequestBody,
|
||||
handler: async (ctx) =>
|
||||
this.wideRoute(
|
||||
ctx,
|
||||
"attempt.complete",
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/attempts/complete",
|
||||
"POST",
|
||||
async () => {
|
||||
const { req, reply, params, body } = ctx;
|
||||
const runnerId = this.runnerIdFromRequest(req);
|
||||
|
||||
// A completion attempt invalidates any pending delayed snapshot
|
||||
// regardless of outcome: the runner has finished executing, so the
|
||||
// suspended state the snapshot was scheduled to capture no longer
|
||||
// exists. Cancel BEFORE the async completion call - the timer
|
||||
// wheel can tick during the await, so cancelling after it leaves
|
||||
// a real window for a due snapshot to dispatch and pause a VM
|
||||
// that has moved on. The runnerId guard keeps a stale duplicate
|
||||
// runner's completion from cancelling a fresh runner's snapshot,
|
||||
// and the runner can't schedule a new suspend until it receives
|
||||
// this route's reply, so nothing legitimate can be cancelled here.
|
||||
this.snapshotService?.cancel(params.runFriendlyId, runnerId);
|
||||
|
||||
const completeResponse = await this.workerClient.completeRunAttempt(
|
||||
params.runFriendlyId,
|
||||
params.snapshotFriendlyId,
|
||||
body,
|
||||
runnerId
|
||||
);
|
||||
|
||||
if (!completeResponse.success) {
|
||||
this.logger.error("Failed to complete run", {
|
||||
params,
|
||||
error: completeResponse.error,
|
||||
});
|
||||
reply.empty(500);
|
||||
return;
|
||||
}
|
||||
|
||||
reply.json(completeResponse.data satisfies WorkloadRunAttemptCompleteResponseBody);
|
||||
return;
|
||||
}
|
||||
),
|
||||
}
|
||||
)
|
||||
.route(
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/heartbeat",
|
||||
"POST",
|
||||
{
|
||||
paramsSchema: WorkloadActionParams,
|
||||
bodySchema: WorkloadHeartbeatRequestBody,
|
||||
handler: async (ctx) =>
|
||||
this.wideRoute(
|
||||
ctx,
|
||||
"heartbeat",
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/heartbeat",
|
||||
"POST",
|
||||
async () => {
|
||||
const { req, reply, params, body } = ctx;
|
||||
const heartbeatResponse = await this.workerClient.heartbeatRun(
|
||||
params.runFriendlyId,
|
||||
params.snapshotFriendlyId,
|
||||
body,
|
||||
this.runnerIdFromRequest(req)
|
||||
);
|
||||
|
||||
if (!heartbeatResponse.success) {
|
||||
this.logger.error("Failed to heartbeat run", {
|
||||
params,
|
||||
error: heartbeatResponse.error,
|
||||
});
|
||||
reply.empty(500);
|
||||
return;
|
||||
}
|
||||
|
||||
reply.json({
|
||||
ok: true,
|
||||
} satisfies WorkloadHeartbeatResponseBody);
|
||||
},
|
||||
{ highFrequency: true }
|
||||
),
|
||||
}
|
||||
)
|
||||
.route(
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/suspend",
|
||||
"GET",
|
||||
{
|
||||
paramsSchema: WorkloadActionParams,
|
||||
handler: async (ctx) =>
|
||||
this.wideRoute(
|
||||
ctx,
|
||||
"suspend",
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/suspend",
|
||||
"GET",
|
||||
async () => {
|
||||
const { reply, params, req } = ctx;
|
||||
const runnerId = this.runnerIdFromRequest(req);
|
||||
const deploymentVersion = this.deploymentVersionFromRequest(req);
|
||||
const projectRef = this.projectRefFromRequest(req);
|
||||
|
||||
this.logger.debug("Suspend request", {
|
||||
params,
|
||||
runnerId,
|
||||
deploymentVersion,
|
||||
projectRef,
|
||||
});
|
||||
|
||||
if (!runnerId || !deploymentVersion || !projectRef) {
|
||||
this.logger.error("Invalid headers for suspend request", {
|
||||
...params,
|
||||
runnerId,
|
||||
deploymentVersion,
|
||||
projectRef,
|
||||
});
|
||||
reply.json(
|
||||
{
|
||||
ok: false,
|
||||
error: "Invalid headers",
|
||||
} satisfies WorkloadSuspendRunResponseBody,
|
||||
false,
|
||||
400
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
if (this.snapshotService) {
|
||||
// Compute mode: delay snapshot to avoid wasted work on short-lived waitpoints.
|
||||
// If the run continues before the delay expires, the snapshot is cancelled.
|
||||
reply.json({ ok: true } satisfies WorkloadSuspendRunResponseBody, false, 202);
|
||||
|
||||
this.snapshotService.schedule(params.runFriendlyId, {
|
||||
runnerId,
|
||||
runFriendlyId: params.runFriendlyId,
|
||||
snapshotFriendlyId: params.snapshotFriendlyId,
|
||||
});
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
if (!this.checkpointClient) {
|
||||
reply.json(
|
||||
{
|
||||
ok: false,
|
||||
error: "Checkpoints disabled",
|
||||
} satisfies WorkloadSuspendRunResponseBody,
|
||||
false,
|
||||
400
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
reply.json(
|
||||
{
|
||||
ok: true,
|
||||
} satisfies WorkloadSuspendRunResponseBody,
|
||||
false,
|
||||
202
|
||||
);
|
||||
|
||||
const suspendResult = await this.checkpointClient.suspendRun({
|
||||
runFriendlyId: params.runFriendlyId,
|
||||
snapshotFriendlyId: params.snapshotFriendlyId,
|
||||
body: {
|
||||
runnerId,
|
||||
runId: params.runFriendlyId,
|
||||
snapshotId: params.snapshotFriendlyId,
|
||||
projectRef,
|
||||
deploymentVersion,
|
||||
},
|
||||
});
|
||||
|
||||
if (!suspendResult) {
|
||||
this.logger.error("Failed to suspend run", { params });
|
||||
return;
|
||||
}
|
||||
}
|
||||
),
|
||||
}
|
||||
)
|
||||
.route(
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/continue",
|
||||
"GET",
|
||||
{
|
||||
paramsSchema: WorkloadActionParams,
|
||||
handler: async (ctx) =>
|
||||
this.wideRoute(
|
||||
ctx,
|
||||
"continue",
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/:snapshotFriendlyId/continue",
|
||||
"GET",
|
||||
async () => {
|
||||
const { req, reply, params } = ctx;
|
||||
this.logger.debug("Run continuation request", { params });
|
||||
|
||||
// Cancel any pending delayed snapshot for this run
|
||||
this.snapshotService?.cancel(params.runFriendlyId);
|
||||
|
||||
const continuationResult = await this.workerClient.continueRunExecution(
|
||||
params.runFriendlyId,
|
||||
params.snapshotFriendlyId,
|
||||
this.runnerIdFromRequest(req)
|
||||
);
|
||||
|
||||
if (!continuationResult.success) {
|
||||
this.logger.error("Failed to continue run execution", { params });
|
||||
reply.json(
|
||||
{
|
||||
ok: false,
|
||||
error: "Failed to continue run execution",
|
||||
},
|
||||
false,
|
||||
400
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
reply.json(continuationResult.data as WorkloadContinueRunExecutionResponseBody);
|
||||
}
|
||||
),
|
||||
}
|
||||
)
|
||||
.route(
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/since/:snapshotFriendlyId",
|
||||
"GET",
|
||||
{
|
||||
paramsSchema: WorkloadActionParams,
|
||||
handler: async (ctx) =>
|
||||
this.wideRoute(
|
||||
ctx,
|
||||
"snapshots.since",
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/snapshots/since/:snapshotFriendlyId",
|
||||
"GET",
|
||||
async () => {
|
||||
const { req, reply, params } = ctx;
|
||||
const sinceSnapshotResponse = await this.workerClient.getSnapshotsSince(
|
||||
params.runFriendlyId,
|
||||
params.snapshotFriendlyId,
|
||||
this.runnerIdFromRequest(req)
|
||||
);
|
||||
|
||||
if (!sinceSnapshotResponse.success) {
|
||||
this.logger.error("Failed to get snapshots since", {
|
||||
runId: params.runFriendlyId,
|
||||
error: sinceSnapshotResponse.error,
|
||||
});
|
||||
reply.empty(500);
|
||||
return;
|
||||
}
|
||||
|
||||
reply.json({
|
||||
snapshots: sinceSnapshotResponse.data.snapshots.map(legacifyCheckpointType),
|
||||
} satisfies WorkloadRunSnapshotsSinceResponseBody);
|
||||
},
|
||||
{ highFrequency: true }
|
||||
),
|
||||
}
|
||||
)
|
||||
.route("/api/v1/workload-actions/deployments/:deploymentId/dequeue", "GET", {
|
||||
paramsSchema: z.object({
|
||||
deploymentId: z.string(),
|
||||
}),
|
||||
|
||||
handler: async (ctx) =>
|
||||
this.wideRoute(
|
||||
ctx,
|
||||
"deployment.dequeue",
|
||||
"/api/v1/workload-actions/deployments/:deploymentId/dequeue",
|
||||
"GET",
|
||||
async () => {
|
||||
const { req, reply, params } = ctx;
|
||||
const dequeueResponse = await this.workerClient.dequeueFromVersion(
|
||||
params.deploymentId,
|
||||
1,
|
||||
this.runnerIdFromRequest(req)
|
||||
);
|
||||
|
||||
if (!dequeueResponse.success) {
|
||||
this.logger.error("Failed to get latest snapshot", {
|
||||
deploymentId: params.deploymentId,
|
||||
error: dequeueResponse.error,
|
||||
});
|
||||
reply.empty(500);
|
||||
return;
|
||||
}
|
||||
|
||||
reply.json(
|
||||
dequeueResponse.data.map(
|
||||
legacifyCheckpointType
|
||||
) satisfies WorkloadDequeueFromVersionResponseBody
|
||||
);
|
||||
}
|
||||
),
|
||||
});
|
||||
|
||||
if (env.SEND_RUN_DEBUG_LOGS) {
|
||||
httpServer.route("/api/v1/workload-actions/runs/:runFriendlyId/logs/debug", "POST", {
|
||||
paramsSchema: WorkloadActionParams.pick({ runFriendlyId: true }),
|
||||
bodySchema: WorkloadDebugLogRequestBody,
|
||||
handler: async (ctx) =>
|
||||
this.wideRoute(
|
||||
ctx,
|
||||
"logs.debug",
|
||||
"/api/v1/workload-actions/runs/:runFriendlyId/logs/debug",
|
||||
"POST",
|
||||
async () => {
|
||||
const { req, reply, params, body } = ctx;
|
||||
reply.empty(204);
|
||||
|
||||
await this.workerClient.sendDebugLog(
|
||||
params.runFriendlyId,
|
||||
body,
|
||||
this.runnerIdFromRequest(req)
|
||||
);
|
||||
},
|
||||
{ highFrequency: true }
|
||||
),
|
||||
});
|
||||
} else {
|
||||
// Disabled: drop immediately without reading/parsing the body and without
|
||||
// any log we can't switch off. Older runners still POST per log line; the
|
||||
// route stays registered (an unregistered route would log "No route match"
|
||||
// per request) but sheds the request at minimal cost. Request metrics still
|
||||
// count it. 204 is non-retryable on the runner client, so no retry storm.
|
||||
httpServer.route("/api/v1/workload-actions/runs/:runFriendlyId/logs/debug", "POST", {
|
||||
skipBodyParsing: true,
|
||||
handler: async (ctx) => {
|
||||
ctx.reply.empty(204);
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
// Snapshot callback endpoint (inbound from compute path)
|
||||
httpServer.route("/api/v1/compute/snapshot-complete", "POST", {
|
||||
bodySchema: SnapshotCallbackPayloadSchema,
|
||||
handler: async (ctx) =>
|
||||
this.wideRoute(
|
||||
ctx,
|
||||
"snapshot.callback",
|
||||
"/api/v1/compute/snapshot-complete",
|
||||
"POST",
|
||||
async () => {
|
||||
const { reply, body } = ctx;
|
||||
if (!this.snapshotService) {
|
||||
reply.empty(404);
|
||||
return;
|
||||
}
|
||||
|
||||
const result = await this.snapshotService.handleCallback(body);
|
||||
reply.empty(result.status);
|
||||
}
|
||||
),
|
||||
});
|
||||
|
||||
return httpServer;
|
||||
}
|
||||
|
||||
private createWebsocketServer() {
|
||||
const io = new Server(this.httpServer.server);
|
||||
|
||||
const websocketServer: Namespace<
|
||||
WorkloadClientToServerEvents,
|
||||
WorkloadServerToClientEvents,
|
||||
DefaultEventsMap,
|
||||
WorkloadClientSocketData
|
||||
> = io.of("/workload");
|
||||
|
||||
websocketServer.on("disconnect", (socket) => {
|
||||
this.logger.verbose("[WS] disconnect", socket.id);
|
||||
});
|
||||
websocketServer.use(async (socket, next) => {
|
||||
const setSocketDataFromHeader = (
|
||||
dataKey: keyof typeof socket.data,
|
||||
headerName: string,
|
||||
required: boolean = true
|
||||
) => {
|
||||
const value = socket.handshake.headers[headerName];
|
||||
|
||||
if (value) {
|
||||
if (Array.isArray(value)) {
|
||||
if (value[0]) {
|
||||
socket.data[dataKey] = value[0];
|
||||
return;
|
||||
}
|
||||
} else {
|
||||
socket.data[dataKey] = value;
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
if (required) {
|
||||
this.logger.error("[WS] missing required header", { headerName });
|
||||
throw new Error("missing header");
|
||||
}
|
||||
};
|
||||
|
||||
try {
|
||||
setSocketDataFromHeader("deploymentId", WORKLOAD_HEADERS.DEPLOYMENT_ID);
|
||||
setSocketDataFromHeader("runnerId", WORKLOAD_HEADERS.RUNNER_ID);
|
||||
} catch (error) {
|
||||
this.logger.error("[WS] setSocketDataFromHeader error", { error });
|
||||
socket.disconnect(true);
|
||||
return;
|
||||
}
|
||||
|
||||
this.logger.debug("[WS] auth success", socket.data);
|
||||
|
||||
next();
|
||||
});
|
||||
websocketServer.on("connection", (socket) => {
|
||||
const socketLogger = this.logger.child({
|
||||
socketId: socket.id,
|
||||
socketData: socket.data,
|
||||
});
|
||||
|
||||
const getSocketMetadata = () => {
|
||||
return {
|
||||
deploymentId: socket.data.deploymentId,
|
||||
runId: socket.data.runFriendlyId,
|
||||
snapshotId: socket.data.snapshotId,
|
||||
runnerId: socket.data.runnerId,
|
||||
};
|
||||
};
|
||||
|
||||
const emitSocketLifecycle = (
|
||||
event: "run_connected" | "run_disconnected",
|
||||
friendlyId: string,
|
||||
disconnectReason?: string
|
||||
) => {
|
||||
emitOneShot({
|
||||
...this.wideEventOpts,
|
||||
op: event === "run_connected" ? "socket.run.connected" : "socket.run.disconnected",
|
||||
kind: "event",
|
||||
populate: (state) => {
|
||||
state.extras.event = event;
|
||||
setMeta(state, "run_id", friendlyId);
|
||||
if (socket.data.deploymentId) {
|
||||
setMeta(state, "deployment_id", socket.data.deploymentId);
|
||||
}
|
||||
if (socket.data.runnerId) setMeta(state, "runner_id", socket.data.runnerId);
|
||||
state.extras.socket_id = socket.id;
|
||||
if (disconnectReason) state.extras.disconnect_reason = disconnectReason;
|
||||
},
|
||||
});
|
||||
};
|
||||
|
||||
const runConnected = (friendlyId: string) => {
|
||||
socketLogger.debug("runConnected", { ...getSocketMetadata() });
|
||||
|
||||
// If there's already a run ID set, we should "disconnect" it from this socket
|
||||
if (socket.data.runFriendlyId && socket.data.runFriendlyId !== friendlyId) {
|
||||
socketLogger.debug("runConnected: disconnecting existing run", {
|
||||
...getSocketMetadata(),
|
||||
newRunId: friendlyId,
|
||||
oldRunId: socket.data.runFriendlyId,
|
||||
});
|
||||
runDisconnected(socket.data.runFriendlyId, "socket_run_replaced");
|
||||
}
|
||||
|
||||
this.runSockets.set(friendlyId, socket);
|
||||
this.emit("runConnected", { run: { friendlyId } });
|
||||
socket.data.runFriendlyId = friendlyId;
|
||||
emitSocketLifecycle("run_connected", friendlyId);
|
||||
};
|
||||
|
||||
const runDisconnected = (friendlyId: string, reason: string) => {
|
||||
socketLogger.debug("runDisconnected", { ...getSocketMetadata() });
|
||||
|
||||
// The run is gone from this runner (crash, exit, or replaced by a new
|
||||
// run), so a pending delayed snapshot for it is stale. Genuine
|
||||
// waitpoint suspensions keep the socket connected, so this doesn't
|
||||
// cancel a snapshot that's still wanted; the runnerId match guards
|
||||
// against a stale duplicate runner cancelling a fresh runner's
|
||||
// snapshot after the run was reassigned. Caveat: socket.data.runnerId
|
||||
// is frozen at the websocket handshake, so after a same-supervisor
|
||||
// restore (new runner id, socket not recreated) this guard refuses
|
||||
// the cancel - a missed cancel, never a wrong one. The
|
||||
// attempt.complete cancel uses the runner's current HTTP header id
|
||||
// and is unaffected.
|
||||
this.snapshotService?.cancel(friendlyId, socket.data.runnerId);
|
||||
|
||||
this.runSockets.delete(friendlyId);
|
||||
this.emit("runDisconnected", { run: { friendlyId } });
|
||||
socket.data.runFriendlyId = undefined;
|
||||
emitSocketLifecycle("run_disconnected", friendlyId, reason);
|
||||
};
|
||||
|
||||
socketLogger.debug("wsServer socket connected", { ...getSocketMetadata() });
|
||||
|
||||
// FIXME: where does this get set?
|
||||
if (socket.data.runFriendlyId) {
|
||||
runConnected(socket.data.runFriendlyId);
|
||||
}
|
||||
|
||||
socket.on("disconnecting", (reason, description) => {
|
||||
socketLogger.verbose("Socket disconnecting", {
|
||||
...getSocketMetadata(),
|
||||
reason,
|
||||
description,
|
||||
});
|
||||
|
||||
if (socket.data.runFriendlyId) {
|
||||
runDisconnected(socket.data.runFriendlyId, `socket_disconnecting:${reason}`);
|
||||
}
|
||||
});
|
||||
|
||||
socket.on("disconnect", (reason, description) => {
|
||||
socketLogger.debug("Socket disconnected", { ...getSocketMetadata(), reason, description });
|
||||
});
|
||||
|
||||
socket.on("error", (error) => {
|
||||
socketLogger.error("Socket error", {
|
||||
...getSocketMetadata(),
|
||||
error: {
|
||||
name: error.name,
|
||||
message: error.message,
|
||||
stack: error.stack,
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
socket.on("run:start", async (message) => {
|
||||
const log = socketLogger.child({
|
||||
eventName: "run:start",
|
||||
...getSocketMetadata(),
|
||||
...message,
|
||||
});
|
||||
|
||||
log.debug("Handling run:start");
|
||||
|
||||
try {
|
||||
runConnected(message.run.friendlyId);
|
||||
} catch (error) {
|
||||
log.error("run:start error", { error });
|
||||
}
|
||||
});
|
||||
|
||||
socket.on("run:stop", async (message) => {
|
||||
const log = socketLogger.child({
|
||||
eventName: "run:stop",
|
||||
...getSocketMetadata(),
|
||||
...message,
|
||||
});
|
||||
|
||||
log.debug("Handling run:stop");
|
||||
|
||||
try {
|
||||
runDisconnected(message.run.friendlyId, "run_stop_message");
|
||||
// Don't delete trace context here - run:stop fires after each snapshot/shutdown
|
||||
// but the run may be restored on a new VM and snapshot again. Trace context is
|
||||
// re-populated on dequeue, and entries are small (4 strings per run).
|
||||
} catch (error) {
|
||||
log.error("run:stop error", { error });
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
return websocketServer;
|
||||
}
|
||||
|
||||
notifyRun({ run }: { run: { friendlyId: string } }) {
|
||||
try {
|
||||
const runSocket = this.runSockets.get(run.friendlyId);
|
||||
|
||||
if (!runSocket) {
|
||||
this.logger.debug("notifyRun: Run socket not found", { run });
|
||||
|
||||
this.workerClient.sendDebugLog(run.friendlyId, {
|
||||
time: new Date(),
|
||||
message: "run:notify socket not found on supervisor",
|
||||
});
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
runSocket.emit("run:notify", { version: "1", run });
|
||||
this.logger.debug("run:notify sent", { run });
|
||||
|
||||
this.workerClient.sendDebugLog(run.friendlyId, {
|
||||
time: new Date(),
|
||||
message: "run:notify supervisor -> runner",
|
||||
});
|
||||
} catch (error) {
|
||||
this.logger.error("Error in notifyRun", { run, error });
|
||||
|
||||
this.workerClient.sendDebugLog(run.friendlyId, {
|
||||
time: new Date(),
|
||||
message: "run:notify error on supervisor",
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
registerRunTraceContext(runFriendlyId: string, ctx: RunTraceContext) {
|
||||
this.snapshotService?.registerTraceContext(runFriendlyId, ctx);
|
||||
}
|
||||
|
||||
async start() {
|
||||
await this.httpServer.start();
|
||||
}
|
||||
|
||||
async stop() {
|
||||
this.snapshotService?.stop();
|
||||
await this.httpServer.stop();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user