chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,27 @@
|
||||
import { WORKER_HEADERS } from "@trigger.dev/core/v3/workers";
|
||||
|
||||
// Secret-bearing headers to drop before logging request headers.
|
||||
// Dependency-free so the redaction is unit-tested directly.
|
||||
export const SENSITIVE_WORKER_HEADERS = new Set([
|
||||
"authorization",
|
||||
"cookie",
|
||||
WORKER_HEADERS.MANAGED_SECRET.toLowerCase(),
|
||||
]);
|
||||
|
||||
/**
|
||||
* Copy `headers` into a plain object, dropping any header whose (lower-cased)
|
||||
* name is in `denylist`. Used before logging request headers.
|
||||
*/
|
||||
export function sanitizeWorkerHeaders(
|
||||
headers: Headers,
|
||||
denylist: ReadonlySet<string> = SENSITIVE_WORKER_HEADERS
|
||||
): Partial<Record<string, string>> {
|
||||
const skip = new Set(Array.from(denylist, (h) => h.toLowerCase()));
|
||||
const sanitized: Partial<Record<string, string>> = {};
|
||||
for (const [key, value] of headers.entries()) {
|
||||
if (!skip.has(key.toLowerCase())) {
|
||||
sanitized[key] = value;
|
||||
}
|
||||
}
|
||||
return sanitized;
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
import { WorkerInstanceGroupType } from "@trigger.dev/database";
|
||||
|
||||
/**
|
||||
* Whether a worker group may be used by the calling project.
|
||||
*
|
||||
* MANAGED groups are shared across projects. UNMANAGED groups are per-project
|
||||
* (masterQueue is `${projectId}-${name}`), so a project may only use an
|
||||
* UNMANAGED group whose `projectId` matches it. Dependency-free so it can be
|
||||
* unit-tested directly.
|
||||
*/
|
||||
export function isWorkerGroupAllowedForProject(
|
||||
workerGroup: { type: WorkerInstanceGroupType; projectId: string | null },
|
||||
projectId: string
|
||||
): boolean {
|
||||
if (workerGroup.type === WorkerInstanceGroupType.UNMANAGED) {
|
||||
return workerGroup.projectId === projectId;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
@@ -0,0 +1,348 @@
|
||||
import type { WorkerInstanceGroup, WorkloadType } from "@trigger.dev/database";
|
||||
import { WorkerInstanceGroupType } from "@trigger.dev/database";
|
||||
import { WithRunEngine } from "../baseService.server";
|
||||
import { isWorkerGroupAllowedForProject } from "./workerGroupAccess";
|
||||
import { WorkerGroupTokenService } from "./workerGroupTokenService.server";
|
||||
import { logger } from "~/services/logger.server";
|
||||
import { FEATURE_FLAG } from "~/v3/featureFlags";
|
||||
import { makeFlag, makeSetFlag } from "~/v3/featureFlags.server";
|
||||
import { isComputeRegionAccessible, resolveComputeAccess } from "~/v3/regionAccess.server";
|
||||
|
||||
export class WorkerGroupService extends WithRunEngine {
|
||||
private readonly defaultNamePrefix = "worker_group";
|
||||
|
||||
async createWorkerGroup({
|
||||
projectId,
|
||||
organizationId,
|
||||
name,
|
||||
description,
|
||||
type,
|
||||
hidden,
|
||||
workloadType,
|
||||
cloudProvider,
|
||||
location,
|
||||
staticIPs,
|
||||
enableFastPath,
|
||||
}: {
|
||||
projectId?: string;
|
||||
organizationId?: string;
|
||||
name?: string;
|
||||
description?: string;
|
||||
type?: WorkerInstanceGroupType;
|
||||
hidden?: boolean;
|
||||
workloadType?: WorkloadType;
|
||||
cloudProvider?: string;
|
||||
location?: string;
|
||||
staticIPs?: string;
|
||||
enableFastPath?: boolean;
|
||||
}) {
|
||||
if (!name) {
|
||||
name = await this.generateWorkerName({ projectId });
|
||||
}
|
||||
|
||||
const tokenService = new WorkerGroupTokenService({
|
||||
prisma: this._prisma,
|
||||
engine: this._engine,
|
||||
});
|
||||
const token = await tokenService.createToken();
|
||||
|
||||
const resolvedType =
|
||||
type ?? (projectId ? WorkerInstanceGroupType.UNMANAGED : WorkerInstanceGroupType.MANAGED);
|
||||
|
||||
const workerGroup = await this._prisma.workerInstanceGroup.create({
|
||||
data: {
|
||||
projectId,
|
||||
organizationId,
|
||||
type: resolvedType,
|
||||
masterQueue: this.generateMasterQueueName({ projectId, name }),
|
||||
tokenId: token.id,
|
||||
description,
|
||||
name,
|
||||
hidden,
|
||||
workloadType,
|
||||
cloudProvider,
|
||||
location,
|
||||
staticIPs,
|
||||
enableFastPath,
|
||||
},
|
||||
});
|
||||
|
||||
if (workerGroup.type === WorkerInstanceGroupType.MANAGED) {
|
||||
const _managedCount = await this._prisma.workerInstanceGroup.count({
|
||||
where: {
|
||||
type: WorkerInstanceGroupType.MANAGED,
|
||||
},
|
||||
});
|
||||
|
||||
const getFlag = makeFlag(this._prisma);
|
||||
const defaultWorkerInstanceGroupId = await getFlag({
|
||||
key: FEATURE_FLAG.defaultWorkerInstanceGroupId,
|
||||
});
|
||||
|
||||
// If there's no global default yet we should set it to the new worker group
|
||||
if (!defaultWorkerInstanceGroupId) {
|
||||
const setFlag = makeSetFlag(this._prisma);
|
||||
await setFlag({
|
||||
key: FEATURE_FLAG.defaultWorkerInstanceGroupId,
|
||||
value: workerGroup.id,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
workerGroup,
|
||||
token,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
This updates a single worker group.
|
||||
The name should never be updated. This would mean changing the masterQueue name which can have unexpected consequences.
|
||||
*/
|
||||
async updateWorkerGroup({
|
||||
projectId,
|
||||
workerGroupId,
|
||||
description,
|
||||
}: {
|
||||
projectId: string;
|
||||
workerGroupId: string;
|
||||
description?: string;
|
||||
}) {
|
||||
const workerGroup = await this._prisma.workerInstanceGroup.findUnique({
|
||||
where: {
|
||||
id: workerGroupId,
|
||||
projectId,
|
||||
},
|
||||
});
|
||||
|
||||
if (!workerGroup) {
|
||||
logger.error("[WorkerGroupService] No worker group found for update", {
|
||||
workerGroupId,
|
||||
description,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
await this._prisma.workerInstanceGroup.update({
|
||||
where: {
|
||||
id: workerGroup.id,
|
||||
},
|
||||
data: {
|
||||
description,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
This lists worker groups.
|
||||
Without a project ID, only shared worker groups will be returned.
|
||||
With a project ID, in addition to all shared worker groups, ones associated with the project will also be returned.
|
||||
*/
|
||||
async listWorkerGroups({ projectId, listHidden }: { projectId?: string; listHidden?: boolean }) {
|
||||
const workerGroups = await this._prisma.workerInstanceGroup.findMany({
|
||||
where: {
|
||||
OR: [
|
||||
{
|
||||
type: WorkerInstanceGroupType.MANAGED,
|
||||
},
|
||||
{
|
||||
projectId,
|
||||
},
|
||||
],
|
||||
AND: listHidden ? [] : [{ hidden: false }],
|
||||
},
|
||||
});
|
||||
|
||||
return workerGroups;
|
||||
}
|
||||
|
||||
async deleteWorkerGroup({
|
||||
projectId,
|
||||
workerGroupId,
|
||||
}: {
|
||||
projectId: string;
|
||||
workerGroupId: string;
|
||||
}) {
|
||||
const workerGroup = await this._prisma.workerInstanceGroup.findUnique({
|
||||
where: {
|
||||
id: workerGroupId,
|
||||
},
|
||||
});
|
||||
|
||||
if (!workerGroup) {
|
||||
logger.error("[WorkerGroupService] WorkerGroup not found for deletion", {
|
||||
workerGroupId,
|
||||
projectId,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
if (workerGroup.projectId !== projectId) {
|
||||
logger.error("[WorkerGroupService] WorkerGroup does not belong to project", {
|
||||
workerGroupId,
|
||||
projectId,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
await this._prisma.workerInstanceGroup.delete({
|
||||
where: {
|
||||
id: workerGroupId,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
async getGlobalDefaultWorkerGroup() {
|
||||
const flags = makeFlag(this._prisma);
|
||||
|
||||
const defaultWorkerInstanceGroupId = await flags({
|
||||
key: FEATURE_FLAG.defaultWorkerInstanceGroupId,
|
||||
});
|
||||
|
||||
if (!defaultWorkerInstanceGroupId) {
|
||||
logger.error("[WorkerGroupService] Default worker group not found in feature flags");
|
||||
return;
|
||||
}
|
||||
|
||||
const workerGroup = await this._prisma.workerInstanceGroup.findUnique({
|
||||
where: {
|
||||
id: defaultWorkerInstanceGroupId,
|
||||
},
|
||||
});
|
||||
|
||||
if (!workerGroup) {
|
||||
logger.error("[WorkerGroupService] Default worker group not found", {
|
||||
defaultWorkerInstanceGroupId,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
return workerGroup;
|
||||
}
|
||||
|
||||
async getDefaultWorkerGroupForProject({
|
||||
projectId,
|
||||
regionOverride,
|
||||
}: {
|
||||
projectId: string;
|
||||
regionOverride?: string;
|
||||
}): Promise<WorkerInstanceGroup | undefined> {
|
||||
const project = await this._prisma.project.findFirst({
|
||||
where: {
|
||||
id: projectId,
|
||||
},
|
||||
include: {
|
||||
defaultWorkerGroup: true,
|
||||
organization: { select: { featureFlags: true } },
|
||||
},
|
||||
});
|
||||
|
||||
if (!project) {
|
||||
throw new Error("Project not found.");
|
||||
}
|
||||
|
||||
// If they've specified a region, we need to check they have access to it
|
||||
if (regionOverride) {
|
||||
const workerGroup = await this._prisma.workerInstanceGroup.findFirst({
|
||||
where: {
|
||||
masterQueue: regionOverride,
|
||||
},
|
||||
});
|
||||
|
||||
if (!workerGroup) {
|
||||
throw new Error(`The region you specified doesn't exist ("${regionOverride}").`);
|
||||
}
|
||||
|
||||
// The masterQueue-only lookup above can resolve another project's
|
||||
// UNMANAGED group, so reject groups not usable by this project
|
||||
// (see isWorkerGroupAllowedForProject).
|
||||
if (!isWorkerGroupAllowedForProject(workerGroup, project.id)) {
|
||||
throw new Error(`The region you specified isn't available to you ("${regionOverride}").`);
|
||||
}
|
||||
|
||||
// If they're restricted, check they have access
|
||||
if (project.allowedWorkerQueues.length > 0) {
|
||||
if (project.allowedWorkerQueues.includes(workerGroup.masterQueue)) {
|
||||
return workerGroup;
|
||||
}
|
||||
|
||||
throw new Error(
|
||||
`You don't have access to this region ("${regionOverride}"). You can use the following regions: ${project.allowedWorkerQueues.join(
|
||||
", "
|
||||
)}.`
|
||||
);
|
||||
}
|
||||
|
||||
if (workerGroup.hidden) {
|
||||
throw new Error(`The region you specified isn't available to you ("${regionOverride}").`);
|
||||
}
|
||||
|
||||
if (workerGroup.workloadType === "MICROVM") {
|
||||
const hasComputeAccess = await resolveComputeAccess(
|
||||
this._prisma,
|
||||
project.organization.featureFlags
|
||||
);
|
||||
|
||||
if (!isComputeRegionAccessible(workerGroup, hasComputeAccess)) {
|
||||
throw new Error(`The region you specified isn't available to you ("${regionOverride}").`);
|
||||
}
|
||||
}
|
||||
|
||||
return workerGroup;
|
||||
}
|
||||
|
||||
if (project.defaultWorkerGroup) {
|
||||
return project.defaultWorkerGroup;
|
||||
}
|
||||
|
||||
return await this.getGlobalDefaultWorkerGroup();
|
||||
}
|
||||
|
||||
async setDefaultWorkerGroupForProject({
|
||||
projectId,
|
||||
workerGroupId,
|
||||
}: {
|
||||
projectId: string;
|
||||
workerGroupId: string;
|
||||
}) {
|
||||
const workerGroup = await this._prisma.workerInstanceGroup.findUnique({
|
||||
where: {
|
||||
id: workerGroupId,
|
||||
},
|
||||
});
|
||||
|
||||
if (!workerGroup) {
|
||||
logger.error("[WorkerGroupService] WorkerGroup not found", {
|
||||
workerGroupId,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
await this._prisma.project.update({
|
||||
where: {
|
||||
id: projectId,
|
||||
},
|
||||
data: {
|
||||
defaultWorkerGroupId: workerGroupId,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
private async generateWorkerName({ projectId }: { projectId?: string }) {
|
||||
const workerGroups = await this._prisma.workerInstanceGroup.count({
|
||||
where: {
|
||||
projectId: projectId ?? null,
|
||||
},
|
||||
});
|
||||
|
||||
return `${this.defaultNamePrefix}_${workerGroups + 1}`;
|
||||
}
|
||||
|
||||
private generateMasterQueueName({ projectId, name }: { projectId?: string; name: string }) {
|
||||
if (!projectId) {
|
||||
return name;
|
||||
}
|
||||
|
||||
return `${projectId}-${name}`;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,612 @@
|
||||
import {
|
||||
createCache,
|
||||
createLRUMemoryStore,
|
||||
DefaultStatefulContext,
|
||||
Namespace,
|
||||
} from "@internal/cache";
|
||||
import type {
|
||||
CheckpointInput,
|
||||
CompleteRunAttemptResult,
|
||||
DequeuedMessage,
|
||||
ExecutionResult,
|
||||
MachinePreset,
|
||||
StartRunAttemptResult,
|
||||
TaskRunExecutionResult,
|
||||
} from "@trigger.dev/core/v3";
|
||||
import { SemanticInternalAttributes } from "@trigger.dev/core/v3";
|
||||
import { fromFriendlyId } from "@trigger.dev/core/v3/isomorphic";
|
||||
import { WORKER_HEADERS, type WorkerQueueClass } from "@trigger.dev/core/v3/workers";
|
||||
import type { RuntimeEnvironment, WorkerInstanceGroup } from "@trigger.dev/database";
|
||||
import { Prisma, WorkerInstanceGroupType } from "@trigger.dev/database";
|
||||
import { SENSITIVE_WORKER_HEADERS, sanitizeWorkerHeaders } from "./sanitizeWorkerHeaders";
|
||||
import { createHash, timingSafeEqual } from "crypto";
|
||||
import { customAlphabet } from "nanoid";
|
||||
import { z } from "zod";
|
||||
import { env } from "~/env.server";
|
||||
import {
|
||||
isWorkerQueueDequeueDisabled,
|
||||
recordBlockedDequeue,
|
||||
} from "~/runEngine/concerns/dequeueGate.server";
|
||||
import { workerQueueForClass } from "~/runEngine/concerns/workerQueueSplit.server";
|
||||
import { generateJWTTokenForEnvironment } from "~/services/apiAuth.server";
|
||||
import { logger } from "~/services/logger.server";
|
||||
import { defaultMachine } from "~/services/platform.v3.server";
|
||||
import { singleton } from "~/utils/singleton";
|
||||
import { resolveVariablesForEnvironment } from "~/v3/environmentVariables/environmentVariablesRepository.server";
|
||||
import { machinePresetFromName } from "~/v3/machinePresets.server";
|
||||
import type { WithRunEngineOptions } from "../baseService.server";
|
||||
import { WithRunEngine } from "../baseService.server";
|
||||
|
||||
const authenticatedWorkerInstanceCache = singleton(
|
||||
"authenticatedWorkerInstanceCache",
|
||||
createAuthenticatedWorkerInstanceCache
|
||||
);
|
||||
|
||||
function createAuthenticatedWorkerInstanceCache() {
|
||||
return createCache({
|
||||
authenticatedWorkerInstance: new Namespace<AuthenticatedWorkerInstance>(
|
||||
new DefaultStatefulContext(),
|
||||
{
|
||||
stores: [createLRUMemoryStore(1000)],
|
||||
fresh: 60_000 * 10, // 10 minutes
|
||||
stale: 60_000 * 11, // 11 minutes
|
||||
}
|
||||
),
|
||||
});
|
||||
}
|
||||
|
||||
export class WorkerGroupTokenService extends WithRunEngine {
|
||||
private readonly tokenPrefix = "tr_wgt_";
|
||||
private readonly tokenLength = 40;
|
||||
private readonly tokenChars = "1234567890abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
|
||||
private readonly tokenGenerator = customAlphabet(this.tokenChars, this.tokenLength);
|
||||
|
||||
async createToken() {
|
||||
const rawToken = await this.generateToken();
|
||||
|
||||
const workerGroupToken = await this._prisma.workerGroupToken.create({
|
||||
data: {
|
||||
tokenHash: rawToken.hash,
|
||||
},
|
||||
});
|
||||
|
||||
return {
|
||||
id: workerGroupToken.id,
|
||||
tokenHash: workerGroupToken.tokenHash,
|
||||
plaintext: rawToken.plaintext,
|
||||
};
|
||||
}
|
||||
|
||||
async findWorkerGroup({ token }: { token: string }) {
|
||||
const tokenHash = await this.hashToken(token);
|
||||
|
||||
const workerGroup = await this._prisma.workerInstanceGroup.findFirst({
|
||||
where: {
|
||||
token: {
|
||||
tokenHash,
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
if (!workerGroup) {
|
||||
logger.warn("[WorkerGroupTokenService] No matching worker group found", { token });
|
||||
return null;
|
||||
}
|
||||
|
||||
return workerGroup;
|
||||
}
|
||||
|
||||
async rotateToken({ workerGroupId }: { workerGroupId: string }) {
|
||||
const workerGroup = await this._prisma.workerInstanceGroup.findFirst({
|
||||
where: {
|
||||
id: workerGroupId,
|
||||
},
|
||||
});
|
||||
|
||||
if (!workerGroup) {
|
||||
logger.error("[WorkerGroupTokenService] WorkerGroup not found", { workerGroupId });
|
||||
return;
|
||||
}
|
||||
|
||||
const rawToken = await this.generateToken();
|
||||
|
||||
const workerGroupToken = await this._prisma.workerGroupToken.update({
|
||||
where: {
|
||||
id: workerGroup.tokenId,
|
||||
},
|
||||
data: {
|
||||
tokenHash: rawToken.hash,
|
||||
},
|
||||
});
|
||||
|
||||
if (!workerGroupToken) {
|
||||
logger.error("[WorkerGroupTokenService] WorkerGroupToken not found", { workerGroupId });
|
||||
return;
|
||||
}
|
||||
|
||||
return {
|
||||
id: workerGroupToken.id,
|
||||
tokenHash: workerGroupToken.tokenHash,
|
||||
plaintext: rawToken.plaintext,
|
||||
};
|
||||
}
|
||||
|
||||
private async hashToken(token: string) {
|
||||
return createHash("sha256").update(token).digest("hex");
|
||||
}
|
||||
|
||||
private async generateToken() {
|
||||
const plaintext = `${this.tokenPrefix}${this.tokenGenerator()}`;
|
||||
const hash = await this.hashToken(plaintext);
|
||||
|
||||
return {
|
||||
plaintext,
|
||||
hash,
|
||||
};
|
||||
}
|
||||
|
||||
async authenticate(request: Request): Promise<AuthenticatedWorkerInstance | undefined> {
|
||||
const token = request.headers.get("Authorization")?.replace("Bearer ", "").trim();
|
||||
|
||||
if (!token) {
|
||||
logger.error("[WorkerGroupTokenService] Token not found in request", {
|
||||
headers: this.sanitizeHeaders(request),
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
if (!token.startsWith(this.tokenPrefix)) {
|
||||
logger.error("[WorkerGroupTokenService] Token does not start with expected prefix", {
|
||||
token,
|
||||
prefix: this.tokenPrefix,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
const instanceName = request.headers.get(WORKER_HEADERS.INSTANCE_NAME);
|
||||
|
||||
if (!instanceName) {
|
||||
logger.error("[WorkerGroupTokenService] Instance name not found in request", {
|
||||
headers: this.sanitizeHeaders(request),
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
const managedWorkerSecret = request.headers.get(WORKER_HEADERS.MANAGED_SECRET);
|
||||
|
||||
if (!managedWorkerSecret) {
|
||||
logger.error("[WorkerGroupTokenService] Managed secret not found in request", {
|
||||
headers: this.sanitizeHeaders(request),
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
const encoder = new TextEncoder();
|
||||
|
||||
const a = encoder.encode(managedWorkerSecret);
|
||||
const b = encoder.encode(env.MANAGED_WORKER_SECRET);
|
||||
|
||||
if (a.byteLength !== b.byteLength) {
|
||||
logger.error("[WorkerGroupTokenService] Managed secret length mismatch", {
|
||||
headers: this.sanitizeHeaders(request),
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
if (!timingSafeEqual(a, b)) {
|
||||
logger.error("[WorkerGroupTokenService] Managed secret mismatch", {
|
||||
headers: this.sanitizeHeaders(request),
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
const cacheKey = ["worker-group-token", token, instanceName];
|
||||
|
||||
const result = await authenticatedWorkerInstanceCache.authenticatedWorkerInstance.swr(
|
||||
cacheKey.join("-"),
|
||||
async () => {
|
||||
const workerGroup = await this.findWorkerGroup({ token });
|
||||
|
||||
if (!workerGroup) {
|
||||
logger.warn("[WorkerGroupTokenService] Worker group not found", { token });
|
||||
return;
|
||||
}
|
||||
|
||||
const workerInstance = await this.getOrCreateWorkerInstance({
|
||||
workerGroup,
|
||||
instanceName,
|
||||
});
|
||||
|
||||
if (!workerInstance) {
|
||||
logger.error("[WorkerGroupTokenService] Unable to get or create worker instance", {
|
||||
workerGroup,
|
||||
instanceName,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
return new AuthenticatedWorkerInstance({
|
||||
prisma: this._prisma,
|
||||
engine: this._engine,
|
||||
type: WorkerInstanceGroupType.MANAGED,
|
||||
name: workerGroup.name,
|
||||
workerGroupId: workerGroup.id,
|
||||
workerInstanceId: workerInstance.id,
|
||||
masterQueue: workerGroup.masterQueue,
|
||||
});
|
||||
}
|
||||
);
|
||||
|
||||
if (result.err) {
|
||||
logger.error("[WorkerGroupTokenService] Failed to authenticate worker instance", {
|
||||
error: result.err,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
return result.val;
|
||||
}
|
||||
|
||||
private async getOrCreateWorkerInstance({
|
||||
workerGroup,
|
||||
instanceName,
|
||||
}: {
|
||||
workerGroup: WorkerInstanceGroup;
|
||||
instanceName: string;
|
||||
}) {
|
||||
const resourceIdentifier = instanceName;
|
||||
|
||||
const workerInstance = await this._prisma.workerInstance.findFirst({
|
||||
where: {
|
||||
workerGroupId: workerGroup.id,
|
||||
resourceIdentifier,
|
||||
},
|
||||
include: {
|
||||
deployment: true,
|
||||
environment: true,
|
||||
},
|
||||
});
|
||||
|
||||
if (workerInstance) {
|
||||
return workerInstance;
|
||||
}
|
||||
|
||||
try {
|
||||
const newWorkerInstance = await this._prisma.workerInstance.create({
|
||||
data: {
|
||||
workerGroupId: workerGroup.id,
|
||||
name: instanceName,
|
||||
resourceIdentifier,
|
||||
},
|
||||
include: {
|
||||
// This will always be empty for shared worker instances, but required for types
|
||||
deployment: true,
|
||||
environment: true,
|
||||
},
|
||||
});
|
||||
|
||||
return newWorkerInstance;
|
||||
} catch (error) {
|
||||
// Gracefully handle race conditions when connecting for the first time
|
||||
if (error instanceof Prisma.PrismaClientKnownRequestError) {
|
||||
// Unique constraint violation
|
||||
if (error.code === "P2002") {
|
||||
try {
|
||||
const existingWorkerInstance = await this._prisma.workerInstance.findFirst({
|
||||
where: {
|
||||
workerGroupId: workerGroup.id,
|
||||
resourceIdentifier,
|
||||
},
|
||||
include: {
|
||||
deployment: true,
|
||||
environment: true,
|
||||
},
|
||||
});
|
||||
|
||||
return existingWorkerInstance;
|
||||
} catch (_error) {
|
||||
logger.error("[WorkerGroupTokenService] Failed to find worker instance", {
|
||||
workerGroup,
|
||||
workerInstance,
|
||||
});
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Strip sensitive headers before logging request headers — see
|
||||
// `sanitizeWorkerHeaders`.
|
||||
private sanitizeHeaders(request: Request, denylist = SENSITIVE_WORKER_HEADERS) {
|
||||
return sanitizeWorkerHeaders(request.headers, denylist);
|
||||
}
|
||||
}
|
||||
|
||||
export const WorkerInstanceEnv = z.enum(["dev", "staging", "prod"]).default("prod");
|
||||
export type WorkerInstanceEnv = z.infer<typeof WorkerInstanceEnv>;
|
||||
|
||||
export type AuthenticatedWorkerInstanceOptions = WithRunEngineOptions<{
|
||||
type: WorkerInstanceGroupType;
|
||||
name: string;
|
||||
workerGroupId: string;
|
||||
workerInstanceId: string;
|
||||
masterQueue: string;
|
||||
}>;
|
||||
|
||||
export class AuthenticatedWorkerInstance extends WithRunEngine {
|
||||
readonly type: WorkerInstanceGroupType;
|
||||
readonly name: string;
|
||||
readonly workerGroupId: string;
|
||||
readonly workerInstanceId: string;
|
||||
readonly masterQueue: string;
|
||||
|
||||
// FIXME: Required for unmanaged workers
|
||||
readonly isLatestDeployment = true;
|
||||
|
||||
constructor(opts: AuthenticatedWorkerInstanceOptions) {
|
||||
super({ prisma: opts.prisma, engine: opts.engine });
|
||||
|
||||
this.type = opts.type;
|
||||
this.name = opts.name;
|
||||
this.workerGroupId = opts.workerGroupId;
|
||||
this.workerInstanceId = opts.workerInstanceId;
|
||||
this.masterQueue = opts.masterQueue;
|
||||
}
|
||||
|
||||
async connect(metadata: Record<string, any>): Promise<void> {
|
||||
await this._prisma.workerInstance.update({
|
||||
where: {
|
||||
id: this.workerInstanceId,
|
||||
},
|
||||
data: {
|
||||
metadata,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
async dequeue({
|
||||
runnerId,
|
||||
queueClass,
|
||||
}: {
|
||||
runnerId?: string;
|
||||
queueClass?: WorkerQueueClass;
|
||||
}): Promise<DequeuedMessage[]> {
|
||||
const workerQueue = workerQueueForClass(this.masterQueue, queueClass);
|
||||
|
||||
if (isWorkerQueueDequeueDisabled(workerQueue)) {
|
||||
recordBlockedDequeue(workerQueue);
|
||||
return [];
|
||||
}
|
||||
|
||||
return await this._engine.dequeueFromWorkerQueue({
|
||||
consumerId: this.workerInstanceId,
|
||||
workerQueue,
|
||||
workerId: this.workerInstanceId,
|
||||
runnerId,
|
||||
});
|
||||
}
|
||||
|
||||
async heartbeatWorkerInstance() {
|
||||
await this._prisma.workerInstance.update({
|
||||
where: {
|
||||
id: this.workerInstanceId,
|
||||
},
|
||||
data: {
|
||||
lastHeartbeatAt: new Date(),
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
async heartbeatRun({
|
||||
runFriendlyId,
|
||||
snapshotFriendlyId,
|
||||
runnerId,
|
||||
}: {
|
||||
runFriendlyId: string;
|
||||
snapshotFriendlyId: string;
|
||||
runnerId?: string;
|
||||
}): Promise<ExecutionResult> {
|
||||
return await this._engine.heartbeatRun({
|
||||
runId: fromFriendlyId(runFriendlyId),
|
||||
snapshotId: fromFriendlyId(snapshotFriendlyId),
|
||||
workerId: this.workerInstanceId,
|
||||
runnerId,
|
||||
});
|
||||
}
|
||||
|
||||
async startRunAttempt({
|
||||
runFriendlyId,
|
||||
snapshotFriendlyId,
|
||||
isWarmStart,
|
||||
runnerId,
|
||||
}: {
|
||||
runFriendlyId: string;
|
||||
snapshotFriendlyId: string;
|
||||
isWarmStart?: boolean;
|
||||
runnerId?: string;
|
||||
}): Promise<
|
||||
StartRunAttemptResult & {
|
||||
envVars: Record<string, string>;
|
||||
}
|
||||
> {
|
||||
const engineResult = await this._engine.startRunAttempt({
|
||||
runId: fromFriendlyId(runFriendlyId),
|
||||
snapshotId: fromFriendlyId(snapshotFriendlyId),
|
||||
isWarmStart,
|
||||
workerId: this.workerInstanceId,
|
||||
runnerId,
|
||||
});
|
||||
|
||||
const defaultMachinePreset = machinePresetFromName(defaultMachine);
|
||||
|
||||
const environment = await this._prisma.runtimeEnvironment.findFirst({
|
||||
where: {
|
||||
id: engineResult.execution.environment.id,
|
||||
},
|
||||
include: {
|
||||
parentEnvironment: true,
|
||||
},
|
||||
});
|
||||
|
||||
const envVars = environment
|
||||
? await this.getEnvVars(
|
||||
environment,
|
||||
engineResult.run.id,
|
||||
engineResult.execution.machine ?? defaultMachinePreset,
|
||||
environment.parentEnvironment ?? undefined,
|
||||
engineResult.run.taskEventStore ?? undefined
|
||||
)
|
||||
: {};
|
||||
|
||||
return {
|
||||
...engineResult,
|
||||
envVars,
|
||||
};
|
||||
}
|
||||
|
||||
async completeRunAttempt({
|
||||
runFriendlyId,
|
||||
snapshotFriendlyId,
|
||||
completion,
|
||||
runnerId,
|
||||
}: {
|
||||
runFriendlyId: string;
|
||||
snapshotFriendlyId: string;
|
||||
completion: TaskRunExecutionResult;
|
||||
runnerId?: string;
|
||||
}): Promise<CompleteRunAttemptResult> {
|
||||
return await this._engine.completeRunAttempt({
|
||||
runId: fromFriendlyId(runFriendlyId),
|
||||
snapshotId: fromFriendlyId(snapshotFriendlyId),
|
||||
completion,
|
||||
workerId: this.workerInstanceId,
|
||||
runnerId,
|
||||
});
|
||||
}
|
||||
|
||||
async getLatestSnapshot({ runFriendlyId }: { runFriendlyId: string }) {
|
||||
return await this._engine.getRunExecutionData({
|
||||
runId: fromFriendlyId(runFriendlyId),
|
||||
});
|
||||
}
|
||||
|
||||
async createCheckpoint({
|
||||
runFriendlyId,
|
||||
snapshotFriendlyId,
|
||||
checkpoint,
|
||||
runnerId,
|
||||
}: {
|
||||
runFriendlyId: string;
|
||||
snapshotFriendlyId: string;
|
||||
checkpoint: CheckpointInput;
|
||||
runnerId?: string;
|
||||
}) {
|
||||
return await this._engine.createCheckpoint({
|
||||
runId: fromFriendlyId(runFriendlyId),
|
||||
snapshotId: fromFriendlyId(snapshotFriendlyId),
|
||||
checkpoint,
|
||||
workerId: this.workerInstanceId,
|
||||
runnerId,
|
||||
});
|
||||
}
|
||||
|
||||
async continueRunExecution({
|
||||
runFriendlyId,
|
||||
snapshotFriendlyId,
|
||||
runnerId,
|
||||
}: {
|
||||
runFriendlyId: string;
|
||||
snapshotFriendlyId: string;
|
||||
runnerId?: string;
|
||||
}) {
|
||||
return await this._engine.continueRunExecution({
|
||||
runId: fromFriendlyId(runFriendlyId),
|
||||
snapshotId: fromFriendlyId(snapshotFriendlyId),
|
||||
workerId: this.workerInstanceId,
|
||||
runnerId,
|
||||
});
|
||||
}
|
||||
|
||||
async getSnapshotsSince({
|
||||
runFriendlyId,
|
||||
snapshotId,
|
||||
}: {
|
||||
runFriendlyId: string;
|
||||
snapshotId: string;
|
||||
}) {
|
||||
return await this._engine.getSnapshotsSince({
|
||||
runId: fromFriendlyId(runFriendlyId),
|
||||
snapshotId: fromFriendlyId(snapshotId),
|
||||
});
|
||||
}
|
||||
|
||||
toJSON(): WorkerGroupTokenAuthenticationResponse {
|
||||
return {
|
||||
type: WorkerInstanceGroupType.MANAGED,
|
||||
name: this.name,
|
||||
workerGroupId: this.workerGroupId,
|
||||
workerInstanceId: this.workerInstanceId,
|
||||
masterQueue: this.masterQueue,
|
||||
};
|
||||
}
|
||||
|
||||
private async getEnvVars(
|
||||
environment: RuntimeEnvironment,
|
||||
runId: string,
|
||||
machinePreset: MachinePreset,
|
||||
parentEnvironment?: RuntimeEnvironment,
|
||||
taskEventStore?: string
|
||||
): Promise<Record<string, string>> {
|
||||
const variables = await resolveVariablesForEnvironment(environment, parentEnvironment);
|
||||
|
||||
const jwt = await generateJWTTokenForEnvironment(environment, {
|
||||
run_id: runId,
|
||||
machine_preset: machinePreset.name,
|
||||
});
|
||||
|
||||
variables.push(
|
||||
...[
|
||||
{ key: "TRIGGER_JWT", value: jwt },
|
||||
{ key: "TRIGGER_RUN_ID", value: runId },
|
||||
{ key: "TRIGGER_MACHINE_PRESET", value: machinePreset.name },
|
||||
]
|
||||
);
|
||||
|
||||
if (taskEventStore) {
|
||||
const resourceAttributes = JSON.stringify({
|
||||
[SemanticInternalAttributes.TASK_EVENT_STORE]: taskEventStore,
|
||||
});
|
||||
|
||||
variables.push(
|
||||
...[
|
||||
{ key: "OTEL_RESOURCE_ATTRIBUTES", value: resourceAttributes },
|
||||
{ key: "TRIGGER_OTEL_RESOURCE_ATTRIBUTES", value: resourceAttributes },
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
return variables.reduce((acc: Record<string, string>, curr) => {
|
||||
acc[curr.key] = curr.value;
|
||||
return acc;
|
||||
}, {});
|
||||
}
|
||||
}
|
||||
|
||||
export type WorkerGroupTokenAuthenticationResponse =
|
||||
| {
|
||||
type: typeof WorkerInstanceGroupType.MANAGED;
|
||||
name: string;
|
||||
workerGroupId: string;
|
||||
workerInstanceId: string;
|
||||
masterQueue: string;
|
||||
}
|
||||
| {
|
||||
type: typeof WorkerInstanceGroupType.UNMANAGED;
|
||||
name: string;
|
||||
workerGroupId: string;
|
||||
workerInstanceId: string;
|
||||
masterQueue: string;
|
||||
environmentId: string;
|
||||
deploymentId: string;
|
||||
};
|
||||
Reference in New Issue
Block a user