diff --git a/apps/webapp/app/components/admin/debugRun.tsx b/apps/webapp/app/components/admin/debugRun.tsx new file mode 100644 index 0000000000..894a2a3bc9 --- /dev/null +++ b/apps/webapp/app/components/admin/debugRun.tsx @@ -0,0 +1,340 @@ +import { useIsImpersonating } from "~/hooks/useOrganizations"; +import { useHasAdminAccess } from "~/hooks/useUser"; +import { Button } from "../primitives/Buttons"; +import { Dialog, DialogContent, DialogHeader, DialogTrigger } from "../primitives/Dialog"; +import { Cog6ToothIcon } from "@heroicons/react/20/solid"; +import { type loader } from "~/routes/resources.taskruns.$runParam.debug"; +import { UseDataFunctionReturn, useTypedFetcher } from "remix-typedjson"; +import { useEffect } from "react"; +import { Spinner } from "../primitives/Spinner"; +import * as Property from "~/components/primitives/PropertyTable"; +import { ClipboardField } from "../primitives/ClipboardField"; +import { MarQSShortKeyProducer } from "~/v3/marqs/marqsKeyProducer"; + +export function AdminDebugRun({ friendlyId }: { friendlyId: string }) { + const hasAdminAccess = useHasAdminAccess(); + const isImpersonating = useIsImpersonating(); + + if (!hasAdminAccess && !isImpersonating) { + return null; + } + + return ( + + + + + + + ); +} + +export function DebugRunDialog({ friendlyId }: { friendlyId: string }) { + return ( + + + + ); +} + +function DebugRunContent({ friendlyId }: { friendlyId: string }) { + const fetcher = useTypedFetcher(); + const isLoading = fetcher.state === "loading"; + + useEffect(() => { + fetcher.load(`/resources/taskruns/${friendlyId}/debug`); + }, [friendlyId]); + + return ( + <> + Debugging run + {isLoading ? ( +
+ +
+ ) : fetcher.data ? ( + + ) : ( + <>Failed to get run debug data + )} + + ); +} + +function DebugRunData({ + run, + queueConcurrencyLimit, + queueCurrentConcurrency, + envConcurrencyLimit, + envCurrentConcurrency, + queueReserveConcurrency, + envReserveConcurrency, +}: UseDataFunctionReturn) { + const keys = new MarQSShortKeyProducer("marqs:"); + + const withPrefix = (key: string) => `marqs:${key}`; + + return ( + + + ID + + + + + + Message key + + + + + + GET message + + + + + + Queue key + + + + + + Get queue set + + + + + + Queue current concurrency key + + + + + + + Get queue current concurrency + + + + + + Queue current concurrency + + {queueCurrentConcurrency ?? "0"} + + + + Queue reserve concurrency key + + + + + + + Get queue reserve concurrency + + + + + + Queue reserve concurrency + + {queueReserveConcurrency ?? "0"} + + + + Queue concurrency limit key + + + + + + GET queue concurrency limit + + + + + + Queue concurrency limit + + {queueConcurrencyLimit ?? "Not set"} + + + + Env current concurrency key + + + + + + Get env current concurrency + + + + + + Env current concurrency + + {envCurrentConcurrency ?? "0"} + + + + Env reserve concurrency key + + + + + + Get env reserve concurrency + + + + + + Env reserve concurrency + + {envReserveConcurrency ?? "0"} + + + + Env concurrency limit key + + + + + + GET env concurrency limit + + + + + + Env concurrency limit + + {envConcurrencyLimit ?? "Not set"} + + + + Shared queue key + + + + + + Get shared queue set + + + + + + ); +} diff --git a/apps/webapp/app/components/runs/v3/BatchStatus.tsx b/apps/webapp/app/components/runs/v3/BatchStatus.tsx index b512ac1b3c..c67b1b4016 100644 --- a/apps/webapp/app/components/runs/v3/BatchStatus.tsx +++ b/apps/webapp/app/components/runs/v3/BatchStatus.tsx @@ -1,16 +1,17 @@ -import { CheckCircleIcon } from "@heroicons/react/20/solid"; +import { CheckCircleIcon, XCircleIcon } from "@heroicons/react/20/solid"; import { BatchTaskRunStatus } from "@trigger.dev/database"; import assertNever from "assert-never"; import { Spinner } from "~/components/primitives/Spinner"; import { cn } from "~/utils/cn"; -export const allBatchStatuses = ["PENDING", "COMPLETED"] as const satisfies Readonly< +export const allBatchStatuses = ["PENDING", "COMPLETED", "ABORTED"] as const satisfies Readonly< Array >; const descriptions: Record = { PENDING: "The batch has child runs that have not yet completed.", COMPLETED: "All the batch child runs have finished.", + ABORTED: "The batch was aborted because some child tasks could not be triggered.", }; export function descriptionForBatchStatus(status: BatchTaskRunStatus): string { @@ -50,6 +51,8 @@ export function BatchStatusIcon({ return ; case "COMPLETED": return ; + case "ABORTED": + return ; default: { assertNever(status); } @@ -62,6 +65,8 @@ export function batchStatusColor(status: BatchTaskRunStatus): string { return "text-pending"; case "COMPLETED": return "text-success"; + case "ABORTED": + return "text-error"; default: { assertNever(status); } @@ -74,6 +79,8 @@ export function batchStatusTitle(status: BatchTaskRunStatus): string { return "In progress"; case "COMPLETED": return "Completed"; + case "ABORTED": + return "Aborted"; default: { assertNever(status); } diff --git a/apps/webapp/app/env.server.ts b/apps/webapp/app/env.server.ts index 310732470e..c644503f42 100644 --- a/apps/webapp/app/env.server.ts +++ b/apps/webapp/app/env.server.ts @@ -357,7 +357,7 @@ const EnvironmentSchema = z.object({ MARQS_AVAILABLE_CAPACITY_BIAS: z.coerce.number().default(0.3), MARQS_QUEUE_AGE_RANDOMIZATION_BIAS: z.coerce.number().default(0.25), MARQS_REUSE_SNAPSHOT_COUNT: z.coerce.number().int().default(0), - MARQS_MAXIMUM_ORG_COUNT: z.coerce.number().int().optional(), + MARQS_MAXIMUM_ENV_COUNT: z.coerce.number().int().optional(), PROD_TASK_HEARTBEAT_INTERVAL_MS: z.coerce.number().int().optional(), diff --git a/apps/webapp/app/presenters/v3/BatchListPresenter.server.ts b/apps/webapp/app/presenters/v3/BatchListPresenter.server.ts index 2317f68a14..ef7d3f2dd4 100644 --- a/apps/webapp/app/presenters/v3/BatchListPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/BatchListPresenter.server.ts @@ -190,7 +190,7 @@ WHERE throw new Error(`Environment not found for Batch ${batch.id}`); } - const hasFinished = batch.status === "COMPLETED"; + const hasFinished = batch.status !== "PENDING"; return { id: batch.id, diff --git a/apps/webapp/app/presenters/v3/SpanPresenter.server.ts b/apps/webapp/app/presenters/v3/SpanPresenter.server.ts index 010efcba12..0f6bb9de4b 100644 --- a/apps/webapp/app/presenters/v3/SpanPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/SpanPresenter.server.ts @@ -118,6 +118,9 @@ export class SpanPresenter extends BasePresenter { metadata: true, metadataType: true, maxAttempts: true, + output: true, + outputType: true, + error: true, project: { include: { organization: true, @@ -162,31 +165,13 @@ export class SpanPresenter extends BasePresenter { const isFinished = isFinalRunStatus(run.status); - const finishedAttempt = isFinished - ? await this._replica.taskRunAttempt.findFirst({ - select: { - output: true, - outputType: true, - error: true, - }, - where: { - status: { in: FINAL_ATTEMPT_STATUSES }, - taskRunId: run.id, - }, - orderBy: { - createdAt: "desc", - }, - }) - : null; - - const output = - finishedAttempt === null - ? undefined - : finishedAttempt.outputType === "application/store" - ? `/resources/packets/${run.runtimeEnvironment.id}/${finishedAttempt.output}` - : typeof finishedAttempt.output !== "undefined" && finishedAttempt.output !== null - ? await prettyPrintPacket(finishedAttempt.output, finishedAttempt.outputType ?? undefined) - : undefined; + const output = !isFinished + ? undefined + : run.outputType === "application/store" + ? `/resources/packets/${run.runtimeEnvironment.id}/${run.output}` + : typeof run.output !== "undefined" && run.output !== null + ? await prettyPrintPacket(run.output, run.outputType ?? undefined) + : undefined; const payload = run.payloadType === "application/store" @@ -196,14 +181,14 @@ export class SpanPresenter extends BasePresenter { : undefined; let error: TaskRunError | undefined = undefined; - if (finishedAttempt?.error) { - const result = TaskRunError.safeParse(finishedAttempt.error); + if (run?.error) { + const result = TaskRunError.safeParse(run.error); if (result.success) { error = result.data; } else { error = { type: "CUSTOM_ERROR", - raw: JSON.stringify(finishedAttempt.error), + raw: JSON.stringify(run.error), }; } } @@ -300,7 +285,7 @@ export class SpanPresenter extends BasePresenter { payload, payloadType: run.payloadType, output, - outputType: finishedAttempt?.outputType ?? "application/json", + outputType: run?.outputType ?? "application/json", error, relationships: { root: run.rootTaskRun diff --git a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.v3.$projectParam.runs.$runParam/route.tsx b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.v3.$projectParam.runs.$runParam/route.tsx index 6b975a1a14..cf368c8473 100644 --- a/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.v3.$projectParam.runs.$runParam/route.tsx +++ b/apps/webapp/app/routes/_app.orgs.$organizationSlug.projects.v3.$projectParam.runs.$runParam/route.tsx @@ -10,7 +10,6 @@ import { MagnifyingGlassPlusIcon, StopCircleIcon, } from "@heroicons/react/20/solid"; -import type { Location } from "@remix-run/react"; import { useLoaderData, useParams, useRevalidator } from "@remix-run/react"; import { LoaderFunctionArgs, SerializeFrom, json } from "@remix-run/server-runtime"; import { Virtualizer } from "@tanstack/react-virtual"; @@ -66,7 +65,9 @@ import { useProject } from "~/hooks/useProject"; import { useReplaceSearchParams } from "~/hooks/useReplaceSearchParams"; import { Shortcut, useShortcutKeys } from "~/hooks/useShortcutKeys"; import { useUser } from "~/hooks/useUser"; -import { Run, RunPresenter } from "~/presenters/v3/RunPresenter.server"; +import { RunPresenter } from "~/presenters/v3/RunPresenter.server"; +import { getImpersonationId } from "~/services/impersonation.server"; +import { getResizableSnapshot } from "~/services/resizablePanel.server"; import { requireUserId } from "~/services/session.server"; import { cn } from "~/utils/cn"; import { lerp } from "~/utils/lerp"; @@ -79,10 +80,8 @@ import { v3RunStreamingPath, v3RunsPath, } from "~/utils/pathBuilder"; -import { SpanView } from "../resources.orgs.$organizationSlug.projects.v3.$projectParam.runs.$runParam.spans.$spanParam/route"; import { useCurrentPlan } from "../_app.orgs.$organizationSlug/route"; -import { getResizableSnapshot } from "~/services/resizablePanel.server"; -import { getImpersonationId } from "~/services/impersonation.server"; +import { SpanView } from "../resources.orgs.$organizationSlug.projects.v3.$projectParam.runs.$runParam.spans.$spanParam/route"; const resizableSettings = { parent: { @@ -205,7 +204,7 @@ export default function Page() { LeadingIcon={ArrowUturnLeftIcon} shortcut={{ key: "R" }} > - Replay run… + Replay run )} +
{run.logsDeletedAt === null ? ( ; + stop: () => void; + failWithError: (error: TaskRunError) => void; }; export type EventRepoConfig = { @@ -916,6 +918,9 @@ export class EventRepository { ] : []; + let isStopped = false; + let failedWithError: TaskRunError | undefined; + const eventBuilder = { traceId, spanId, @@ -933,10 +938,20 @@ export class EventRepository { } } }, + stop: () => { + isStopped = true; + }, + failWithError: (error: TaskRunError) => { + failedWithError = error; + }, }; const result = await callback(eventBuilder, traceContext, propagatedContext?.traceparent); + if (isStopped) { + return result; + } + const duration = process.hrtime.bigint() - start; const metadata = { @@ -970,13 +985,14 @@ export class EventRepository { parentId, tracestate, duration: options.incomplete ? 0 : duration, - isPartial: options.incomplete, + isPartial: failedWithError ? false : options.incomplete, + isError: !!failedWithError, message: message, serviceName: "api server", serviceNamespace: "trigger.dev", level: "TRACE", kind: options.kind, - status: "OK", + status: failedWithError ? "ERROR" : "OK", startTime, environmentId: options.environment.id, environmentType: options.environment.type, @@ -1004,6 +1020,17 @@ export class EventRepository { payload: options.attributes.payload, payloadType: options.attributes.payloadType, idempotencyKey: options.attributes.idempotencyKey, + events: failedWithError + ? [ + { + name: "exception", + time: new Date(), + properties: { + exception: createExceptionPropertiesFromError(failedWithError), + }, + }, + ] + : undefined, }; if (options.immediate) { diff --git a/apps/webapp/app/v3/marqs/concurrencyMonitor.server.ts b/apps/webapp/app/v3/marqs/concurrencyMonitor.server.ts index cc091a6a08..f68d342d31 100644 --- a/apps/webapp/app/v3/marqs/concurrencyMonitor.server.ts +++ b/apps/webapp/app/v3/marqs/concurrencyMonitor.server.ts @@ -103,18 +103,16 @@ export class MarqsConcurrencyMonitor { async #processKey(key: string, redis: Redis) { key = this.keys.stripKeyPrefix(key); - const orgKey = this.keys.orgCurrentConcurrencyKeyFromQueue(key); const envKey = this.keys.envCurrentConcurrencyKeyFromQueue(key); let runIds: string[] = []; try { // Next, we need to get all the items from the key, and any parent keys (org, env, queue) using sunion. - runIds = await redis.sunion(orgKey, envKey, key); + runIds = await redis.sunion(envKey, key); } catch (e) { this._logger.error("[MarqsConcurrencyMonitor] error during sunion", { key, - orgKey, envKey, runIds, error: e, @@ -136,7 +134,6 @@ export class MarqsConcurrencyMonitor { if (completedRunIds.length === 0) { this._logger.debug("[MarqsConcurrencyMonitor] no completed runs found", { key, - orgKey, envKey, runIds, durationMs, @@ -147,7 +144,6 @@ export class MarqsConcurrencyMonitor { this._logger.debug("[MarqsConcurrencyMonitor] removing completed runs from queue", { key, - orgKey, envKey, completedRunIds, durationMs, @@ -160,7 +156,6 @@ export class MarqsConcurrencyMonitor { const pipeline = redis.pipeline(); pipeline.srem(key, ...completedRunIds); - pipeline.srem(orgKey, ...completedRunIds); pipeline.srem(envKey, ...completedRunIds); try { @@ -168,7 +163,6 @@ export class MarqsConcurrencyMonitor { } catch (e) { this._logger.error("[MarqsConcurrencyMonitor] error removing completed runs from queue", { key, - orgKey, envKey, completedRunIds, error: e, diff --git a/apps/webapp/app/v3/marqs/envPriorityDequeuingStrategy.server.ts b/apps/webapp/app/v3/marqs/envPriorityDequeuingStrategy.server.ts new file mode 100644 index 0000000000..7acaf48818 --- /dev/null +++ b/apps/webapp/app/v3/marqs/envPriorityDequeuingStrategy.server.ts @@ -0,0 +1,95 @@ +import { EnvQueues, MarQSFairDequeueStrategy, MarQSKeyProducer } from "./types"; + +export type EnvPriorityDequeuingStrategyOptions = { + keys: MarQSKeyProducer; + delegate: MarQSFairDequeueStrategy; +}; + +export class EnvPriorityDequeuingStrategy implements MarQSFairDequeueStrategy { + private _delegate: MarQSFairDequeueStrategy; + + constructor(private options: EnvPriorityDequeuingStrategyOptions) { + this._delegate = options.delegate; + } + + async distributeFairQueuesFromParentQueue( + parentQueue: string, + consumerId: string + ): Promise> { + const envQueues = await this._delegate.distributeFairQueuesFromParentQueue( + parentQueue, + consumerId + ); + + return this.#sortQueuesInEnvironmentsByPriority(envQueues); + } + + #sortQueuesInEnvironmentsByPriority(envs: EnvQueues[]): EnvQueues[] { + return envs.map((env) => { + return this.#sortQueuesInEnvironmentByPriority(env); + }); + } + + // Sorts the queues by priority. A higher priority means the queue should be dequeued first. + // All the queues with the same priority should keep the order they were in the original list. + // So that means if all the queues have the same priority, the order should be preserved. + #sortQueuesInEnvironmentByPriority(env: EnvQueues): EnvQueues { + const queues = env.queues; + + // Group queues by their base name (without priority) + const queueGroups = new Map(); + + queues.forEach((queue) => { + const descriptor = this.options.keys.queueDescriptorFromQueue(queue); + const baseQueueName = this.options.keys.queueKey( + descriptor.organization, + descriptor.environment, + descriptor.name, + descriptor.concurrencyKey + ); + + if (!queueGroups.has(baseQueueName)) { + queueGroups.set(baseQueueName, []); + } + + queueGroups.get(baseQueueName)!.push(queue); + }); + + // For each group, keep only the highest priority queue + const resultQueues: string[] = []; + queueGroups.forEach((groupQueues) => { + const sortedGroupQueues = [...groupQueues].sort((a, b) => { + const aPriority = this.#getQueuePriority(a); + const bPriority = this.#getQueuePriority(b); + + if (aPriority === bPriority) { + return 0; + } + + return bPriority - aPriority; + }); + + resultQueues.push(sortedGroupQueues[0]); + }); + + // Sort the final result by priority + const sortedQueues = resultQueues.sort((a, b) => { + const aPriority = this.#getQueuePriority(a); + const bPriority = this.#getQueuePriority(b); + + if (aPriority === bPriority) { + return 0; + } + + return bPriority - aPriority; + }); + + return { envId: env.envId, queues: sortedQueues }; + } + + #getQueuePriority(queue: string): number { + const queueRecord = this.options.keys.queueDescriptorFromQueue(queue); + + return queueRecord.priority ?? 0; + } +} diff --git a/apps/webapp/app/v3/marqs/fairDequeuingStrategy.server.ts b/apps/webapp/app/v3/marqs/fairDequeuingStrategy.server.ts index 6d971ce3c1..986f75c525 100644 --- a/apps/webapp/app/v3/marqs/fairDequeuingStrategy.server.ts +++ b/apps/webapp/app/v3/marqs/fairDequeuingStrategy.server.ts @@ -3,7 +3,7 @@ import { createCache, DefaultStatefulContext, Namespace, Cache as UnkeyCache } f import { MemoryStore } from "@unkey/cache/stores"; import { randomUUID } from "crypto"; import { Redis } from "ioredis"; -import { MarQSFairDequeueStrategy, MarQSKeyProducer } from "./types"; +import { EnvQueues, MarQSFairDequeueStrategy, MarQSKeyProducer } from "./types"; import seedrandom from "seedrandom"; import { Tracer } from "@opentelemetry/api"; import { startSpan } from "../tracing.server"; @@ -33,7 +33,6 @@ export type FairDequeuingStrategyBiases = { export type FairDequeuingStrategyOptions = { redis: Redis; keys: MarQSKeyProducer; - defaultOrgConcurrency: number; defaultEnvConcurrency: number; parentQueueLimit: number; tracer: Tracer; @@ -44,19 +43,19 @@ export type FairDequeuingStrategyOptions = { */ biases?: FairDequeuingStrategyBiases; reuseSnapshotCount?: number; - maximumOrgCount?: number; + maximumEnvCount?: number; }; type FairQueueConcurrency = { current: number; limit: number; + reserve: number; }; type FairQueue = { id: string; age: number; org: string; env: string }; type FairQueueSnapshot = { id: string; - orgs: Record; envs: Record; queues: Array; }; @@ -73,7 +72,6 @@ type WeightedQueue = { const emptyFairQueueSnapshot: FairQueueSnapshot = { id: "empty", - orgs: {}, envs: {}, queues: [], }; @@ -113,7 +111,7 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy { async distributeFairQueuesFromParentQueue( parentQueue: string, consumerId: string - ): Promise> { + ): Promise> { return await startSpan( this.options.tracer, "distributeFairQueuesFromParentQueue", @@ -124,7 +122,6 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy { const snapshot = await this.#createQueueSnapshot(parentQueue, consumerId); span.setAttributes({ - snapshot_org_count: Object.keys(snapshot.orgs).length, snapshot_env_count: Object.keys(snapshot.envs).length, snapshot_queue_count: snapshot.queues.length, }); @@ -135,21 +132,27 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy { return []; } - const shuffledQueues = this.#shuffleQueuesByEnv(snapshot); + const envQueues = this.#shuffleQueuesByEnv(snapshot); - span.setAttribute("shuffled_queue_count", shuffledQueues.length); + span.setAttribute( + "shuffled_queue_count", + envQueues.reduce((sum, env) => sum + env.queues.length, 0) + ); - if (shuffledQueues[0]) { - span.setAttribute("winning_env", this.options.keys.envIdFromQueue(shuffledQueues[0])); - span.setAttribute("winning_org", this.options.keys.orgIdFromQueue(shuffledQueues[0])); + if (envQueues[0]?.queues[0]) { + span.setAttribute("winning_env", envQueues[0].envId); + span.setAttribute( + "winning_org", + this.options.keys.orgIdFromQueue(envQueues[0].queues[0]) + ); } - return shuffledQueues; + return envQueues; } ); } - #shuffleQueuesByEnv(snapshot: FairQueueSnapshot): Array { + #shuffleQueuesByEnv(snapshot: FairQueueSnapshot): Array { const envs = Object.keys(snapshot.envs); const biases = this.options.biases ?? defaultBiases; @@ -215,7 +218,8 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy { } // Helper method to maintain DRY principle - #orderQueuesByEnvs(envs: string[], snapshot: FairQueueSnapshot): Array { + // Update return type + #orderQueuesByEnvs(envs: string[], snapshot: FairQueueSnapshot): Array { const queuesByEnv = snapshot.queues.reduce((acc, queue) => { if (!acc[queue.env]) { acc[queue.env] = []; @@ -224,15 +228,20 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy { return acc; }, {} as Record>); - const queues = envs.reduce((acc, envId) => { + return envs.reduce((acc, envId) => { if (queuesByEnv[envId]) { - // Instead of sorting by age, use weighted random selection - acc.push(...this.#weightedRandomQueueOrder(queuesByEnv[envId])); + // Get ordered queues for this env + const orderedQueues = this.#weightedRandomQueueOrder(queuesByEnv[envId]); + // Only add the env if it has queues + if (orderedQueues.length > 0) { + acc.push({ + envId, + queues: orderedQueues.map((queue) => queue.id), + }); + } } return acc; - }, [] as Array); - - return queues.map((queue) => queue.id); + }, [] as Array); } #weightedRandomQueueOrder(queues: FairQueue[]): FairQueue[] { @@ -344,74 +353,32 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy { return emptyFairQueueSnapshot; } - // Apply org selection if maximumOrgCount is specified - let selectedOrgIds: Set; - if (this.options.maximumOrgCount && this.options.maximumOrgCount > 0) { - selectedOrgIds = this.#selectTopOrgs(queues, this.options.maximumOrgCount); - // Filter queues to only include selected orgs - queues = queues.filter((queue) => selectedOrgIds.has(queue.org)); + // Apply env selection if maximumEnvCount is specified + let selectedEnvIds: Set; + if (this.options.maximumEnvCount && this.options.maximumEnvCount > 0) { + selectedEnvIds = this.#selectTopEnvs(queues, this.options.maximumEnvCount); + // Filter queues to only include selected envs + queues = queues.filter((queue) => selectedEnvIds.has(queue.env)); - span.setAttribute("selected_org_count", selectedOrgIds.size); + span.setAttribute("selected_env_count", selectedEnvIds.size); } span.setAttribute("selected_queue_count", queues.length); - const orgIds = new Set(); const envIds = new Set(); - const envIdToOrgId = new Map(); for (const queue of queues) { - orgIds.add(queue.org); envIds.add(queue.env); - - envIdToOrgId.set(queue.env, queue.org); } - const orgs = await Promise.all( - Array.from(orgIds).map(async (orgId) => { - return { id: orgId, concurrency: await this.#getOrgConcurrency(orgId) }; - }) - ); - - const orgsAtFullConcurrency = orgs.filter( - (org) => org.concurrency.current >= org.concurrency.limit - ); - - const orgIdsAtFullConcurrency = new Set(orgsAtFullConcurrency.map((org) => org.id)); - - const orgsSnapshot = orgs.reduce((acc, org) => { - if (!orgIdsAtFullConcurrency.has(org.id)) { - acc[org.id] = org; - } - - return acc; - }, {} as Record); - - span.setAttributes({ - org_count: orgs.length, - orgs_at_full_concurrency_count: orgsAtFullConcurrency.length, - orgs_snapshot_count: Object.keys(orgsSnapshot).length, - }); - - if (Object.keys(orgsSnapshot).length === 0) { - return emptyFairQueueSnapshot; - } - - const envsWithoutFullOrgs = Array.from(envIds).filter( - (envId) => !orgIdsAtFullConcurrency.has(envIdToOrgId.get(envId)!) - ); - const envs = await Promise.all( - envsWithoutFullOrgs.map(async (envId) => { - return { - id: envId, - concurrency: await this.#getEnvConcurrency(envId, envIdToOrgId.get(envId)!), - }; + Array.from(envIds).map(async (envId) => { + return { id: envId, concurrency: await this.#getEnvConcurrency(envId) }; }) ); const envsAtFullConcurrency = envs.filter( - (env) => env.concurrency.current >= env.concurrency.limit + (env) => env.concurrency.current >= env.concurrency.limit + env.concurrency.reserve ); const envIdsAtFullConcurrency = new Set(envsAtFullConcurrency.map((env) => env.id)); @@ -420,7 +387,6 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy { if (!envIdsAtFullConcurrency.has(env.id)) { acc[env.id] = env; } - return acc; }, {} as Record); @@ -429,14 +395,10 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy { envs_at_full_concurrency_count: envsAtFullConcurrency.length, }); - const queuesSnapshot = queues.filter( - (queue) => - !orgIdsAtFullConcurrency.has(queue.org) && !envIdsAtFullConcurrency.has(queue.env) - ); + const queuesSnapshot = queues.filter((queue) => !envIdsAtFullConcurrency.has(queue.env)); const snapshot = { id: randomUUID(), - orgs: orgsSnapshot, envs: envsSnapshot, queues: queuesSnapshot, }; @@ -455,82 +417,67 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy { }); } - #selectTopOrgs(queues: FairQueue[], maximumOrgCount: number): Set { - // Group queues by org - const queuesByOrg = queues.reduce((acc, queue) => { - if (!acc[queue.org]) { - acc[queue.org] = []; + #selectTopEnvs(queues: FairQueue[], maximumEnvCount: number): Set { + // Group queues by env + const queuesByEnv = queues.reduce((acc, queue) => { + if (!acc[queue.env]) { + acc[queue.env] = []; } - acc[queue.org].push(queue); + acc[queue.env].push(queue); return acc; }, {} as Record); - // Calculate average age for each org - const orgAverageAges = Object.entries(queuesByOrg).map(([orgId, orgQueues]) => { - const averageAge = orgQueues.reduce((sum, q) => sum + q.age, 0) / orgQueues.length; - return { orgId, averageAge }; + // Calculate average age for each env + const envAverageAges = Object.entries(queuesByEnv).map(([envId, envQueues]) => { + const averageAge = envQueues.reduce((sum, q) => sum + q.age, 0) / envQueues.length; + return { envId, averageAge }; }); // Perform weighted shuffle based on average ages - const maxAge = Math.max(...orgAverageAges.map((o) => o.averageAge)); - const weightedOrgs = orgAverageAges.map((org) => ({ - orgId: org.orgId, - weight: org.averageAge / maxAge, // Normalize weights + const maxAge = Math.max(...envAverageAges.map((e) => e.averageAge)); + const weightedEnvs = envAverageAges.map((env) => ({ + envId: env.envId, + weight: env.averageAge / maxAge, // Normalize weights })); - // Select top N orgs using weighted shuffle - const selectedOrgs = new Set(); - let remainingOrgs = [...weightedOrgs]; - let totalWeight = remainingOrgs.reduce((sum, org) => sum + org.weight, 0); + // Select top N envs using weighted shuffle + const selectedEnvs = new Set(); + let remainingEnvs = [...weightedEnvs]; + let totalWeight = remainingEnvs.reduce((sum, env) => sum + env.weight, 0); - while (selectedOrgs.size < maximumOrgCount && remainingOrgs.length > 0) { + while (selectedEnvs.size < maximumEnvCount && remainingEnvs.length > 0) { let random = this._rng() * totalWeight; let index = 0; - while (random > 0 && index < remainingOrgs.length) { - random -= remainingOrgs[index].weight; + while (random > 0 && index < remainingEnvs.length) { + random -= remainingEnvs[index].weight; index++; } index = Math.max(0, index - 1); - selectedOrgs.add(remainingOrgs[index].orgId); - totalWeight -= remainingOrgs[index].weight; - remainingOrgs.splice(index, 1); + selectedEnvs.add(remainingEnvs[index].envId); + totalWeight -= remainingEnvs[index].weight; + remainingEnvs.splice(index, 1); } - return selectedOrgs; + return selectedEnvs; } - async #getOrgConcurrency(orgId: string): Promise { - return await startSpan(this.options.tracer, "getOrgConcurrency", async (span) => { - span.setAttribute("org_id", orgId); - - const [currentValue, limitValue] = await Promise.all([ - this.#getOrgCurrentConcurrency(orgId), - this.#getOrgConcurrencyLimit(orgId), - ]); - - span.setAttribute("current_value", currentValue); - span.setAttribute("limit_value", limitValue); - - return { current: currentValue, limit: limitValue }; - }); - } - - async #getEnvConcurrency(envId: string, orgId: string): Promise { + async #getEnvConcurrency(envId: string): Promise { return await startSpan(this.options.tracer, "getEnvConcurrency", async (span) => { - span.setAttribute("org_id", orgId); span.setAttribute("env_id", envId); - const [currentValue, limitValue] = await Promise.all([ + const [currentValue, limitValue, reserveValue] = await Promise.all([ this.#getEnvCurrentConcurrency(envId), this.#getEnvConcurrencyLimit(envId), + this.#getEnvReserveConcurrency(envId), ]); span.setAttribute("current_value", currentValue); span.setAttribute("limit_value", limitValue); + span.setAttribute("reserve_value", reserveValue); - return { current: currentValue, limit: limitValue }; + return { current: currentValue, limit: limitValue, reserve: reserveValue }; }); } @@ -570,31 +517,31 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy { }); } - async #getOrgConcurrencyLimit(orgId: string) { - return await startSpan(this.options.tracer, "getOrgConcurrencyLimit", async (span) => { - span.setAttribute("org_id", orgId); + async #getEnvConcurrencyLimit(envId: string) { + return await startSpan(this.options.tracer, "getEnvConcurrencyLimit", async (span) => { + span.setAttribute("env_id", envId); - const key = this.options.keys.orgConcurrencyLimitKey(orgId); + const key = this.options.keys.envConcurrencyLimitKey(envId); const result = await this._cache.concurrencyLimit.swr(key, async () => { const value = await this.options.redis.get(key); if (!value) { - return this.options.defaultOrgConcurrency; + return this.options.defaultEnvConcurrency; } return Number(value); }); - return result.val ?? this.options.defaultOrgConcurrency; + return result.val ?? this.options.defaultEnvConcurrency; }); } - async #getOrgCurrentConcurrency(orgId: string) { - return await startSpan(this.options.tracer, "getOrgCurrentConcurrency", async (span) => { - span.setAttribute("org_id", orgId); + async #getEnvCurrentConcurrency(envId: string) { + return await startSpan(this.options.tracer, "getEnvCurrentConcurrency", async (span) => { + span.setAttribute("env_id", envId); - const key = this.options.keys.orgCurrentConcurrencyKey(orgId); + const key = this.options.keys.envCurrentConcurrencyKey(envId); const result = await this.options.redis.scard(key); @@ -604,31 +551,11 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy { }); } - async #getEnvConcurrencyLimit(envId: string) { - return await startSpan(this.options.tracer, "getEnvConcurrencyLimit", async (span) => { - span.setAttribute("env_id", envId); - - const key = this.options.keys.envConcurrencyLimitKey(envId); - - const result = await this._cache.concurrencyLimit.swr(key, async () => { - const value = await this.options.redis.get(key); - - if (!value) { - return this.options.defaultEnvConcurrency; - } - - return Number(value); - }); - - return result.val ?? this.options.defaultEnvConcurrency; - }); - } - - async #getEnvCurrentConcurrency(envId: string) { - return await startSpan(this.options.tracer, "getEnvCurrentConcurrency", async (span) => { + async #getEnvReserveConcurrency(envId: string) { + return await startSpan(this.options.tracer, "getEnvReserveConcurrency", async (span) => { span.setAttribute("env_id", envId); - const key = this.options.keys.envCurrentConcurrencyKey(envId); + const key = this.options.keys.envReserveConcurrencyKey(envId); const result = await this.options.redis.scard(key); @@ -643,7 +570,7 @@ export class NoopFairDequeuingStrategy implements MarQSFairDequeueStrategy { async distributeFairQueuesFromParentQueue( parentQueue: string, consumerId: string - ): Promise> { + ): Promise> { return []; } } diff --git a/apps/webapp/app/v3/marqs/index.server.ts b/apps/webapp/app/v3/marqs/index.server.ts index 341f003464..cef0e95024 100644 --- a/apps/webapp/app/v3/marqs/index.server.ts +++ b/apps/webapp/app/v3/marqs/index.server.ts @@ -22,15 +22,19 @@ import { concurrencyTracker } from "../services/taskRunConcurrencyTracker.server import { attributesFromAuthenticatedEnv, tracer } from "../tracer.server"; import { AsyncWorker } from "./asyncWorker.server"; import { FairDequeuingStrategy } from "./fairDequeuingStrategy.server"; -import { MarQSShortKeyProducer } from "./marqsKeyProducer.server"; +import { MarQSShortKeyProducer } from "./marqsKeyProducer"; import { + EnqueueMessageReserveConcurrencyOptions, MarQSFairDequeueStrategy, MarQSKeyProducer, + MarQSKeyProducerEnv, MessagePayload, MessageQueueSubscriber, VisibilityTimeoutStrategy, } from "./types"; import { V3LegacyRunEngineWorkerVisibilityTimeout } from "./v3VisibilityTimeout.server"; +import { flattenAttributes } from "@trigger.dev/core/v3"; +import { EnvPriorityDequeuingStrategy } from "./envPriorityDequeuingStrategy.server"; const KEY_PREFIX = "marqs:"; @@ -61,6 +65,11 @@ export type MarQSOptions = { subscriber?: MessageQueueSubscriber; }; +export const MarQSPriorityLevel = { + resume: 100, + retry: 10, +} as const; + /** * MarQS - Multitenant Asynchronous Reliable Queueing System (pronounced "markus") */ @@ -99,32 +108,31 @@ export class MarQS { } public async updateEnvConcurrencyLimits(env: AuthenticatedEnvironment) { + const envConcurrencyLimitKey = this.keys.envConcurrencyLimitKey(env); + + logger.debug("Updating env concurrency limits", { + envConcurrencyLimitKey, + service: this.name, + }); + await this.#callUpdateGlobalConcurrencyLimits({ - envConcurrencyLimitKey: this.keys.envConcurrencyLimitKey(env), - orgConcurrencyLimitKey: this.keys.orgConcurrencyLimitKey(env), + envConcurrencyLimitKey, envConcurrencyLimit: env.maximumConcurrencyLimit, - orgConcurrencyLimit: env.organization.maximumConcurrencyLimit, }); } - public async getQueueConcurrencyLimit(env: AuthenticatedEnvironment, queue: string) { + public async getQueueConcurrencyLimit(env: MarQSKeyProducerEnv, queue: string) { const result = await this.redis.get(this.keys.queueConcurrencyLimitKey(env, queue)); return result ? Number(result) : undefined; } - public async getEnvConcurrencyLimit(env: AuthenticatedEnvironment) { + public async getEnvConcurrencyLimit(env: MarQSKeyProducerEnv) { const result = await this.redis.get(this.keys.envConcurrencyLimitKey(env)); return result ? Number(result) : this.options.defaultEnvConcurrency; } - public async getOrgConcurrencyLimit(env: AuthenticatedEnvironment) { - const result = await this.redis.get(this.keys.orgConcurrencyLimitKey(env)); - - return result ? Number(result) : this.options.defaultOrgConcurrency; - } - public async lengthOfQueue( env: AuthenticatedEnvironment, queue: string, @@ -133,7 +141,7 @@ export class MarQS { return this.redis.zcard(this.keys.queueKey(env, queue, concurrencyKey)); } - public async lengthOfEnvQueue(env: AuthenticatedEnvironment) { + public async lengthOfEnvQueue(env: MarQSKeyProducerEnv) { return this.redis.zcard(this.keys.envQueueKey(env)); } @@ -158,19 +166,29 @@ export class MarQS { } public async currentConcurrencyOfQueue( - env: AuthenticatedEnvironment, + env: MarQSKeyProducerEnv, + queue: string, + concurrencyKey?: string + ) { + return this.redis.scard(this.keys.queueCurrentConcurrencyKey(env, queue, concurrencyKey)); + } + + public async reserveConcurrencyOfQueue( + env: MarQSKeyProducerEnv, queue: string, concurrencyKey?: string ) { - return this.redis.scard(this.keys.currentConcurrencyKey(env, queue, concurrencyKey)); + return this.redis.scard( + this.keys.queueReserveConcurrencyKeyFromQueue(this.keys.queueKey(env, queue, concurrencyKey)) + ); } - public async currentConcurrencyOfEnvironment(env: AuthenticatedEnvironment) { + public async currentConcurrencyOfEnvironment(env: MarQSKeyProducerEnv) { return this.redis.scard(this.keys.envCurrentConcurrencyKey(env)); } - public async currentConcurrencyOfOrg(env: AuthenticatedEnvironment) { - return this.redis.scard(this.keys.orgCurrentConcurrencyKey(env)); + public async reserveConcurrencyOfEnvironment(env: MarQSKeyProducerEnv) { + return this.redis.scard(this.keys.envReserveConcurrencyKey(env.id)); } public async enqueueMessage( @@ -179,12 +197,14 @@ export class MarQS { messageId: string, messageData: Record, concurrencyKey?: string, - timestamp?: number + timestamp?: number | Date, + reserve?: EnqueueMessageReserveConcurrencyOptions, + priority?: number ) { return await this.#trace( "enqueueMessage", async (span) => { - const messageQueue = this.keys.queueKey(env, queue, concurrencyKey); + const messageQueue = this.keys.queueKey(env, queue, concurrencyKey, priority); const parentQueue = this.keys.envSharedQueueKey(env); @@ -195,7 +215,12 @@ export class MarQS { data: messageData, queue: messageQueue, concurrencyKey, - timestamp: timestamp ?? Date.now(), + timestamp: + typeof timestamp === "undefined" + ? Date.now() + : typeof timestamp === "number" + ? timestamp + : timestamp.getTime(), messageId, parentQueue, }; @@ -207,9 +232,18 @@ export class MarQS { [SemanticAttributes.PARENT_QUEUE]: parentQueue, }); - await this.#callEnqueueMessage(messagePayload); + if (reserve) { + span.setAttribute("reserve_message_id", reserve.messageId); + span.setAttribute("reserve_recursive_queue", reserve.recursiveQueue); + } + + const result = await this.#callEnqueueMessage(messagePayload, reserve); - await this.options.subscriber?.messageEnqueued(messagePayload); + if (result) { + await this.options.subscriber?.messageEnqueued(messagePayload); + } + + return result; }, { kind: SpanKind.PRODUCER, @@ -223,6 +257,75 @@ export class MarQS { ); } + public async replaceMessage( + messageId: string, + messageData: Record, + timestamp?: number, + priority?: number, + inplace?: boolean + ) { + return this.#trace( + "replaceMessage", + async (span) => { + const oldMessage = await this.readMessage(messageId); + + if (!oldMessage) { + return; + } + + const queue = this.keys.queueKeyFromQueue(oldMessage.queue, priority); + + span.setAttributes({ + [SemanticAttributes.QUEUE]: queue, + [SemanticAttributes.MESSAGE_ID]: oldMessage.messageId, + [SemanticAttributes.CONCURRENCY_KEY]: oldMessage.concurrencyKey, + [SemanticAttributes.PARENT_QUEUE]: oldMessage.parentQueue, + }); + + const traceContext = { + traceparent: oldMessage.data.traceparent, + tracestate: oldMessage.data.tracestate, + }; + + const newMessage: MessagePayload = { + version: "1", + // preserve original trace context + data: { ...oldMessage.data, ...messageData, ...traceContext, queue }, + queue, + concurrencyKey: oldMessage.concurrencyKey, + timestamp: timestamp ?? Date.now(), + messageId, + parentQueue: oldMessage.parentQueue, + }; + + if (inplace) { + await this.#callReplaceMessage(newMessage); + return; + } + + await this.options.visibilityTimeoutStrategy.cancelHeartbeat(messageId); + + await this.#callAcknowledgeMessage({ + parentQueue: oldMessage.parentQueue, + messageQueue: oldMessage.queue, + messageId, + }); + + await this.#callEnqueueMessage(newMessage); + + await this.options.subscriber?.messageReplaced(newMessage); + }, + { + kind: SpanKind.CONSUMER, + attributes: { + [SEMATTRS_MESSAGING_OPERATION]: "replace", + [SEMATTRS_MESSAGE_ID]: messageId, + [SEMATTRS_MESSAGING_SYSTEM]: "marqs", + }, + } + ); + } + public async dequeueMessageInEnv(env: AuthenticatedEnvironment) { return this.#trace( "dequeueMessageInEnv", @@ -233,24 +336,21 @@ export class MarQS { span.setAttribute(SemanticAttributes.CONSUMER_ID, env.id); // Get prioritized list of queues to try - const queues = + const environments = await this.options.envQueuePriorityStrategy.distributeFairQueuesFromParentQueue( parentQueue, env.id ); + const queues = environments.flatMap((e) => e.queues); + + span.setAttribute("env_count", environments.length); span.setAttribute("queue_count", queues.length); for (const messageQueue of queues) { const messageData = await this.#callDequeueMessage({ messageQueue, parentQueue, - concurrencyLimitKey: this.keys.concurrencyLimitKeyFromQueue(messageQueue), - currentConcurrencyKey: this.keys.currentConcurrencyKeyFromQueue(messageQueue), - envConcurrencyLimitKey: this.keys.envConcurrencyLimitKeyFromQueue(messageQueue), - envCurrentConcurrencyKey: this.keys.envCurrentConcurrencyKeyFromQueue(messageQueue), - orgConcurrencyLimitKey: this.keys.orgConcurrencyLimitKeyFromQueue(messageQueue), - orgCurrentConcurrencyKey: this.keys.orgCurrentConcurrencyKeyFromQueue(messageQueue), }); if (!messageData) { @@ -280,11 +380,7 @@ export class MarQS { await this.#callAcknowledgeMessage({ parentQueue, - messageKey: this.keys.messageKey(messageData.messageId), messageQueue: messageQueue, - concurrencyKey: this.keys.currentConcurrencyKeyFromQueue(messageQueue), - envConcurrencyKey: this.keys.envCurrentConcurrencyKeyFromQueue(messageQueue), - orgConcurrencyKey: this.keys.orgCurrentConcurrencyKeyFromQueue(messageQueue), messageId: messageData.messageId, }); @@ -327,70 +423,78 @@ export class MarQS { span.setAttribute(SemanticAttributes.PARENT_QUEUE, parentQueue); // Get prioritized list of queues to try - const queues = await this.options.queuePriorityStrategy.distributeFairQueuesFromParentQueue( - parentQueue, - consumerId - ); + const envQueues = + await this.options.queuePriorityStrategy.distributeFairQueuesFromParentQueue( + parentQueue, + consumerId + ); - span.setAttribute("queue_count", queues.length); + span.setAttribute("environment_count", envQueues.length); - if (queues.length === 0) { + if (envQueues.length === 0) { return; } + let attemptedEnvs = 0; + let attemptedQueues = 0; + // Try each queue in order until we successfully dequeue a message - for (const messageQueue of queues) { - try { - const messageData = await this.#callDequeueMessage({ - messageQueue, - parentQueue, - concurrencyLimitKey: this.keys.concurrencyLimitKeyFromQueue(messageQueue), - currentConcurrencyKey: this.keys.currentConcurrencyKeyFromQueue(messageQueue), - envConcurrencyLimitKey: this.keys.envConcurrencyLimitKeyFromQueue(messageQueue), - envCurrentConcurrencyKey: this.keys.envCurrentConcurrencyKeyFromQueue(messageQueue), - orgConcurrencyLimitKey: this.keys.orgConcurrencyLimitKeyFromQueue(messageQueue), - orgCurrentConcurrencyKey: this.keys.orgCurrentConcurrencyKeyFromQueue(messageQueue), - }); + for (const env of envQueues) { + attemptedEnvs++; - if (!messageData) { - continue; // Try next queue if no message was dequeued - } + for (const messageQueue of env.queues) { + attemptedQueues++; - const message = await this.readMessage(messageData.messageId); - - if (message) { - const ageOfMessageInMs = Date.now() - message.timestamp; - - span.setAttributes({ - [SEMATTRS_MESSAGE_ID]: message.messageId, - [SemanticAttributes.QUEUE]: message.queue, - [SemanticAttributes.MESSAGE_ID]: message.messageId, - [SemanticAttributes.CONCURRENCY_KEY]: message.concurrencyKey, - [SemanticAttributes.PARENT_QUEUE]: message.parentQueue, - age_in_seconds: ageOfMessageInMs / 1000, - attempted_queues: queues.indexOf(messageQueue) + 1, // How many queues we tried before success - message_timestamp: message.timestamp, - message_age: Date.now() - message.timestamp, + try { + const messageData = await this.#callDequeueMessage({ + messageQueue, + parentQueue, }); - await this.options.subscriber?.messageDequeued(message); - - await this.options.visibilityTimeoutStrategy.heartbeat( - messageData.messageId, - this.visibilityTimeoutInMs - ); - - return message; + if (!messageData) { + continue; // Try next queue if no message was dequeued + } + + const message = await this.readMessage(messageData.messageId); + + if (message) { + const ageOfMessageInMs = Date.now() - message.timestamp; + + span.setAttributes({ + [SEMATTRS_MESSAGE_ID]: message.messageId, + [SemanticAttributes.QUEUE]: message.queue, + [SemanticAttributes.MESSAGE_ID]: message.messageId, + [SemanticAttributes.CONCURRENCY_KEY]: message.concurrencyKey, + [SemanticAttributes.PARENT_QUEUE]: message.parentQueue, + age_in_seconds: ageOfMessageInMs / 1000, + attempted_queues: attemptedQueues, // How many queues we tried before success + attempted_envs: attemptedEnvs, // How many environments we tried before success + message_timestamp: message.timestamp, + message_age: Date.now() - message.timestamp, + ...flattenAttributes(message.data, "message"), + }); + + await this.options.subscriber?.messageDequeued(message); + + await this.options.visibilityTimeoutStrategy.heartbeat( + messageData.messageId, + this.visibilityTimeoutInMs + ); + + return message; + } + } catch (error) { + // Log error but continue trying other queues + logger.warn(`[${this.name}] Failed to dequeue from queue ${messageQueue}`, { error }); + continue; } - } catch (error) { - // Log error but continue trying other queues - logger.warn(`[${this.name}] Failed to dequeue from queue ${messageQueue}`, { error }); - continue; } } // If we get here, we tried all queues but couldn't dequeue a message - span.setAttribute("attempted_queues", queues.length); + span.setAttribute("attempted_queues", attemptedQueues); + span.setAttribute("attempted_envs", attemptedEnvs); + return; }, { @@ -430,11 +534,7 @@ export class MarQS { await this.#callAcknowledgeMessage({ parentQueue: message.parentQueue, - messageKey: this.keys.messageKey(messageId), messageQueue: message.queue, - concurrencyKey: this.keys.currentConcurrencyKeyFromQueue(message.queue), - envConcurrencyKey: this.keys.envCurrentConcurrencyKeyFromQueue(message.queue), - orgConcurrencyKey: this.keys.orgCurrentConcurrencyKeyFromQueue(message.queue), messageId, }); @@ -451,69 +551,72 @@ export class MarQS { ); } - public async replaceMessage( + /** + * Negative acknowledge a message, which will requeue the message. + * Returns whether it went back into the queue or not. + */ + public async nackMessage( messageId: string, - messageData: Record, - timestamp?: number, - inplace?: boolean + retryAt: number = Date.now(), + updates?: Record ) { return this.#trace( - "replaceMessage", + "nackMessage", async (span) => { - const oldMessage = await this.readMessage(messageId); + const message = await this.readMessage(messageId); - if (!oldMessage) { - return; + if (!message) { + logger.debug(`[${this.name}].nackMessage() message not found`, { + messageId, + retryAt, + updates, + service: this.name, + }); + return false; } - span.setAttributes({ - [SemanticAttributes.QUEUE]: oldMessage.queue, - [SemanticAttributes.MESSAGE_ID]: oldMessage.messageId, - [SemanticAttributes.CONCURRENCY_KEY]: oldMessage.concurrencyKey, - [SemanticAttributes.PARENT_QUEUE]: oldMessage.parentQueue, - }); + const nackCount = await this.#getNackCount(messageId); - const traceContext = { - traceparent: oldMessage.data.traceparent, - tracestate: oldMessage.data.tracestate, - }; + span.setAttribute("nack_count", nackCount); - const newMessage: MessagePayload = { - version: "1", - // preserve original trace context - data: { ...oldMessage.data, ...messageData, ...traceContext }, - queue: oldMessage.queue, - concurrencyKey: oldMessage.concurrencyKey, - timestamp: timestamp ?? Date.now(), - messageId, - parentQueue: oldMessage.parentQueue, - }; + if (nackCount >= this.options.maximumNackCount) { + logger.debug(`[${this.name}].nackMessage() maximum nack count reached`, { + messageId, + retryAt, + updates, + service: this.name, + }); - if (inplace) { - await this.#callReplaceMessage(newMessage); - return; + span.setAttribute("maximum_nack_count_reached", true); + + // If we have reached the maximum nack count, we will ack the message + await this.acknowledgeMessage(messageId, "maximum nack count reached"); + return false; + } + + span.setAttributes({ + [SemanticAttributes.QUEUE]: message.queue, + [SemanticAttributes.MESSAGE_ID]: message.messageId, + [SemanticAttributes.CONCURRENCY_KEY]: message.concurrencyKey, + [SemanticAttributes.PARENT_QUEUE]: message.parentQueue, + }); + + if (updates) { + await this.replaceMessage(messageId, updates, retryAt, undefined, true); } await this.options.visibilityTimeoutStrategy.cancelHeartbeat(messageId); - await this.#callAcknowledgeMessage({ - parentQueue: oldMessage.parentQueue, - messageKey: this.keys.messageKey(messageId), - messageQueue: oldMessage.queue, - concurrencyKey: this.keys.currentConcurrencyKeyFromQueue(oldMessage.queue), - envConcurrencyKey: this.keys.envCurrentConcurrencyKeyFromQueue(oldMessage.queue), - orgConcurrencyKey: this.keys.orgCurrentConcurrencyKeyFromQueue(oldMessage.queue), - messageId, - }); + await this.#callNackMessage(messageId, message, retryAt); - await this.#callEnqueueMessage(newMessage); + await this.options.subscriber?.messageNacked(message); - await this.options.subscriber?.messageReplaced(newMessage); + return true; }, { kind: SpanKind.CONSUMER, attributes: { - [SEMATTRS_MESSAGING_OPERATION]: "replace", + [SEMATTRS_MESSAGING_OPERATION]: "nack", [SEMATTRS_MESSAGE_ID]: messageId, [SEMATTRS_MESSAGING_SYSTEM]: "marqs", }, @@ -542,65 +645,6 @@ export class MarQS { ); } - public async releaseConcurrency(messageId: string, releaseForRun: boolean = false) { - return this.#trace( - "releaseConcurrency", - async (span) => { - span.setAttributes({ - [SemanticAttributes.MESSAGE_ID]: messageId, - }); - - const message = await this.readMessage(messageId); - - if (!message) { - logger.log(`[${this.name}].releaseConcurrency() message not found`, { - messageId, - releaseForRun, - service: this.name, - }); - return; - } - - span.setAttributes({ - [SemanticAttributes.QUEUE]: message.queue, - [SemanticAttributes.MESSAGE_ID]: message.messageId, - [SemanticAttributes.CONCURRENCY_KEY]: message.concurrencyKey, - [SemanticAttributes.PARENT_QUEUE]: message.parentQueue, - }); - - const concurrencyKey = this.keys.currentConcurrencyKeyFromQueue(message.queue); - const envConcurrencyKey = this.keys.envCurrentConcurrencyKeyFromQueue(message.queue); - const orgConcurrencyKey = this.keys.orgCurrentConcurrencyKeyFromQueue(message.queue); - - logger.debug("Calling releaseConcurrency", { - messageId, - queue: message.queue, - concurrencyKey, - envConcurrencyKey, - orgConcurrencyKey, - service: this.name, - releaseForRun, - }); - - return this.redis.releaseConcurrency( - //don't release the for the run, it breaks concurrencyLimits - releaseForRun ? concurrencyKey : "", - envConcurrencyKey, - orgConcurrencyKey, - message.messageId - ); - }, - { - kind: SpanKind.CONSUMER, - attributes: { - [SEMATTRS_MESSAGING_OPERATION]: "releaseConcurrency", - [SEMATTRS_MESSAGE_ID]: messageId, - [SEMATTRS_MESSAGING_SYSTEM]: "marqs", - }, - } - ); - } - async #trace( name: string, fn: (span: Span) => Promise, @@ -637,103 +681,20 @@ export class MarQS { ); } - /** - * Negative acknowledge a message, which will requeue the message. - * Returns whether it went back into the queue or not. - */ - public async nackMessage( - messageId: string, - retryAt: number = Date.now(), - updates?: Record - ) { - return this.#trace( - "nackMessage", - async (span) => { - const message = await this.readMessage(messageId); + async #getNackCount(messageId: string): Promise { + const result = await this.redis.get(this.keys.nackCounterKey(messageId)); - if (!message) { - logger.debug(`[${this.name}].nackMessage() message not found`, { - messageId, - retryAt, - updates, - service: this.name, - }); - return false; - } + return result ? Number(result) : 0; + } - const nackCount = await this.#getNackCount(messageId); + // This should increment by the number of seconds, but with a max value of Date.now() + visibilityTimeoutInMs + public async heartbeatMessage(messageId: string) { + await this.options.visibilityTimeoutStrategy.heartbeat(messageId, this.visibilityTimeoutInMs); + } - span.setAttribute("nack_count", nackCount); - - if (nackCount >= this.options.maximumNackCount) { - logger.debug(`[${this.name}].nackMessage() maximum nack count reached`, { - messageId, - retryAt, - updates, - service: this.name, - }); - - span.setAttribute("maximum_nack_count_reached", true); - - // If we have reached the maximum nack count, we will ack the message - await this.acknowledgeMessage(messageId, "maximum nack count reached"); - return false; - } - - span.setAttributes({ - [SemanticAttributes.QUEUE]: message.queue, - [SemanticAttributes.MESSAGE_ID]: message.messageId, - [SemanticAttributes.CONCURRENCY_KEY]: message.concurrencyKey, - [SemanticAttributes.PARENT_QUEUE]: message.parentQueue, - }); - - if (updates) { - await this.replaceMessage(messageId, updates, retryAt, true); - } - - await this.options.visibilityTimeoutStrategy.cancelHeartbeat(messageId); - - await this.#callNackMessage({ - messageKey: this.keys.messageKey(messageId), - messageQueue: message.queue, - parentQueue: message.parentQueue, - concurrencyKey: this.keys.currentConcurrencyKeyFromQueue(message.queue), - envConcurrencyKey: this.keys.envCurrentConcurrencyKeyFromQueue(message.queue), - orgConcurrencyKey: this.keys.orgCurrentConcurrencyKeyFromQueue(message.queue), - nackCounterKey: this.keys.nackCounterKey(messageId), - messageId, - messageScore: retryAt, - }); - - await this.options.subscriber?.messageNacked(message); - - return true; - }, - { - kind: SpanKind.CONSUMER, - attributes: { - [SEMATTRS_MESSAGING_OPERATION]: "nack", - [SEMATTRS_MESSAGE_ID]: messageId, - [SEMATTRS_MESSAGING_SYSTEM]: "marqs", - }, - } - ); - } - - async #getNackCount(messageId: string): Promise { - const result = await this.redis.get(this.keys.nackCounterKey(messageId)); - - return result ? Number(result) : 0; - } - - // This should increment by the number of seconds, but with a max value of Date.now() + visibilityTimeoutInMs - public async heartbeatMessage(messageId: string) { - await this.options.visibilityTimeoutStrategy.heartbeat(messageId, this.visibilityTimeoutInMs); - } - - get visibilityTimeoutInMs() { - return this.options.visibilityTimeoutInMs ?? 300000; // 5 minutes - } + get visibilityTimeoutInMs() { + return this.options.visibilityTimeoutInMs ?? 300000; // 5 minutes + } async readMessage(messageId: string) { return this.#trace( @@ -928,67 +889,247 @@ export class MarQS { }); } - async #callEnqueueMessage(message: MessagePayload) { - const concurrencyKey = this.keys.currentConcurrencyKeyFromQueue(message.queue); - const envConcurrencyKey = this.keys.envCurrentConcurrencyKeyFromQueue(message.queue); - const orgConcurrencyKey = this.keys.orgCurrentConcurrencyKeyFromQueue(message.queue); + async #callEnqueueMessage( + message: MessagePayload, + reserve?: EnqueueMessageReserveConcurrencyOptions + ) { + const queueKey = message.queue; + const parentQueueKey = message.parentQueue; + const messageKey = this.keys.messageKey(message.messageId); + const queueCurrentConcurrencyKey = this.keys.queueCurrentConcurrencyKeyFromQueue(message.queue); + const queueReserveConcurrencyKey = this.keys.queueReserveConcurrencyKeyFromQueue(message.queue); + const envCurrentConcurrencyKey = this.keys.envCurrentConcurrencyKeyFromQueue(message.queue); + const envReserveConcurrencyKey = this.keys.envReserveConcurrencyKeyFromQueue(message.queue); + const envQueueKey = this.keys.envQueueKeyFromQueue(message.queue); + + const queueName = message.queue; + const messageId = message.messageId; + const messageData = JSON.stringify(message); + const messageScore = String(message.timestamp); + + if (!reserve) { + logger.debug("Calling enqueueMessage", { + service: this.name, + queueKey, + parentQueueKey, + messageKey, + queueCurrentConcurrencyKey, + queueReserveConcurrencyKey, + envCurrentConcurrencyKey, + envReserveConcurrencyKey, + envQueueKey, + queueName, + messageId, + messageData, + messageScore, + }); - logger.debug("Calling enqueueMessage", { - messagePayload: message, - concurrencyKey, - envConcurrencyKey, - orgConcurrencyKey, - service: this.name, - }); + const result = await this.redis.enqueueMessage( + queueKey, + parentQueueKey, + messageKey, + queueCurrentConcurrencyKey, + queueReserveConcurrencyKey, + envCurrentConcurrencyKey, + envReserveConcurrencyKey, + envQueueKey, + queueName, + messageId, + messageData, + messageScore + ); - return this.redis.enqueueMessage( - message.queue, - message.parentQueue, - this.keys.messageKey(message.messageId), - concurrencyKey, - envConcurrencyKey, - orgConcurrencyKey, - this.keys.envQueueKeyFromQueue(message.queue), - message.queue, - message.messageId, - JSON.stringify(message), - String(message.timestamp) - ); + logger.debug("enqueueMessage result", { + service: this.name, + queueKey, + parentQueueKey, + messageKey, + queueCurrentConcurrencyKey, + queueReserveConcurrencyKey, + envCurrentConcurrencyKey, + envReserveConcurrencyKey, + envQueueKey, + queueName, + messageId, + messageData, + messageScore, + result, + }); + + return true; + } + + const envConcurrencyLimitKey = this.keys.envConcurrencyLimitKeyFromQueue(message.queue); + const reserveMessageId = reserve.messageId; + const defaultEnvConcurrencyLimit = String(this.options.defaultEnvConcurrency); + + if (!reserve.recursiveQueue) { + logger.debug("Calling enqueueMessageWithReservingConcurrency", { + service: this.name, + queueKey, + parentQueueKey, + messageKey, + queueCurrentConcurrencyKey, + queueReserveConcurrencyKey, + envCurrentConcurrencyKey, + envReserveConcurrencyKey, + envConcurrencyLimitKey, + envQueueKey, + queueName, + messageId, + messageData, + messageScore, + reserveMessageId, + defaultEnvConcurrencyLimit, + }); + + const result = await this.redis.enqueueMessageWithReservingConcurrency( + queueKey, + parentQueueKey, + messageKey, + queueCurrentConcurrencyKey, + queueReserveConcurrencyKey, + envCurrentConcurrencyKey, + envReserveConcurrencyKey, + envConcurrencyLimitKey, + envQueueKey, + queueName, + messageId, + messageData, + messageScore, + reserveMessageId, + defaultEnvConcurrencyLimit + ); + + logger.debug("enqueueMessageWithReservingConcurrency result", { + service: this.name, + queueKey, + parentQueueKey, + messageKey, + queueCurrentConcurrencyKey, + queueReserveConcurrencyKey, + envCurrentConcurrencyKey, + envReserveConcurrencyKey, + envConcurrencyLimitKey, + envQueueKey, + queueName, + messageId, + messageData, + messageScore, + reserveMessageId, + defaultEnvConcurrencyLimit, + result, + }); + + return true; + } else { + const queueConcurrencyLimitKey = this.keys.queueConcurrencyLimitKeyFromQueue(message.queue); + + logger.debug("Calling enqueueMessageWithReservingConcurrencyForRecursiveQueue", { + service: this.name, + queueKey, + parentQueueKey, + messageKey, + queueCurrentConcurrencyKey, + queueReserveConcurrencyKey, + queueConcurrencyLimitKey, + envCurrentConcurrencyKey, + envReserveConcurrencyKey, + envConcurrencyLimitKey, + envQueueKey, + queueName, + messageId, + messageData, + messageScore, + reserveMessageId, + defaultEnvConcurrencyLimit, + }); + + const result = await this.redis.enqueueMessageWithReservingConcurrencyOnRecursiveQueue( + queueKey, + parentQueueKey, + messageKey, + queueCurrentConcurrencyKey, + queueReserveConcurrencyKey, + queueConcurrencyLimitKey, + envCurrentConcurrencyKey, + envReserveConcurrencyKey, + envConcurrencyLimitKey, + envQueueKey, + queueName, + messageId, + messageData, + messageScore, + reserveMessageId, + defaultEnvConcurrencyLimit + ); + + logger.debug("enqueueMessageWithReservingConcurrencyOnRecursiveQueue result", { + service: this.name, + queueKey, + parentQueueKey, + messageKey, + queueCurrentConcurrencyKey, + queueReserveConcurrencyKey, + queueConcurrencyLimitKey, + envCurrentConcurrencyKey, + envReserveConcurrencyKey, + envConcurrencyLimitKey, + envQueueKey, + queueName, + messageId, + messageData, + messageScore, + reserveMessageId, + defaultEnvConcurrencyLimit, + result, + }); + + return !!result; + } } async #callDequeueMessage({ messageQueue, parentQueue, - concurrencyLimitKey, - envConcurrencyLimitKey, - orgConcurrencyLimitKey, - currentConcurrencyKey, - envCurrentConcurrencyKey, - orgCurrentConcurrencyKey, }: { messageQueue: string; parentQueue: string; - concurrencyLimitKey: string; - envConcurrencyLimitKey: string; - orgConcurrencyLimitKey: string; - currentConcurrencyKey: string; - envCurrentConcurrencyKey: string; - orgCurrentConcurrencyKey: string; }) { + const queueConcurrencyLimitKey = this.keys.queueConcurrencyLimitKeyFromQueue(messageQueue); + const queueCurrentConcurrencyKey = this.keys.queueCurrentConcurrencyKeyFromQueue(messageQueue); + const envConcurrencyLimitKey = this.keys.envConcurrencyLimitKeyFromQueue(messageQueue); + const envCurrentConcurrencyKey = this.keys.envCurrentConcurrencyKeyFromQueue(messageQueue); + const envReserveConcurrencyKey = this.keys.envReserveConcurrencyKeyFromQueue(messageQueue); + const queueReserveConcurrencyKey = this.keys.queueReserveConcurrencyKeyFromQueue(messageQueue); + const envQueueKey = this.keys.envQueueKeyFromQueue(messageQueue); + + logger.debug("Calling dequeueMessage", { + messageQueue, + parentQueue, + queueConcurrencyLimitKey, + envConcurrencyLimitKey, + queueCurrentConcurrencyKey, + queueReserveConcurrencyKey, + envCurrentConcurrencyKey, + envReserveConcurrencyKey, + envQueueKey, + service: this.name, + }); + const result = await this.redis.dequeueMessage( messageQueue, parentQueue, - concurrencyLimitKey, + queueConcurrencyLimitKey, envConcurrencyLimitKey, - orgConcurrencyLimitKey, - currentConcurrencyKey, + queueCurrentConcurrencyKey, + queueReserveConcurrencyKey, envCurrentConcurrencyKey, - orgCurrentConcurrencyKey, - this.keys.envQueueKeyFromQueue(messageQueue), + envReserveConcurrencyKey, + envQueueKey, messageQueue, String(Date.now()), - String(this.options.defaultEnvConcurrency), - String(this.options.defaultOrgConcurrency) + String(this.options.defaultEnvConcurrency) ); if (!result) { @@ -1024,29 +1165,28 @@ export class MarQS { async #callAcknowledgeMessage({ parentQueue, - messageKey, messageQueue, - concurrencyKey, - envConcurrencyKey, - orgConcurrencyKey, messageId, }: { parentQueue: string; - messageKey: string; messageQueue: string; - concurrencyKey: string; - envConcurrencyKey: string; - orgConcurrencyKey: string; messageId: string; }) { + const messageKey = this.keys.messageKey(messageId); + const concurrencyKey = this.keys.queueCurrentConcurrencyKeyFromQueue(messageQueue); + const envConcurrencyKey = this.keys.envCurrentConcurrencyKeyFromQueue(messageQueue); + const envReserveConcurrencyKey = this.keys.envReserveConcurrencyKeyFromQueue(messageQueue); + const queueReserveConcurrencyKey = this.keys.queueReserveConcurrencyKeyFromQueue(messageQueue); + const envQueueKey = this.keys.envQueueKeyFromQueue(messageQueue); + logger.debug("Calling acknowledgeMessage", { messageKey, messageQueue, concurrencyKey, envConcurrencyKey, - orgConcurrencyKey, messageId, parentQueue, + envQueueKey, service: this.name, }); @@ -1055,106 +1195,63 @@ export class MarQS { messageKey, messageQueue, concurrencyKey, + queueReserveConcurrencyKey, envConcurrencyKey, - orgConcurrencyKey, - this.keys.envQueueKeyFromQueue(messageQueue), + envReserveConcurrencyKey, + envQueueKey, messageId, messageQueue ); } - async #callNackMessage({ - messageKey, - messageQueue, - parentQueue, - concurrencyKey, - envConcurrencyKey, - orgConcurrencyKey, - nackCounterKey, - messageId, - messageScore, - }: { - messageKey: string; - messageQueue: string; - parentQueue: string; - concurrencyKey: string; - envConcurrencyKey: string; - orgConcurrencyKey: string; - nackCounterKey: string; - messageId: string; - messageScore: number; - }) { + async #callNackMessage(messageId: string, message: MessagePayload, messageScore: number) { + const messageKey = this.keys.messageKey(message.messageId); + const queueKey = message.queue; + const parentQueueKey = message.parentQueue; + const queueCurrentConcurrencyKey = this.keys.queueCurrentConcurrencyKeyFromQueue(message.queue); + const envCurrentConcurrencyKey = this.keys.envCurrentConcurrencyKeyFromQueue(message.queue); + const nackCounterKey = this.keys.nackCounterKey(message.messageId); + const envQueueKey = this.keys.envQueueKeyFromQueue(message.queue); + const queueName = message.queue; + logger.debug("Calling nackMessage", { messageKey, - messageQueue, - parentQueue, - concurrencyKey, - envConcurrencyKey, - orgConcurrencyKey, + queueKey, + parentQueueKey, + queueCurrentConcurrencyKey, + envCurrentConcurrencyKey, nackCounterKey, messageId, messageScore, + envQueueKey, service: this.name, }); return this.redis.nackMessage( messageKey, - messageQueue, - parentQueue, - concurrencyKey, - envConcurrencyKey, - orgConcurrencyKey, - this.keys.envQueueKeyFromQueue(messageQueue), + queueKey, + parentQueueKey, + queueCurrentConcurrencyKey, + envCurrentConcurrencyKey, + envQueueKey, nackCounterKey, - messageQueue, + queueName, messageId, String(Date.now()), String(messageScore) ); } - async #callCalculateQueueCurrentConcurrencies({ - currentConcurrencyKey, - currentEnvConcurrencyKey, - currentOrgConcurrencyKey, - }: { - currentConcurrencyKey: string; - currentEnvConcurrencyKey: string; - currentOrgConcurrencyKey: string; - }) { - const currentConcurrencies = await this.redis.calculateQueueCurrentConcurrencies( - currentConcurrencyKey, - currentEnvConcurrencyKey, - currentOrgConcurrencyKey - ); - - const orgCurrent = Number(currentConcurrencies[0]); - const envCurrent = Number(currentConcurrencies[1]); - const queueCurrent = Number(currentConcurrencies[2]); - - return { - queue: queueCurrent, - env: envCurrent, - org: orgCurrent, - }; - } - #callUpdateGlobalConcurrencyLimits({ envConcurrencyLimitKey, - orgConcurrencyLimitKey, envConcurrencyLimit, - orgConcurrencyLimit, }: { envConcurrencyLimitKey: string; - orgConcurrencyLimitKey: string; envConcurrencyLimit: number; - orgConcurrencyLimit: number; }) { return this.redis.updateGlobalConcurrencyLimits( envConcurrencyLimitKey, - orgConcurrencyLimitKey, - String(envConcurrencyLimit), - String(orgConcurrencyLimit) + String(envConcurrencyLimit) ); } @@ -1190,15 +1287,16 @@ export class MarQS { #registerCommands() { this.redis.defineCommand("enqueueMessage", { - numberOfKeys: 7, + numberOfKeys: 8, lua: ` -local queue = KEYS[1] -local parentQueue = KEYS[2] +local queueKey = KEYS[1] +local parentQueueKey = KEYS[2] local messageKey = KEYS[3] -local concurrencyKey = KEYS[4] -local envCurrentConcurrencyKey = KEYS[5] -local orgCurrentConcurrencyKey = KEYS[6] -local envQueue = KEYS[7] +local queueCurrentConcurrencyKey = KEYS[4] +local queueReserveConcurrencyKey = KEYS[5] +local envCurrentConcurrencyKey = KEYS[6] +local envReserveConcurrencyKey = KEYS[7] +local envQueueKey = KEYS[8] local queueName = ARGV[1] local messageId = ARGV[2] @@ -1209,73 +1307,199 @@ local messageScore = ARGV[4] redis.call('SET', messageKey, messageData) -- Add the message to the queue -redis.call('ZADD', queue, messageScore, messageId) +redis.call('ZADD', queueKey, messageScore, messageId) -- Add the message to the env queue -redis.call('ZADD', envQueue, messageScore, messageId) +redis.call('ZADD', envQueueKey, messageScore, messageId) -- Rebalance the parent queue -local earliestMessage = redis.call('ZRANGE', queue, 0, 0, 'WITHSCORES') +local earliestMessage = redis.call('ZRANGE', queueKey, 0, 0, 'WITHSCORES') if #earliestMessage == 0 then - redis.call('ZREM', parentQueue, queueName) + redis.call('ZREM', parentQueueKey, queueName) else - redis.call('ZADD', parentQueue, earliestMessage[2], queueName) + redis.call('ZADD', parentQueueKey, earliestMessage[2], queueName) end -- Update the concurrency keys -redis.call('SREM', concurrencyKey, messageId) +redis.call('SREM', queueCurrentConcurrencyKey, messageId) redis.call('SREM', envCurrentConcurrencyKey, messageId) -redis.call('SREM', orgCurrentConcurrencyKey, messageId) +redis.call('SREM', envReserveConcurrencyKey, messageId) +redis.call('SREM', queueReserveConcurrencyKey, messageId) + +return true + `, + }); + + this.redis.defineCommand("enqueueMessageWithReservingConcurrency", { + numberOfKeys: 9, + lua: ` +local queueKey = KEYS[1] +local parentQueueKey = KEYS[2] +local messageKey = KEYS[3] +local queueCurrentConcurrencyKey = KEYS[4] +local queueReserveConcurrencyKey = KEYS[5] +local envCurrentConcurrencyKey = KEYS[6] +local envReserveConcurrencyKey = KEYS[7] +local envConcurrencyLimitKey = KEYS[8] +local envQueueKey = KEYS[9] + +local queueName = ARGV[1] +local messageId = ARGV[2] +local messageData = ARGV[3] +local messageScore = ARGV[4] +local reserveMessageId = ARGV[5] +local defaultEnvConcurrencyLimit = ARGV[6] + +-- Write the message to the message key +redis.call('SET', messageKey, messageData) + +-- Add the message to the queue +redis.call('ZADD', queueKey, messageScore, messageId) + +-- Add the message to the env queue +redis.call('ZADD', envQueueKey, messageScore, messageId) + +-- Rebalance the parent queue +local earliestMessage = redis.call('ZRANGE', queueKey, 0, 0, 'WITHSCORES') +if #earliestMessage == 0 then + redis.call('ZREM', parentQueueKey, queueName) +else + redis.call('ZADD', parentQueueKey, earliestMessage[2], queueName) +end + +-- Update the concurrency keys +redis.call('SREM', queueCurrentConcurrencyKey, messageId) +redis.call('SREM', envCurrentConcurrencyKey, messageId) +redis.call('SREM', envReserveConcurrencyKey, messageId) +redis.call('SREM', queueReserveConcurrencyKey, messageId) + +-- Reserve the concurrency for the message +local envReserveConcurrencyLimit = tonumber(redis.call('GET', envConcurrencyLimitKey) or defaultEnvConcurrencyLimit) +-- Count the number of messages in the reserve concurrency set +local envReserveConcurrency = tonumber(redis.call('SCARD', envReserveConcurrencyKey) or '0') + +-- If there is space, add the messaageId to the env reserve concurrency set +if envReserveConcurrency < envReserveConcurrencyLimit then + redis.call('SADD', envReserveConcurrencyKey, reserveMessageId) +end + +return true + `, + }); + + this.redis.defineCommand("enqueueMessageWithReservingConcurrencyOnRecursiveQueue", { + numberOfKeys: 10, + lua: ` +local queueKey = KEYS[1] +local parentQueueKey = KEYS[2] +local messageKey = KEYS[3] +local queueCurrentConcurrencyKey = KEYS[4] +local queueReserveConcurrencyKey = KEYS[5] +local queueConcurrencyLimitKey = KEYS[6] +local envCurrentConcurrencyKey = KEYS[7] +local envReserveConcurrencyKey = KEYS[8] +local envConcurrencyLimitKey = KEYS[9] +local envQueueKey = KEYS[10] + +local queueName = ARGV[1] +local messageId = ARGV[2] +local messageData = ARGV[3] +local messageScore = ARGV[4] +local reserveMessageId = ARGV[5] +local defaultEnvConcurrencyLimit = ARGV[6] + +-- Get the env reserve concurrency limit because we need it to calculate the max reserve concurrency +-- for the specific queue +local envReserveConcurrencyLimit = tonumber(redis.call('GET', envConcurrencyLimitKey) or defaultEnvConcurrencyLimit) + +-- Count the number of messages in the queue reserve concurrency set +local queueReserveConcurrency = tonumber(redis.call('SCARD', queueReserveConcurrencyKey) or '0') +local queueConcurrencyLimit = tonumber(redis.call('GET', queueConcurrencyLimitKey) or '1000000') + +local queueReserveConcurrencyLimit = math.min(queueConcurrencyLimit, envReserveConcurrencyLimit) + +-- If we cannot add the reserve concurrency, then we have to return false +if queueReserveConcurrency >= queueReserveConcurrencyLimit then + return false +end + +-- Write the message to the message key +redis.call('SET', messageKey, messageData) + +-- Add the message to the queue +redis.call('ZADD', queueKey, messageScore, messageId) + +-- Add the message to the env queue +redis.call('ZADD', envQueueKey, messageScore, messageId) + +-- Rebalance the parent queue +local earliestMessage = redis.call('ZRANGE', queueKey, 0, 0, 'WITHSCORES') +if #earliestMessage == 0 then + redis.call('ZREM', parentQueueKey, queueName) +else + redis.call('ZADD', parentQueueKey, earliestMessage[2], queueName) +end + +-- Update the concurrency keys +redis.call('SREM', queueCurrentConcurrencyKey, messageId) +redis.call('SREM', envCurrentConcurrencyKey, messageId) +redis.call('SREM', envReserveConcurrencyKey, messageId) +redis.call('SREM', queueReserveConcurrencyKey, messageId) + +-- Count the number of messages in the env reserve concurrency set +local envReserveConcurrency = tonumber(redis.call('SCARD', envReserveConcurrencyKey) or '0') + +-- If there is space, add the messaageId to the env reserve concurrency set +if envReserveConcurrency < envReserveConcurrencyLimit then + redis.call('SADD', envReserveConcurrencyKey, reserveMessageId) +end + +redis.call('SADD', queueReserveConcurrencyKey, reserveMessageId) + +return true `, }); this.redis.defineCommand("dequeueMessage", { numberOfKeys: 9, lua: ` --- Keys: childQueue, parentQueue, concurrencyLimitKey, envConcurrencyLimitKey, orgConcurrencyLimitKey, currentConcurrencyKey, envCurrentConcurrencyKey, orgCurrentConcurrencyKey -local childQueue = KEYS[1] -local parentQueue = KEYS[2] -local concurrencyLimitKey = KEYS[3] +local queueKey = KEYS[1] +local parentQueueKey = KEYS[2] +local queueConcurrencyLimitKey = KEYS[3] local envConcurrencyLimitKey = KEYS[4] -local orgConcurrencyLimitKey = KEYS[5] -local currentConcurrencyKey = KEYS[6] +local queueCurrentConcurrencyKey = KEYS[5] +local queueReserveConcurrencyKey = KEYS[6] local envCurrentConcurrencyKey = KEYS[7] -local orgCurrentConcurrencyKey = KEYS[8] +local envReserveConcurrencyKey = KEYS[8] local envQueueKey = KEYS[9] --- Args: childQueueName, currentTime, defaultEnvConcurrencyLimit, defaultOrgConcurrencyLimit -local childQueueName = ARGV[1] +local queueName = ARGV[1] local currentTime = tonumber(ARGV[2]) local defaultEnvConcurrencyLimit = ARGV[3] -local defaultOrgConcurrencyLimit = ARGV[4] - --- Check current org concurrency against the limit -local orgCurrentConcurrency = tonumber(redis.call('SCARD', orgCurrentConcurrencyKey) or '0') -local orgConcurrencyLimit = tonumber(redis.call('GET', orgConcurrencyLimitKey) or defaultOrgConcurrencyLimit) - -if orgCurrentConcurrency >= orgConcurrencyLimit then - return nil -end -- Check current env concurrency against the limit local envCurrentConcurrency = tonumber(redis.call('SCARD', envCurrentConcurrencyKey) or '0') local envConcurrencyLimit = tonumber(redis.call('GET', envConcurrencyLimitKey) or defaultEnvConcurrencyLimit) +local envReserveConcurrency = tonumber(redis.call('SCARD', envReserveConcurrencyKey) or '0') +local totalEnvConcurrencyLimit = envConcurrencyLimit + envReserveConcurrency -if envCurrentConcurrency >= envConcurrencyLimit then +if envCurrentConcurrency >= totalEnvConcurrencyLimit then return nil end -- Check current queue concurrency against the limit -local currentConcurrency = tonumber(redis.call('SCARD', currentConcurrencyKey) or '0') -local concurrencyLimit = tonumber(redis.call('GET', concurrencyLimitKey) or envConcurrencyLimit) +local queueCurrentConcurrency = tonumber(redis.call('SCARD', queueCurrentConcurrencyKey) or '0') +local queueConcurrencyLimit = math.min(tonumber(redis.call('GET', queueConcurrencyLimitKey) or '1000000'), envConcurrencyLimit) +local queueReserveConcurrency = tonumber(redis.call('SCARD', queueReserveConcurrencyKey) or '0') +local totalQueueConcurrencyLimit = queueConcurrencyLimit + queueReserveConcurrency -- Check condition only if concurrencyLimit exists -if currentConcurrency >= concurrencyLimit then +if queueCurrentConcurrency >= totalQueueConcurrencyLimit then return nil end -- Attempt to dequeue the next message -local messages = redis.call('ZRANGEBYSCORE', childQueue, '-inf', currentTime, 'WITHSCORES', 'LIMIT', 0, 1) +local messages = redis.call('ZRANGEBYSCORE', queueKey, '-inf', currentTime, 'WITHSCORES', 'LIMIT', 0, 1) if #messages == 0 then return nil @@ -1284,19 +1508,24 @@ end local messageId = messages[1] local messageScore = tonumber(messages[2]) --- Move message to timeout queue and update concurrency -redis.call('ZREM', childQueue, messageId) +-- Remove the message from the queue and update concurrency +redis.call('ZREM', queueKey, messageId) redis.call('ZREM', envQueueKey, messageId) -redis.call('SADD', currentConcurrencyKey, messageId) +redis.call('SADD', queueCurrentConcurrencyKey, messageId) redis.call('SADD', envCurrentConcurrencyKey, messageId) -redis.call('SADD', orgCurrentConcurrencyKey, messageId) + +-- Remove the message from the reserve concurrency set +redis.call('SREM', envReserveConcurrencyKey, messageId) + +-- Remove the message from the queue reserve concurrency set +redis.call('SREM', queueReserveConcurrencyKey, messageId) -- Rebalance the parent queue -local earliestMessage = redis.call('ZRANGE', childQueue, 0, 0, 'WITHSCORES') +local earliestMessage = redis.call('ZRANGE', queueKey, 0, 0, 'WITHSCORES') if #earliestMessage == 0 then - redis.call('ZREM', parentQueue, childQueueName) + redis.call('ZREM', parentQueueKey, queueName) else - redis.call('ZADD', parentQueue, earliestMessage[2], childQueueName) + redis.call('ZADD', parentQueueKey, earliestMessage[2], queueName) end return {messageId, messageScore} -- Return message details @@ -1323,70 +1552,63 @@ redis.call('SET', messageKey, messageData, 'GET') }); this.redis.defineCommand("acknowledgeMessage", { - numberOfKeys: 7, + numberOfKeys: 8, lua: ` --- Keys: parentQueue, messageKey, messageQueue, concurrencyKey, envCurrentConcurrencyKey, orgCurrentConcurrencyKey -local parentQueue = KEYS[1] +local parentQueueKey = KEYS[1] local messageKey = KEYS[2] -local messageQueue = KEYS[3] -local concurrencyKey = KEYS[4] -local envCurrentConcurrencyKey = KEYS[5] -local orgCurrentConcurrencyKey = KEYS[6] -local envQueueKey = KEYS[7] +local queueKey = KEYS[3] +local queueConcurrencyKey = KEYS[4] +local queueReserveConcurrencyKey = KEYS[5] +local envCurrentConcurrencyKey = KEYS[6] +local envReserveConcurrencyKey = KEYS[7] +local envQueueKey = KEYS[8] --- Args: messageId, messageQueueName local messageId = ARGV[1] -local messageQueueName = ARGV[2] - --- Remove the message from the message key -redis.call('DEL', messageKey) +local queueName = ARGV[2] -- Remove the message from the queue -redis.call('ZREM', messageQueue, messageId) - --- Remove the message from the env queue -redis.call('ZREM', envQueueKey, messageId) +redis.call('ZREM', queueKey, messageId) -- Rebalance the parent queue -local earliestMessage = redis.call('ZRANGE', messageQueue, 0, 0, 'WITHSCORES') +local earliestMessage = redis.call('ZRANGE', queueKey, 0, 0, 'WITHSCORES') if #earliestMessage == 0 then - redis.call('ZREM', parentQueue, messageQueueName) + redis.call('ZREM', parentQueueKey, queueName) else - redis.call('ZADD', parentQueue, earliestMessage[2], messageQueueName) + redis.call('ZADD', parentQueueKey, earliestMessage[2], queueName) end -- Update the concurrency keys -redis.call('SREM', concurrencyKey, messageId) +redis.call('SREM', queueConcurrencyKey, messageId) redis.call('SREM', envCurrentConcurrencyKey, messageId) -redis.call('SREM', orgCurrentConcurrencyKey, messageId) +redis.call('SREM', envReserveConcurrencyKey, messageId) +redis.call('SREM', queueReserveConcurrencyKey, messageId) +redis.call('ZREM', envQueueKey, messageId) +redis.call('DEL', messageKey) `, }); this.redis.defineCommand("nackMessage", { - numberOfKeys: 8, + numberOfKeys: 7, lua: ` local messageKey = KEYS[1] -local childQueueKey = KEYS[2] +local queueKey = KEYS[2] local parentQueueKey = KEYS[3] -local concurrencyKey = KEYS[4] -local envConcurrencyKey = KEYS[5] -local orgConcurrencyKey = KEYS[6] -local envQueueKey = KEYS[7] -local nackCounterKey = KEYS[8] +local queueCurrentConcurrencyKey = KEYS[4] +local envCurrentConcurrencyKey = KEYS[5] +local envQueueKey = KEYS[6] +local nackCounterKey = KEYS[7] --- Args: childQueueName, messageId, currentTime, messageScore -local childQueueName = ARGV[1] +local queueName = ARGV[1] local messageId = ARGV[2] local currentTime = tonumber(ARGV[3]) local messageScore = tonumber(ARGV[4]) --- Update the concurrency keys -redis.call('SREM', concurrencyKey, messageId) -redis.call('SREM', envConcurrencyKey, messageId) -redis.call('SREM', orgConcurrencyKey, messageId) +-- Update the current concurrency keys +redis.call('SREM', queueCurrentConcurrencyKey, messageId) +redis.call('SREM', envCurrentConcurrencyKey, messageId) -- Enqueue the message into the queue -redis.call('ZADD', childQueueKey, messageScore, messageId) +redis.call('ZADD', queueKey, messageScore, messageId) -- Enqueue the message into the env queue redis.call('ZADD', envQueueKey, messageScore, messageId) @@ -1396,64 +1618,22 @@ redis.call('INCR', nackCounterKey) redis.call('EXPIRE', nackCounterKey, 2592000) -- Rebalance the parent queue -local earliestMessage = redis.call('ZRANGE', childQueueKey, 0, 0, 'WITHSCORES') +local earliestMessage = redis.call('ZRANGE', queueKey, 0, 0, 'WITHSCORES') if #earliestMessage == 0 then - redis.call('ZREM', parentQueueKey, childQueueName) + redis.call('ZREM', parentQueueKey, queueName) else - redis.call('ZADD', parentQueueKey, earliestMessage[2], childQueueName) + redis.call('ZADD', parentQueueKey, earliestMessage[2], queueName) end `, }); - this.redis.defineCommand("releaseConcurrency", { - numberOfKeys: 3, - lua: ` -local concurrencyKey = KEYS[1] -local envCurrentConcurrencyKey = KEYS[2] -local orgCurrentConcurrencyKey = KEYS[3] - -local messageId = ARGV[1] - --- Update the concurrency keys -if concurrencyKey ~= "" then - redis.call('SREM', concurrencyKey, messageId) -end -redis.call('SREM', envCurrentConcurrencyKey, messageId) -redis.call('SREM', orgCurrentConcurrencyKey, messageId) -`, - }); - - this.redis.defineCommand("calculateQueueCurrentConcurrencies", { - numberOfKeys: 3, - lua: ` --- Keys: currentConcurrencyKey, currentEnvConcurrencyKey, currentOrgConcurrencyKey -local currentConcurrencyKey = KEYS[1] -local currentEnvConcurrencyKey = KEYS[2] -local currentOrgConcurrencyKey = KEYS[3] - -local currentOrgConcurrency = tonumber(redis.call('SCARD', currentOrgConcurrencyKey) or '0') - -local currentEnvConcurrency = tonumber(redis.call('SCARD', currentEnvConcurrencyKey) or '0') - -local currentConcurrency = tonumber(redis.call('SCARD', currentConcurrencyKey) or '0') - -return { currentOrgConcurrency, currentEnvConcurrency, currentConcurrency } - `, - }); - this.redis.defineCommand("updateGlobalConcurrencyLimits", { - numberOfKeys: 2, + numberOfKeys: 1, lua: ` --- Keys: envConcurrencyLimitKey, orgConcurrencyLimitKey local envConcurrencyLimitKey = KEYS[1] -local orgConcurrencyLimitKey = KEYS[2] - --- Args: envConcurrencyLimit, orgConcurrencyLimit local envConcurrencyLimit = ARGV[1] -local orgConcurrencyLimit = ARGV[2] redis.call('SET', envConcurrencyLimitKey, envConcurrencyLimit) -redis.call('SET', orgConcurrencyLimitKey, orgConcurrencyLimit) `, }); @@ -1493,34 +1673,73 @@ end declare module "ioredis" { interface RedisCommander { enqueueMessage( - queue: string, - parentQueue: string, + queueKey: string, + parentQueueKey: string, messageKey: string, - concurrencyKey: string, - envConcurrencyKey: string, - orgConcurrencyKey: string, - envQueue: string, + queueCurrentConcurrencyKey: string, + queueReserveConcurrencyKey: string, + envCurrentConcurrencyKey: string, + envReserveConcurrencyKey: string, + envQueueKey: string, queueName: string, messageId: string, messageData: string, messageScore: string, - callback?: Callback - ): Result; + callback?: Callback + ): Result; + + enqueueMessageWithReservingConcurrency( + queueKey: string, + parentQueueKey: string, + messageKey: string, + queueCurrentConcurrencyKey: string, + queueReserveConcurrencyKey: string, + envCurrentConcurrencyKey: string, + envReserveConcurrencyKey: string, + envConcurrencyLimitKey: string, + envQueueKey: string, + queueName: string, + messageId: string, + messageData: string, + messageScore: string, + reserveMessageId: string, + defaultEnvConcurrencyLimit: string, + callback?: Callback + ): Result; + + enqueueMessageWithReservingConcurrencyOnRecursiveQueue( + queueKey: string, + parentQueueKey: string, + messageKey: string, + queueCurrentConcurrencyKey: string, + queueReserveConcurrencyKey: string, + queueConcurrencyLimitKey: string, + envCurrentConcurrencyKey: string, + envReserveConcurrencyKey: string, + envConcurrencyLimitKey: string, + envQueueKey: string, + queueName: string, + messageId: string, + messageData: string, + messageScore: string, + reserveMessageId: string, + defaultEnvConcurrencyLimit: string, + callback?: Callback + ): Result; dequeueMessage( - childQueue: string, - parentQueue: string, - concurrencyLimitKey: string, + queueKey: string, + parentQueueKey: string, + queueConcurrencyLimitKey: string, envConcurrencyLimitKey: string, - orgConcurrencyLimitKey: string, - currentConcurrencyKey: string, + queueCurrentConcurrencyKey: string, + queueReserveConcurrencyKey: string, envCurrentConcurrencyKey: string, - orgCurrentConcurrencyKey: string, + envReserveConcurrencyKey: string, envQueueKey: string, - childQueueName: string, + queueName: string, currentTime: string, defaultEnvConcurrencyLimit: string, - defaultOrgConcurrencyLimit: string, callback?: Callback<[string, string]> ): Result<[string, string] | null, Context>; @@ -1535,8 +1754,9 @@ declare module "ioredis" { messageKey: string, messageQueue: string, concurrencyKey: string, + queueReserveConcurrencyKey: string, envConcurrencyKey: string, - orgConcurrencyKey: string, + envReserveConcurrencyKey: string, envQueueKey: string, messageId: string, messageQueueName: string, @@ -1545,33 +1765,22 @@ declare module "ioredis" { nackMessage( messageKey: string, - childQueueKey: string, + queueKey: string, parentQueueKey: string, - concurrencyKey: string, - envConcurrencyKey: string, - orgConcurrencyKey: string, + queueCurrentConcurrencyKey: string, + envCurrentConcurrencyKey: string, envQueueKey: string, nackCounterKey: string, - childQueueName: string, + queueName: string, messageId: string, currentTime: string, messageScore: string, callback?: Callback ): Result; - releaseConcurrency( - concurrencyKey: string, - envConcurrencyKey: string, - orgConcurrencyKey: string, - messageId: string, - callback?: Callback - ): Result; - updateGlobalConcurrencyLimits( envConcurrencyLimitKey: string, - orgConcurrencyLimitKey: string, envConcurrencyLimit: string, - orgConcurrencyLimit: string, callback?: Callback ): Result; @@ -1582,80 +1791,75 @@ declare module "ioredis" { currentScore: string, callback?: Callback ): Result; - - calculateQueueCurrentConcurrencies( - currentConcurrencyKey: string, - currentEnvConcurrencyKey: string, - currentOrgConcurrencyKey: string, - callback?: Callback - ): Result; } } export const marqs = singleton("marqs", getMarQSClient); function getMarQSClient() { - if (env.V3_ENABLED) { - if (env.REDIS_HOST && env.REDIS_PORT) { - const redisOptions = { - keyPrefix: KEY_PREFIX, - port: env.REDIS_PORT, - host: env.REDIS_HOST, - username: env.REDIS_USERNAME, - password: env.REDIS_PASSWORD, - enableAutoPipelining: true, - ...(env.REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }), - }; - - const redis = new Redis(redisOptions); - const keysProducer = new MarQSShortKeyProducer(KEY_PREFIX); - - return new MarQS({ - name: "marqs", - tracer: trace.getTracer("marqs"), - keysProducer, - visibilityTimeoutStrategy: new V3LegacyRunEngineWorkerVisibilityTimeout(), - queuePriorityStrategy: new FairDequeuingStrategy({ - tracer: tracer, - redis, - parentQueueLimit: env.MARQS_SHARED_QUEUE_LIMIT, - keys: keysProducer, - defaultEnvConcurrency: env.DEFAULT_ENV_EXECUTION_CONCURRENCY_LIMIT, - defaultOrgConcurrency: env.DEFAULT_ORG_EXECUTION_CONCURRENCY_LIMIT, - biases: { - concurrencyLimitBias: env.MARQS_CONCURRENCY_LIMIT_BIAS, - availableCapacityBias: env.MARQS_AVAILABLE_CAPACITY_BIAS, - queueAgeRandomization: env.MARQS_QUEUE_AGE_RANDOMIZATION_BIAS, - }, - reuseSnapshotCount: env.MARQS_REUSE_SNAPSHOT_COUNT, - maximumOrgCount: env.MARQS_MAXIMUM_ORG_COUNT, - }), - envQueuePriorityStrategy: new FairDequeuingStrategy({ - tracer: tracer, - redis, - parentQueueLimit: env.MARQS_DEV_QUEUE_LIMIT, - keys: keysProducer, - defaultEnvConcurrency: env.DEFAULT_ENV_EXECUTION_CONCURRENCY_LIMIT, - defaultOrgConcurrency: env.DEFAULT_ORG_EXECUTION_CONCURRENCY_LIMIT, - biases: { - concurrencyLimitBias: 0.0, - availableCapacityBias: 0.0, - queueAgeRandomization: 0.1, - }, - }), - workers: 1, + if (!env.REDIS_HOST || !env.REDIS_PORT) { + throw new Error( + "Could not initialize Trigger.dev because process.env.REDIS_HOST and process.env.REDIS_PORT are required to be set." + ); + } + + const redisOptions = { + keyPrefix: KEY_PREFIX, + port: env.REDIS_PORT, + host: env.REDIS_HOST, + username: env.REDIS_USERNAME, + password: env.REDIS_PASSWORD, + enableAutoPipelining: true, + ...(env.REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }), + }; + + const redis = new Redis(redisOptions); + const keysProducer = new MarQSShortKeyProducer(KEY_PREFIX); + + return new MarQS({ + name: "marqs", + tracer: trace.getTracer("marqs"), + keysProducer, + visibilityTimeoutStrategy: new V3LegacyRunEngineWorkerVisibilityTimeout(), + queuePriorityStrategy: new EnvPriorityDequeuingStrategy({ + keys: keysProducer, + delegate: new FairDequeuingStrategy({ + tracer: tracer, redis, + parentQueueLimit: env.MARQS_SHARED_QUEUE_LIMIT, + keys: keysProducer, defaultEnvConcurrency: env.DEFAULT_ENV_EXECUTION_CONCURRENCY_LIMIT, - defaultOrgConcurrency: env.DEFAULT_ORG_EXECUTION_CONCURRENCY_LIMIT, - visibilityTimeoutInMs: env.MARQS_VISIBILITY_TIMEOUT_MS, - enableRebalancing: !env.MARQS_DISABLE_REBALANCING, - maximumNackCount: env.MARQS_MAXIMUM_NACK_COUNT, - subscriber: concurrencyTracker, - }); - } else { - console.warn( - "Could not initialize MarQS because process.env.REDIS_HOST and process.env.REDIS_PORT are required to be set. Trigger.dev v3 will not work without this." - ); - } - } + biases: { + concurrencyLimitBias: env.MARQS_CONCURRENCY_LIMIT_BIAS, + availableCapacityBias: env.MARQS_AVAILABLE_CAPACITY_BIAS, + queueAgeRandomization: env.MARQS_QUEUE_AGE_RANDOMIZATION_BIAS, + }, + reuseSnapshotCount: env.MARQS_REUSE_SNAPSHOT_COUNT, + maximumEnvCount: env.MARQS_MAXIMUM_ENV_COUNT, + }), + }), + envQueuePriorityStrategy: new EnvPriorityDequeuingStrategy({ + keys: keysProducer, + delegate: new FairDequeuingStrategy({ + tracer: tracer, + redis, + parentQueueLimit: env.MARQS_DEV_QUEUE_LIMIT, + keys: keysProducer, + defaultEnvConcurrency: env.DEFAULT_ENV_EXECUTION_CONCURRENCY_LIMIT, + biases: { + concurrencyLimitBias: 0.0, + availableCapacityBias: 0.0, + queueAgeRandomization: 0.1, + }, + }), + }), + workers: 1, + redis, + defaultEnvConcurrency: env.DEFAULT_ENV_EXECUTION_CONCURRENCY_LIMIT, + defaultOrgConcurrency: env.DEFAULT_ORG_EXECUTION_CONCURRENCY_LIMIT, + visibilityTimeoutInMs: env.MARQS_VISIBILITY_TIMEOUT_MS, + enableRebalancing: !env.MARQS_DISABLE_REBALANCING, + maximumNackCount: env.MARQS_MAXIMUM_NACK_COUNT, + subscriber: concurrencyTracker, + }); } diff --git a/apps/webapp/app/v3/marqs/marqsKeyProducer.server.ts b/apps/webapp/app/v3/marqs/marqsKeyProducer.server.ts deleted file mode 100644 index 227a0a9dd2..0000000000 --- a/apps/webapp/app/v3/marqs/marqsKeyProducer.server.ts +++ /dev/null @@ -1,227 +0,0 @@ -import { AuthenticatedEnvironment } from "~/services/apiAuth.server"; -import { MarQSKeyProducer } from "./types"; - -const constants = { - SHARED_QUEUE: "sharedQueue", - CURRENT_CONCURRENCY_PART: "currentConcurrency", - CONCURRENCY_LIMIT_PART: "concurrency", - DISABLED_CONCURRENCY_LIMIT_PART: "disabledConcurrency", - ENV_PART: "env", - ORG_PART: "org", - QUEUE_PART: "queue", - CONCURRENCY_KEY_PART: "ck", - MESSAGE_PART: "message", -} as const; - -export class MarQSShortKeyProducer implements MarQSKeyProducer { - constructor(private _prefix: string) {} - - sharedQueueScanPattern() { - return `${this._prefix}*${constants.SHARED_QUEUE}`; - } - - queueCurrentConcurrencyScanPattern() { - return `${this._prefix}${constants.ORG_PART}:*:${constants.ENV_PART}:*:queue:*:${constants.CURRENT_CONCURRENCY_PART}`; - } - - stripKeyPrefix(key: string): string { - if (key.startsWith(this._prefix)) { - return key.slice(this._prefix.length); - } - - return key; - } - - queueConcurrencyLimitKey(env: AuthenticatedEnvironment, queue: string) { - return [this.queueKey(env, queue), constants.CONCURRENCY_LIMIT_PART].join(":"); - } - - envConcurrencyLimitKey(envId: string): string; - envConcurrencyLimitKey(env: AuthenticatedEnvironment): string; - envConcurrencyLimitKey(envOrId: AuthenticatedEnvironment | string): string { - return [ - this.envKeySection(typeof envOrId === "string" ? envOrId : envOrId.id), - constants.CONCURRENCY_LIMIT_PART, - ].join(":"); - } - - orgConcurrencyLimitKey(orgId: string): string; - orgConcurrencyLimitKey(env: AuthenticatedEnvironment): string; - orgConcurrencyLimitKey(envOrOrgId: AuthenticatedEnvironment | string) { - return [ - this.orgKeySection(typeof envOrOrgId === "string" ? envOrOrgId : envOrOrgId.organizationId), - constants.CONCURRENCY_LIMIT_PART, - ].join(":"); - } - - queueKey(orgId: string, envId: string, queue: string, concurrencyKey?: string): string; - queueKey(env: AuthenticatedEnvironment, queue: string, concurrencyKey?: string): string; - queueKey( - envOrOrgId: AuthenticatedEnvironment | string, - queueOrEnvId: string, - queueOrConcurrencyKey: string, - concurrencyKey?: string - ): string { - if (typeof envOrOrgId === "string") { - return [ - this.orgKeySection(envOrOrgId), - this.envKeySection(queueOrEnvId), - this.queueSection(queueOrConcurrencyKey), - ] - .concat(concurrencyKey ? this.concurrencyKeySection(concurrencyKey) : []) - .join(":"); - } else { - return [ - this.orgKeySection(envOrOrgId.organizationId), - this.envKeySection(envOrOrgId.id), - this.queueSection(queueOrEnvId), - ] - .concat(queueOrConcurrencyKey ? this.concurrencyKeySection(queueOrConcurrencyKey) : []) - .join(":"); - } - } - - envSharedQueueKey(env: AuthenticatedEnvironment) { - if (env.type === "DEVELOPMENT") { - return [ - this.orgKeySection(env.organizationId), - this.envKeySection(env.id), - constants.SHARED_QUEUE, - ].join(":"); - } - - return this.sharedQueueKey(); - } - - sharedQueueKey(): string { - return constants.SHARED_QUEUE; - } - - concurrencyLimitKeyFromQueue(queue: string) { - const concurrencyQueueName = queue.replace(/:ck:.+$/, ""); - - return `${concurrencyQueueName}:${constants.CONCURRENCY_LIMIT_PART}`; - } - - currentConcurrencyKeyFromQueue(queue: string) { - return `${queue}:${constants.CURRENT_CONCURRENCY_PART}`; - } - - currentConcurrencyKey( - env: AuthenticatedEnvironment, - queue: string, - concurrencyKey?: string - ): string { - return [this.queueKey(env, queue, concurrencyKey), constants.CURRENT_CONCURRENCY_PART].join( - ":" - ); - } - - disabledConcurrencyLimitKeyFromQueue(queue: string) { - const orgId = this.normalizeQueue(queue).split(":")[1]; - - return this.disabledConcurrencyLimitKey(orgId); - } - - disabledConcurrencyLimitKey(orgId: string) { - return `${constants.ORG_PART}:${orgId}:${constants.DISABLED_CONCURRENCY_LIMIT_PART}`; - } - - orgConcurrencyLimitKeyFromQueue(queue: string) { - const orgId = this.normalizeQueue(queue).split(":")[1]; - - return `${constants.ORG_PART}:${orgId}:${constants.CONCURRENCY_LIMIT_PART}`; - } - - orgCurrentConcurrencyKeyFromQueue(queue: string) { - const orgId = this.normalizeQueue(queue).split(":")[1]; - - return `${constants.ORG_PART}:${orgId}:${constants.CURRENT_CONCURRENCY_PART}`; - } - - envConcurrencyLimitKeyFromQueue(queue: string) { - const envId = this.normalizeQueue(queue).split(":")[3]; - - return `${constants.ENV_PART}:${envId}:${constants.CONCURRENCY_LIMIT_PART}`; - } - - envCurrentConcurrencyKeyFromQueue(queue: string) { - const envId = this.normalizeQueue(queue).split(":")[3]; - - return `${constants.ENV_PART}:${envId}:${constants.CURRENT_CONCURRENCY_PART}`; - } - - orgCurrentConcurrencyKey(orgId: string): string; - orgCurrentConcurrencyKey(env: AuthenticatedEnvironment): string; - orgCurrentConcurrencyKey(envOrOrgId: AuthenticatedEnvironment | string): string { - return [ - this.orgKeySection(typeof envOrOrgId === "string" ? envOrOrgId : envOrOrgId.organizationId), - constants.CURRENT_CONCURRENCY_PART, - ].join(":"); - } - - envCurrentConcurrencyKey(envId: string): string; - envCurrentConcurrencyKey(env: AuthenticatedEnvironment): string; - envCurrentConcurrencyKey(envOrId: AuthenticatedEnvironment | string): string { - return [ - this.envKeySection(typeof envOrId === "string" ? envOrId : envOrId.id), - constants.CURRENT_CONCURRENCY_PART, - ].join(":"); - } - - envQueueKeyFromQueue(queue: string) { - const envId = this.normalizeQueue(queue).split(":")[3]; - - return `${constants.ENV_PART}:${envId}:${constants.QUEUE_PART}`; - } - - envQueueKey(env: AuthenticatedEnvironment): string { - return [constants.ENV_PART, this.shortId(env.id), constants.QUEUE_PART].join(":"); - } - - messageKey(messageId: string) { - return `${constants.MESSAGE_PART}:${messageId}`; - } - - nackCounterKey(messageId: string): string { - return `${constants.MESSAGE_PART}:${messageId}:nacks`; - } - - orgIdFromQueue(queue: string) { - return this.normalizeQueue(queue).split(":")[1]; - } - - envIdFromQueue(queue: string) { - return this.normalizeQueue(queue).split(":")[3]; - } - - private shortId(id: string) { - // Return the last 12 characters of the id - return id.slice(-12); - } - - private envKeySection(envId: string) { - return `${constants.ENV_PART}:${this.shortId(envId)}`; - } - - private orgKeySection(orgId: string) { - return `${constants.ORG_PART}:${this.shortId(orgId)}`; - } - - private queueSection(queue: string) { - return `${constants.QUEUE_PART}:${queue}`; - } - - private concurrencyKeySection(concurrencyKey: string) { - return `${constants.CONCURRENCY_KEY_PART}:${concurrencyKey}`; - } - - // This removes the leading prefix from the queue name if it exists - private normalizeQueue(queue: string) { - if (queue.startsWith(this._prefix)) { - return queue.slice(this._prefix.length); - } - - return queue; - } -} diff --git a/apps/webapp/app/v3/marqs/marqsKeyProducer.ts b/apps/webapp/app/v3/marqs/marqsKeyProducer.ts new file mode 100644 index 0000000000..a4fbfb6a62 --- /dev/null +++ b/apps/webapp/app/v3/marqs/marqsKeyProducer.ts @@ -0,0 +1,312 @@ +import { MarQSKeyProducer, MarQSKeyProducerEnv, QueueDescriptor } from "./types"; + +const constants = { + SHARED_QUEUE: "sharedQueue", + CURRENT_CONCURRENCY_PART: "currentConcurrency", + CONCURRENCY_LIMIT_PART: "concurrency", + DISABLED_CONCURRENCY_LIMIT_PART: "disabledConcurrency", + ENV_PART: "env", + ORG_PART: "org", + QUEUE_PART: "queue", + CONCURRENCY_KEY_PART: "ck", + MESSAGE_PART: "message", + RESERVE_CONCURRENCY_PART: "reserveConcurrency", + PRIORITY_PART: "priority", +} as const; + +const ORG_REGEX = /org:([^:]+):/; +const ENV_REGEX = /env:([^:]+):/; +const QUEUE_REGEX = /queue:([^:]+)(?::|$)/; +const CONCURRENCY_KEY_REGEX = /ck:([^:]+)(?::|$)/; +const PRIORITY_REGEX = /priority:(\d+)(?::|$)/; + +export class MarQSShortKeyProducer implements MarQSKeyProducer { + constructor(private _prefix: string) {} + + sharedQueueScanPattern() { + return `${this._prefix}*${constants.SHARED_QUEUE}`; + } + + queueCurrentConcurrencyScanPattern() { + return `${this._prefix}${constants.ORG_PART}:*:${constants.ENV_PART}:*:queue:*:${constants.CURRENT_CONCURRENCY_PART}`; + } + + stripKeyPrefix(key: string): string { + if (key.startsWith(this._prefix)) { + return key.slice(this._prefix.length); + } + + return key; + } + + queueConcurrencyLimitKey(env: MarQSKeyProducerEnv, queue: string) { + return [this.queueKey(env, queue), constants.CONCURRENCY_LIMIT_PART].join(":"); + } + + envConcurrencyLimitKey(envId: string): string; + envConcurrencyLimitKey(env: MarQSKeyProducerEnv): string; + envConcurrencyLimitKey(envOrId: MarQSKeyProducerEnv | string): string { + return [ + this.envKeySection(typeof envOrId === "string" ? envOrId : envOrId.id), + constants.CONCURRENCY_LIMIT_PART, + ].join(":"); + } + + queueKey( + orgId: string, + envId: string, + queue: string, + concurrencyKey?: string, + priority?: number + ): string; + queueKey( + env: MarQSKeyProducerEnv, + queue: string, + concurrencyKey?: string, + priority?: number + ): string; + queueKey( + envOrOrgId: MarQSKeyProducerEnv | string, + queueOrEnvId: string, + queueOrConcurrencyKey: string, + concurrencyKeyOrPriority?: string | number, + priority?: number + ): string { + if (typeof envOrOrgId === "string") { + return [ + this.orgKeySection(envOrOrgId), + this.envKeySection(queueOrEnvId), + this.queueSection(queueOrConcurrencyKey), + ] + .concat( + typeof concurrencyKeyOrPriority === "string" + ? this.concurrencyKeySection(concurrencyKeyOrPriority) + : [] + ) + .concat(typeof priority === "number" && priority ? this.prioritySection(priority) : []) + .join(":"); + } else { + return [ + this.orgKeySection(envOrOrgId.organizationId), + this.envKeySection(envOrOrgId.id), + this.queueSection(queueOrEnvId), + ] + .concat(queueOrConcurrencyKey ? this.concurrencyKeySection(queueOrConcurrencyKey) : []) + .concat( + typeof concurrencyKeyOrPriority === "number" && concurrencyKeyOrPriority + ? this.prioritySection(concurrencyKeyOrPriority) + : [] + ) + .join(":"); + } + } + + queueKeyFromQueue(queue: string, priority?: number): string { + const descriptor = this.queueDescriptorFromQueue(queue); + + return this.queueKey( + descriptor.organization, + descriptor.environment, + descriptor.name, + descriptor.concurrencyKey, + descriptor.priority ?? priority + ); + } + + envSharedQueueKey(env: MarQSKeyProducerEnv) { + if (env.type === "DEVELOPMENT") { + return [ + this.orgKeySection(env.organizationId), + this.envKeySection(env.id), + constants.SHARED_QUEUE, + ].join(":"); + } + + return this.sharedQueueKey(); + } + + sharedQueueKey(): string { + return constants.SHARED_QUEUE; + } + + queueConcurrencyLimitKeyFromQueue(queue: string) { + const descriptor = this.queueDescriptorFromQueue(queue); + + return this.queueConcurrencyLimitKeyFromDescriptor(descriptor); + } + + queueCurrentConcurrencyKeyFromQueue(queue: string) { + const descriptor = this.queueDescriptorFromQueue(queue); + return this.currentConcurrencyKeyFromDescriptor(descriptor); + } + + queueReserveConcurrencyKeyFromQueue(queue: string) { + const descriptor = this.queueDescriptorFromQueue(queue); + + return this.queueReserveConcurrencyKeyFromDescriptor(descriptor); + } + + queueCurrentConcurrencyKey( + env: MarQSKeyProducerEnv, + queue: string, + concurrencyKey?: string + ): string { + return [this.queueKey(env, queue, concurrencyKey), constants.CURRENT_CONCURRENCY_PART].join( + ":" + ); + } + + envConcurrencyLimitKeyFromQueue(queue: string) { + const descriptor = this.queueDescriptorFromQueue(queue); + + return `${constants.ENV_PART}:${descriptor.environment}:${constants.CONCURRENCY_LIMIT_PART}`; + } + + envCurrentConcurrencyKeyFromQueue(queue: string) { + const descriptor = this.queueDescriptorFromQueue(queue); + + return `${constants.ENV_PART}:${descriptor.environment}:${constants.CURRENT_CONCURRENCY_PART}`; + } + + envReserveConcurrencyKeyFromQueue(queue: string) { + const descriptor = this.queueDescriptorFromQueue(queue); + + return this.envReserveConcurrencyKey(descriptor.environment); + } + + envReserveConcurrencyKey(envId: string): string { + return `${constants.ENV_PART}:${this.shortId(envId)}:${constants.RESERVE_CONCURRENCY_PART}`; + } + + envCurrentConcurrencyKey(envId: string): string; + envCurrentConcurrencyKey(env: MarQSKeyProducerEnv): string; + envCurrentConcurrencyKey(envOrId: MarQSKeyProducerEnv | string): string { + return [ + this.envKeySection(typeof envOrId === "string" ? envOrId : envOrId.id), + constants.CURRENT_CONCURRENCY_PART, + ].join(":"); + } + + envQueueKeyFromQueue(queue: string) { + const descriptor = this.queueDescriptorFromQueue(queue); + + return `${constants.ENV_PART}:${descriptor.environment}:${constants.QUEUE_PART}`; + } + + envQueueKey(env: MarQSKeyProducerEnv): string { + return [constants.ENV_PART, this.shortId(env.id), constants.QUEUE_PART].join(":"); + } + + messageKey(messageId: string) { + return `${constants.MESSAGE_PART}:${messageId}`; + } + + nackCounterKey(messageId: string): string { + return `${constants.MESSAGE_PART}:${messageId}:nacks`; + } + + orgIdFromQueue(queue: string) { + const descriptor = this.queueDescriptorFromQueue(queue); + + return descriptor.organization; + } + + envIdFromQueue(queue: string) { + const descriptor = this.queueDescriptorFromQueue(queue); + + return descriptor.environment; + } + + queueDescriptorFromQueue(queue: string): QueueDescriptor { + const match = queue.match(QUEUE_REGEX); + + if (!match) { + throw new Error(`Invalid queue: ${queue}, no queue name found`); + } + + const [, queueName] = match; + + const envMatch = queue.match(ENV_REGEX); + + if (!envMatch) { + throw new Error(`Invalid queue: ${queue}, no environment found`); + } + + const [, envId] = envMatch; + + const orgMatch = queue.match(ORG_REGEX); + + if (!orgMatch) { + throw new Error(`Invalid queue: ${queue}, no organization found`); + } + + const [, orgId] = orgMatch; + + const concurrencyKeyMatch = queue.match(CONCURRENCY_KEY_REGEX); + + const concurrencyKey = concurrencyKeyMatch ? concurrencyKeyMatch[1] : undefined; + + const priorityMatch = queue.match(PRIORITY_REGEX); + + const priority = priorityMatch ? parseInt(priorityMatch[1], 10) : undefined; + + return { + name: queueName, + environment: envId, + organization: orgId, + concurrencyKey, + priority, + }; + } + + private shortId(id: string) { + // Return the last 12 characters of the id + return id.slice(-12); + } + + private envKeySection(envId: string) { + return `${constants.ENV_PART}:${this.shortId(envId)}`; + } + + private orgKeySection(orgId: string) { + return `${constants.ORG_PART}:${this.shortId(orgId)}`; + } + + private queueSection(queue: string) { + return `${constants.QUEUE_PART}:${queue}`; + } + + private concurrencyKeySection(concurrencyKey: string) { + return `${constants.CONCURRENCY_KEY_PART}:${concurrencyKey}`; + } + + private prioritySection(priority: number) { + return `${constants.PRIORITY_PART}:${priority}`; + } + + private currentConcurrencyKeyFromDescriptor(descriptor: QueueDescriptor) { + return [ + this.queueKey( + descriptor.organization, + descriptor.environment, + descriptor.name, + descriptor.concurrencyKey + ), + constants.CURRENT_CONCURRENCY_PART, + ].join(":"); + } + + private queueReserveConcurrencyKeyFromDescriptor(descriptor: QueueDescriptor) { + return [ + this.queueKey(descriptor.organization, descriptor.environment, descriptor.name), + constants.RESERVE_CONCURRENCY_PART, + ].join(":"); + } + + private queueConcurrencyLimitKeyFromDescriptor(descriptor: QueueDescriptor) { + return [ + this.queueKey(descriptor.organization, descriptor.environment, descriptor.name), + constants.CONCURRENCY_LIMIT_PART, + ].join(":"); + } +} diff --git a/apps/webapp/app/v3/marqs/types.ts b/apps/webapp/app/v3/marqs/types.ts index e769a71206..339f1bded9 100644 --- a/apps/webapp/app/v3/marqs/types.ts +++ b/apps/webapp/app/v3/marqs/types.ts @@ -1,57 +1,85 @@ +import type { RuntimeEnvironmentType } from "@trigger.dev/database"; import { z } from "zod"; -import { type AuthenticatedEnvironment } from "~/services/apiAuth.server"; export type QueueRange = { offset: number; count: number }; +export type QueueDescriptor = { + organization: string; + environment: string; + name: string; + concurrencyKey?: string; + priority?: number; +}; + +export type MarQSKeyProducerEnv = { + id: string; + organizationId: string; + type: RuntimeEnvironmentType; +}; + export interface MarQSKeyProducer { - queueConcurrencyLimitKey(env: AuthenticatedEnvironment, queue: string): string; + queueConcurrencyLimitKey(env: MarQSKeyProducerEnv, queue: string): string; envConcurrencyLimitKey(envId: string): string; - envConcurrencyLimitKey(env: AuthenticatedEnvironment): string; + envConcurrencyLimitKey(env: MarQSKeyProducerEnv): string; - orgConcurrencyLimitKey(orgId: string): string; - orgConcurrencyLimitKey(env: AuthenticatedEnvironment): string; + envCurrentConcurrencyKey(envId: string): string; + envCurrentConcurrencyKey(env: MarQSKeyProducerEnv): string; - orgCurrentConcurrencyKey(orgId: string): string; - orgCurrentConcurrencyKey(env: AuthenticatedEnvironment): string; + envReserveConcurrencyKey(envId: string): string; - envCurrentConcurrencyKey(envId: string): string; - envCurrentConcurrencyKey(env: AuthenticatedEnvironment): string; + queueKey( + orgId: string, + envId: string, + queue: string, + concurrencyKey?: string, + priority?: number + ): string; + queueKey( + env: MarQSKeyProducerEnv, + queue: string, + concurrencyKey?: string, + priority?: number + ): string; - queueKey(orgId: string, envId: string, queue: string, concurrencyKey?: string): string; - queueKey(env: AuthenticatedEnvironment, queue: string, concurrencyKey?: string): string; + queueKeyFromQueue(queue: string, priority?: number): string; - envQueueKey(env: AuthenticatedEnvironment): string; - envSharedQueueKey(env: AuthenticatedEnvironment): string; + envQueueKey(env: MarQSKeyProducerEnv): string; + envSharedQueueKey(env: MarQSKeyProducerEnv): string; sharedQueueKey(): string; sharedQueueScanPattern(): string; queueCurrentConcurrencyScanPattern(): string; - concurrencyLimitKeyFromQueue(queue: string): string; - currentConcurrencyKeyFromQueue(queue: string): string; - currentConcurrencyKey( - env: AuthenticatedEnvironment, + queueConcurrencyLimitKeyFromQueue(queue: string): string; + queueCurrentConcurrencyKeyFromQueue(queue: string): string; + queueCurrentConcurrencyKey( + env: MarQSKeyProducerEnv, queue: string, concurrencyKey?: string ): string; - disabledConcurrencyLimitKey(orgId: string): string; - disabledConcurrencyLimitKeyFromQueue(queue: string): string; - orgConcurrencyLimitKeyFromQueue(queue: string): string; - orgCurrentConcurrencyKeyFromQueue(queue: string): string; envConcurrencyLimitKeyFromQueue(queue: string): string; envCurrentConcurrencyKeyFromQueue(queue: string): string; + envReserveConcurrencyKeyFromQueue(queue: string): string; envQueueKeyFromQueue(queue: string): string; messageKey(messageId: string): string; nackCounterKey(messageId: string): string; stripKeyPrefix(key: string): string; orgIdFromQueue(queue: string): string; envIdFromQueue(queue: string): string; + + queueReserveConcurrencyKeyFromQueue(queue: string): string; + queueDescriptorFromQueue(queue: string): QueueDescriptor; } +export type EnvQueues = { + envId: string; + queues: string[]; +}; + export interface MarQSFairDequeueStrategy { distributeFairQueuesFromParentQueue( parentQueue: string, consumerId: string - ): Promise>; + ): Promise>; } export const MessagePayload = z.object({ @@ -78,3 +106,8 @@ export interface VisibilityTimeoutStrategy { heartbeat(messageId: string, timeoutInMs: number): Promise; cancelHeartbeat(messageId: string): Promise; } + +export type EnqueueMessageReserveConcurrencyOptions = { + messageId: string; + recursiveQueue: boolean; +}; diff --git a/apps/webapp/app/v3/marqs/v2.server.ts b/apps/webapp/app/v3/marqs/v2.server.ts index fee98846b4..3492f42c57 100644 --- a/apps/webapp/app/v3/marqs/v2.server.ts +++ b/apps/webapp/app/v3/marqs/v2.server.ts @@ -9,7 +9,7 @@ import { PerformRunExecutionV3Service } from "~/services/runs/performRunExecutio import { singleton } from "~/utils/singleton"; import { generateFriendlyId } from "../friendlyIdentifiers"; import { MarQS } from "./index.server"; -import { MarQSShortKeyProducer } from "./marqsKeyProducer.server"; +import { MarQSShortKeyProducer } from "./marqsKeyProducer"; import { RequeueV2Message } from "./requeueV2Message.server"; import { VisibilityTimeoutStrategy } from "./types"; import Redis from "ioredis"; @@ -80,7 +80,6 @@ function getMarQSClient() { parentQueueLimit: 100, keys: new MarQSV2KeyProducer(KEY_PREFIX), defaultEnvConcurrency: env.V2_MARQS_DEFAULT_ENV_CONCURRENCY, - defaultOrgConcurrency: env.DEFAULT_ORG_EXECUTION_CONCURRENCY_LIMIT, }), envQueuePriorityStrategy: new NoopFairDequeuingStrategy(), // We don't use this in v2, since all queues go through the shared queue workers: 0, diff --git a/apps/webapp/app/v3/services/batchTriggerV3.server.ts b/apps/webapp/app/v3/services/batchTriggerV3.server.ts index fdbc8d9c64..4ffd59aacc 100644 --- a/apps/webapp/app/v3/services/batchTriggerV3.server.ts +++ b/apps/webapp/app/v3/services/batchTriggerV3.server.ts @@ -456,18 +456,17 @@ export class BatchTriggerV3Service extends BaseService { error: result.error, }); - await this.#enqueueBatchTaskRun({ - batchId: batch.id, - processingId: "0", - range: { - start: result.workingIndex, - count: PROCESSING_BATCH_SIZE, + await this._prisma.batchTaskRun.update({ + where: { + id: batch.id, + }, + data: { + status: "ABORTED", + completedAt: new Date(), }, - attemptCount: 0, - strategy: "sequential", }); - return batch; + throw result.error; } // Update the batch to be sealed diff --git a/apps/webapp/app/v3/services/completeAttempt.server.ts b/apps/webapp/app/v3/services/completeAttempt.server.ts index 235ae87352..ae5beb5bbc 100644 --- a/apps/webapp/app/v3/services/completeAttempt.server.ts +++ b/apps/webapp/app/v3/services/completeAttempt.server.ts @@ -22,7 +22,7 @@ import { env } from "~/env.server"; import { AuthenticatedEnvironment } from "~/services/apiAuth.server"; import { logger } from "~/services/logger.server"; import { safeJsonParse } from "~/utils/json"; -import { marqs } from "~/v3/marqs/index.server"; +import { marqs, MarQSPriorityLevel } from "~/v3/marqs/index.server"; import { createExceptionPropertiesFromError, eventRepository } from "../eventRepository.server"; import { FailedTaskRunRetryHelper } from "../failedTaskRun.server"; import { FAILED_RUN_STATUSES, isFinalAttemptStatus, isFinalRunStatus } from "../taskStatus"; @@ -476,7 +476,8 @@ export class CompleteAttemptService extends BaseService { checkpointEventId: this.opts.supportsRetryCheckpoints ? checkpointEventId : undefined, retryCheckpointsDisabled: !this.opts.supportsRetryCheckpoints, }, - executionRetry.timestamp + executionRetry.timestamp, + MarQSPriorityLevel.retry ); }; @@ -614,8 +615,13 @@ export class CompleteAttemptService extends BaseService { }); if (environment.type === "DEVELOPMENT") { - // This is already an EXECUTE message so we can just NACK - await marqs?.nackMessage(taskRunAttempt.taskRunId, executionRetry.timestamp); + marqs.replaceMessage( + taskRunAttempt.taskRunId, + {}, + executionRetry.timestamp, + MarQSPriorityLevel.retry + ); + return "RETRIED"; } diff --git a/apps/webapp/app/v3/services/createCheckpoint.server.ts b/apps/webapp/app/v3/services/createCheckpoint.server.ts index bbd8618898..15cdf697d9 100644 --- a/apps/webapp/app/v3/services/createCheckpoint.server.ts +++ b/apps/webapp/app/v3/services/createCheckpoint.server.ts @@ -2,7 +2,7 @@ import { CoordinatorToPlatformMessages, ManualCheckpointMetadata } from "@trigge import type { InferSocketMessageSchema } from "@trigger.dev/core/v3/zodSocket"; import type { Checkpoint, CheckpointRestoreEvent } from "@trigger.dev/database"; import { logger } from "~/services/logger.server"; -import { marqs } from "~/v3/marqs/index.server"; +import { marqs, MarQSPriorityLevel } from "~/v3/marqs/index.server"; import { generateFriendlyId } from "../friendlyIdentifiers"; import { isFreezableAttemptStatus, isFreezableRunStatus } from "../taskStatus"; import { BaseService } from "./baseService.server"; @@ -174,7 +174,8 @@ export class CreateCheckpointService extends BaseService { resumableAttemptId: attempt.id, checkpointEventId: checkpointEvent.id, }, - restoreAtUnixTimeMs + restoreAtUnixTimeMs, + MarQSPriorityLevel.resume ); return { @@ -302,6 +303,7 @@ export class CreateCheckpointService extends BaseService { checkpointEventId: checkpointEvent.id, }, undefined, + undefined, true ); diff --git a/apps/webapp/app/v3/services/enqueueDelayedRun.server.ts b/apps/webapp/app/v3/services/enqueueDelayedRun.server.ts index 6306df4765..fd0f22f3c2 100644 --- a/apps/webapp/app/v3/services/enqueueDelayedRun.server.ts +++ b/apps/webapp/app/v3/services/enqueueDelayedRun.server.ts @@ -6,6 +6,7 @@ import { BaseService } from "./baseService.server"; import { ExpireEnqueuedRunService } from "./expireEnqueuedRun.server"; import { commonWorker } from "../commonWorker.server"; import { workerQueue } from "~/services/worker.server"; +import { enqueueRun } from "./enqueueRun.server"; export class EnqueueDelayedRunService extends BaseService { public static async enqueue(runId: string, runAt?: Date) { @@ -44,6 +45,24 @@ export class EnqueueDelayedRunService extends BaseService { project: true, }, }, + dependency: { + include: { + dependentBatchRun: { + include: { + dependentTaskAttempt: { + include: { + taskRun: true, + }, + }, + }, + }, + dependentAttempt: { + include: { + taskRun: true, + }, + }, + }, + }, }, }); @@ -83,18 +102,12 @@ export class EnqueueDelayedRunService extends BaseService { } }); - await marqs?.enqueueMessage( - run.runtimeEnvironment, - run.queue, - run.id, - { - type: "EXECUTE", - taskIdentifier: run.taskIdentifier, - projectId: run.runtimeEnvironment.projectId, - environmentId: run.runtimeEnvironment.id, - environmentType: run.runtimeEnvironment.type, - }, - run.concurrencyKey ?? undefined - ); + await enqueueRun({ + env: run.runtimeEnvironment, + run: run, + dependentRun: + run.dependency?.dependentAttempt?.taskRun ?? + run.dependency?.dependentBatchRun?.dependentTaskAttempt?.taskRun, + }); } } diff --git a/apps/webapp/app/v3/services/enqueueRun.server.ts b/apps/webapp/app/v3/services/enqueueRun.server.ts new file mode 100644 index 0000000000..cb091d70d9 --- /dev/null +++ b/apps/webapp/app/v3/services/enqueueRun.server.ts @@ -0,0 +1,66 @@ +import { TaskRunError, TaskRunErrorCodes } from "@trigger.dev/core/v3/schemas"; +import { TaskRun } from "@trigger.dev/database"; +import { AuthenticatedEnvironment } from "~/services/apiAuth.server"; +import { marqs } from "../marqs/index.server"; + +export type EnqueueRunOptions = { + env: AuthenticatedEnvironment; + run: TaskRun; + dependentRun?: { queue: string; id: string }; +}; + +export type EnqueueRunResult = + | { + ok: true; + } + | { + ok: false; + error: TaskRunError; + }; + +export async function enqueueRun({ + env, + run, + dependentRun, +}: EnqueueRunOptions): Promise { + // If this is a triggerAndWait or batchTriggerAndWait, + // we need to add the parent run to the reserve concurrency set + // to free up concurrency for the children to run + // In the case of a recursive queue, reserving concurrency can fail, which means there is a deadlock and we need to fail the run + + // TODO: reserveConcurrency can fail because of a deadlock, we need to handle that case + const wasEnqueued = await marqs.enqueueMessage( + env, + run.queue, + run.id, + { + type: "EXECUTE", + taskIdentifier: run.taskIdentifier, + projectId: env.projectId, + environmentId: env.id, + environmentType: env.type, + }, + run.concurrencyKey ?? undefined, + run.queueTimestamp ?? undefined, + dependentRun + ? { messageId: dependentRun.id, recursiveQueue: dependentRun.queue === run.queue } + : undefined + ); + + if (!wasEnqueued) { + const error = { + type: "INTERNAL_ERROR", + code: TaskRunErrorCodes.RECURSIVE_WAIT_DEADLOCK, + message: `This run will never execute because it was triggered recursively and the task has no remaining concurrency available`, + } satisfies TaskRunError; + + return { + ok: false, + error, + }; + } + + return { + ok: true, + }; +} diff --git a/apps/webapp/app/v3/services/resumeBatchRun.server.ts b/apps/webapp/app/v3/services/resumeBatchRun.server.ts index 6a794e238c..68f5a59cf7 100644 --- a/apps/webapp/app/v3/services/resumeBatchRun.server.ts +++ b/apps/webapp/app/v3/services/resumeBatchRun.server.ts @@ -1,6 +1,6 @@ import { PrismaClientOrTransaction } from "~/db.server"; import { workerQueue } from "~/services/worker.server"; -import { marqs } from "~/v3/marqs/index.server"; +import { marqs, MarQSPriorityLevel } from "~/v3/marqs/index.server"; import { BaseService } from "./baseService.server"; import { logger } from "~/services/logger.server"; import { BatchTaskRun } from "@trigger.dev/database"; @@ -152,6 +152,8 @@ export class ResumeBatchRunService extends BaseService { queue: true, taskIdentifier: true, concurrencyKey: true, + createdAt: true, + queueTimestamp: true, }, }, }, @@ -186,6 +188,7 @@ export class ResumeBatchRunService extends BaseService { dependentTaskAttemptId: dependentTaskAttempt.id, }); + // TODO: use the new priority queue thingie await marqs?.enqueueMessage( environment, dependentRun.queue, @@ -200,7 +203,10 @@ export class ResumeBatchRunService extends BaseService { environmentId: environment.id, environmentType: environment.type, }, - dependentRun.concurrencyKey ?? undefined + dependentRun.concurrencyKey ?? undefined, + dependentRun.queueTimestamp ?? dependentRun.createdAt, + undefined, + MarQSPriorityLevel.resume ); return "COMPLETED"; @@ -246,16 +252,25 @@ export class ResumeBatchRunService extends BaseService { hasCheckpointEvent: !!batchRun.checkpointEventId, }); - await marqs?.replaceMessage(dependentRun.id, { - type: "RESUME", - completedAttemptIds: batchRun.items.map((item) => item.taskRunAttemptId).filter(Boolean), - resumableAttemptId: dependentTaskAttempt.id, - checkpointEventId: batchRun.checkpointEventId ?? undefined, - taskIdentifier: dependentTaskAttempt.taskRun.taskIdentifier, - projectId: environment.projectId, - environmentId: environment.id, - environmentType: environment.type, - }); + await marqs?.replaceMessage( + dependentRun.id, + { + type: "RESUME", + completedAttemptIds: batchRun.items + .map((item) => item.taskRunAttemptId) + .filter(Boolean), + resumableAttemptId: dependentTaskAttempt.id, + checkpointEventId: batchRun.checkpointEventId ?? undefined, + taskIdentifier: dependentTaskAttempt.taskRun.taskIdentifier, + projectId: environment.projectId, + environmentId: environment.id, + environmentType: environment.type, + }, + ( + dependentTaskAttempt.taskRun.queueTimestamp ?? dependentTaskAttempt.taskRun.createdAt + ).getTime(), + MarQSPriorityLevel.resume + ); return "COMPLETED"; } else { diff --git a/apps/webapp/app/v3/services/resumeTaskDependency.server.ts b/apps/webapp/app/v3/services/resumeTaskDependency.server.ts index bdbece9778..3151de18a0 100644 --- a/apps/webapp/app/v3/services/resumeTaskDependency.server.ts +++ b/apps/webapp/app/v3/services/resumeTaskDependency.server.ts @@ -1,6 +1,6 @@ import { PrismaClientOrTransaction } from "~/db.server"; import { workerQueue } from "~/services/worker.server"; -import { marqs } from "~/v3/marqs/index.server"; +import { MarQS, marqs, MarQSPriorityLevel } from "~/v3/marqs/index.server"; import { BaseService } from "./baseService.server"; import { logger } from "~/services/logger.server"; @@ -49,6 +49,8 @@ export class ResumeTaskDependencyService extends BaseService { runId: dependentRun.id, } ); + + // TODO: use the new priority queue thingie await marqs?.enqueueMessage( dependency.taskRun.runtimeEnvironment, dependentRun.queue, @@ -64,7 +66,9 @@ export class ResumeTaskDependencyService extends BaseService { environmentType: dependency.taskRun.runtimeEnvironment.type, }, dependentRun.concurrencyKey ?? undefined, - dependentRun.createdAt.getTime() + dependentRun.queueTimestamp ?? dependentRun.createdAt, + undefined, + MarQSPriorityLevel.resume ); } else { logger.debug("Task dependency resume: Attempt is not paused or there's no checkpoint event", { @@ -97,7 +101,8 @@ export class ResumeTaskDependencyService extends BaseService { environmentId: dependency.taskRun.runtimeEnvironment.id, environmentType: dependency.taskRun.runtimeEnvironment.type, }, - dependentRun.createdAt.getTime() + (dependentRun.queueTimestamp ?? dependentRun.createdAt).getTime(), + MarQSPriorityLevel.resume ); } } diff --git a/apps/webapp/app/v3/services/triggerTask.server.ts b/apps/webapp/app/v3/services/triggerTask.server.ts index b7df21cde5..91a201e9cf 100644 --- a/apps/webapp/app/v3/services/triggerTask.server.ts +++ b/apps/webapp/app/v3/services/triggerTask.server.ts @@ -1,34 +1,36 @@ import { IOPacket, + packetRequiresOffloading, QueueOptions, SemanticInternalAttributes, + taskRunErrorEnhancer, + taskRunErrorToString, TriggerTaskRequestBody, - packetRequiresOffloading, } from "@trigger.dev/core/v3"; +import { parseNaturalLanguageDuration } from "@trigger.dev/core/v3/apps"; +import { Prisma, TaskRun } from "@trigger.dev/database"; import { env } from "~/env.server"; +import { sanitizeQueueName } from "~/models/taskQueue.server"; +import { createTag, MAX_TAGS_PER_RUN } from "~/models/taskRunTag.server"; import { AuthenticatedEnvironment } from "~/services/apiAuth.server"; import { autoIncrementCounter } from "~/services/autoIncrementCounter.server"; -import { workerQueue } from "~/services/worker.server"; +import { logger } from "~/services/logger.server"; +import { getEntitlement } from "~/services/platform.v3.server"; +import { resolveIdempotencyKeyTTL } from "~/utils/idempotencyKeys.server"; +import { handleMetadataPacket } from "~/utils/packets"; import { marqs } from "~/v3/marqs/index.server"; import { eventRepository } from "../eventRepository.server"; import { generateFriendlyId } from "../friendlyIdentifiers"; -import { uploadPacketToObjectStore } from "../r2.server"; -import { startActiveSpan } from "../tracer.server"; -import { getEntitlement } from "~/services/platform.v3.server"; -import { BaseService, ServiceValidationError } from "./baseService.server"; -import { logger } from "~/services/logger.server"; -import { isFinalAttemptStatus, isFinalRunStatus } from "../taskStatus"; -import { createTag, MAX_TAGS_PER_RUN } from "~/models/taskRunTag.server"; import { findCurrentWorkerFromEnvironment } from "../models/workerDeployment.server"; -import { handleMetadataPacket } from "~/utils/packets"; -import { parseNaturalLanguageDuration } from "@trigger.dev/core/v3/apps"; -import { ExpireEnqueuedRunService } from "./expireEnqueuedRun.server"; import { guardQueueSizeLimitsForEnv } from "../queueSizeLimits.server"; +import { uploadPacketToObjectStore } from "../r2.server"; +import { isFinalAttemptStatus, isFinalRunStatus } from "../taskStatus"; +import { startActiveSpan } from "../tracer.server"; import { clampMaxDuration } from "../utils/maxDuration"; -import { resolveIdempotencyKeyTTL } from "~/utils/idempotencyKeys.server"; -import { Prisma, TaskRun } from "@trigger.dev/database"; -import { sanitizeQueueName } from "~/models/taskQueue.server"; +import { BaseService, ServiceValidationError } from "./baseService.server"; import { EnqueueDelayedRunService } from "./enqueueDelayedRun.server"; +import { enqueueRun } from "./enqueueRun.server"; +import { ExpireEnqueuedRunService } from "./expireEnqueuedRun.server"; import { getTaskEventStore } from "../taskEventStore.server"; export type TriggerTaskServiceOptions = { @@ -186,6 +188,8 @@ export class TriggerTaskService extends BaseService { taskIdentifier: true, rootTaskRunId: true, depth: true, + queueTimestamp: true, + queue: true, }, }, }, @@ -242,6 +246,8 @@ export class TriggerTaskService extends BaseService { taskIdentifier: true, rootTaskRunId: true, depth: true, + queueTimestamp: true, + queue: true, }, }, }, @@ -294,7 +300,7 @@ export class TriggerTaskService extends BaseService { : undefined; try { - return await eventRepository.traceEvent( + const result = await eventRepository.traceEvent( taskId, { context: options.traceContext, @@ -367,6 +373,12 @@ export class TriggerTaskService extends BaseService { ? dependentBatchRun.dependentTaskAttempt.taskRun.depth + 1 : 0; + const queueTimestamp = + dependentAttempt?.taskRun.queueTimestamp ?? + dependentBatchRun?.dependentTaskAttempt?.taskRun.queueTimestamp ?? + delayUntil ?? + new Date(); + const taskRun = await tx.taskRun.create({ data: { status: delayUntil ? "DELAYED" : "PENDING", @@ -394,6 +406,7 @@ export class TriggerTaskService extends BaseService { isTest: body.options?.test ?? false, delayUntil, queuedAt: delayUntil ? undefined : new Date(), + queueTimestamp, maxAttempts: body.options?.maxAttempts, taskEventStore: getTaskEventStore(), ttl, @@ -547,44 +560,61 @@ export class TriggerTaskService extends BaseService { this._prisma ); - //release the concurrency for the env and org, if part of a (batch)triggerAndWait - if (dependentAttempt) { - const isSameTask = dependentAttempt.taskRun.taskIdentifier === taskId; - await marqs?.releaseConcurrency(dependentAttempt.taskRun.id, isSameTask); - } - if (dependentBatchRun?.dependentTaskAttempt) { - const isSameTask = - dependentBatchRun.dependentTaskAttempt.taskRun.taskIdentifier === taskId; - await marqs?.releaseConcurrency( - dependentBatchRun.dependentTaskAttempt.taskRun.id, - isSameTask - ); - } - if (!run) { return; } - // We need to enqueue the task run into the appropriate queue. This is done after the tx completes to prevent a race condition where the task run hasn't been created yet by the time we dequeue. + // Now enqueue the run if it's not delayed if (run.status === "PENDING") { - await marqs?.enqueueMessage( - environment, - run.queue, - run.id, - { - type: "EXECUTE", - taskIdentifier: taskId, - projectId: environment.projectId, - environmentId: environment.id, - environmentType: environment.type, - }, - body.options?.concurrencyKey - ); + const enqueueResult = await enqueueRun({ + env: environment, + run, + dependentRun: + dependentAttempt?.taskRun ?? dependentBatchRun?.dependentTaskAttempt?.taskRun, + }); + + if (!enqueueResult.ok) { + // Now we need to fail the run with enqueueResult.error and make sure and + // set the traced event to failed as well + await this._prisma.taskRun.update({ + where: { id: run.id }, + data: { + status: "SYSTEM_FAILURE", + completedAt: new Date(), + error: enqueueResult.error, + }, + }); + + event.failWithError(enqueueResult.error); + + return { + run, + isCached: false, + error: enqueueResult.error, + }; + } } return { run, isCached: false }; } ); + + if (result?.error) { + throw new ServiceValidationError( + taskRunErrorToString(taskRunErrorEnhancer(result.error)) + ); + } + + const run = result?.run; + + if (!run) { + return; + } + + return { + run, + isCached: result?.isCached, + }; } catch (error) { // Detect a prisma transaction Unique constraint violation if (error instanceof Prisma.PrismaClientKnownRequestError) { diff --git a/apps/webapp/test/envPriorityDequeueingStrategy.test.ts b/apps/webapp/test/envPriorityDequeueingStrategy.test.ts new file mode 100644 index 0000000000..8335136923 --- /dev/null +++ b/apps/webapp/test/envPriorityDequeueingStrategy.test.ts @@ -0,0 +1,388 @@ +import { describe, expect, it } from "vitest"; +import type { EnvQueues, MarQSFairDequeueStrategy } from "~/v3/marqs/types.js"; +import { EnvPriorityDequeuingStrategy } from "../app/v3/marqs/envPriorityDequeuingStrategy.server.js"; +import { createKeyProducer } from "./utils/marqs.js"; + +const keyProducer = createKeyProducer("test"); + +describe("EnvPriorityDequeuingStrategy", () => { + class TestDelegate implements MarQSFairDequeueStrategy { + constructor(private queues: EnvQueues[]) {} + + async distributeFairQueuesFromParentQueue(): Promise> { + return this.queues; + } + } + + describe("distributeFairQueuesFromParentQueue", () => { + it("should preserve order when all queues have the same priority", async () => { + const inputQueues: EnvQueues[] = [ + { + envId: "env1", + queues: [ + "org:org1:env:env1:queue:queue1:priority:1", + "org:org1:env:env1:queue:queue2:priority:1", + "org:org1:env:env1:queue:queue3:priority:1", + ], + }, + ]; + + const delegate = new TestDelegate(inputQueues); + const strategy = new EnvPriorityDequeuingStrategy({ + delegate, + keys: keyProducer, + }); + + const result = await strategy.distributeFairQueuesFromParentQueue("parentQueue", "consumer1"); + + expect(result).toEqual(inputQueues); + expect(result[0].queues).toEqual(inputQueues[0].queues); + }); + + it("should sort queues by priority in descending order", async () => { + const inputQueues: EnvQueues[] = [ + { + envId: "env1", + queues: [ + "org:org1:env:env1:queue:queue1:priority:1", + "org:org1:env:env1:queue:queue2:priority:3", + "org:org1:env:env1:queue:queue3:priority:2", + ], + }, + ]; + + const delegate = new TestDelegate(inputQueues); + const strategy = new EnvPriorityDequeuingStrategy({ + delegate, + keys: keyProducer, + }); + + const result = await strategy.distributeFairQueuesFromParentQueue("parentQueue", "consumer1"); + + expect(result[0].queues).toEqual([ + "org:org1:env:env1:queue:queue2:priority:3", + "org:org1:env:env1:queue:queue3:priority:2", + "org:org1:env:env1:queue:queue1:priority:1", + ]); + }); + + it("should handle queues without priority by treating them as priority 0", async () => { + const inputQueues: EnvQueues[] = [ + { + envId: "env1", + queues: [ + "org:org1:env:env1:queue:queue1", + "org:org1:env:env1:queue:queue2:priority:2", + "org:org1:env:env1:queue:queue3", + ], + }, + ]; + + const delegate = new TestDelegate(inputQueues); + const strategy = new EnvPriorityDequeuingStrategy({ + delegate, + keys: keyProducer, + }); + + const result = await strategy.distributeFairQueuesFromParentQueue("parentQueue", "consumer1"); + + expect(result[0].queues).toEqual([ + "org:org1:env:env1:queue:queue2:priority:2", + "org:org1:env:env1:queue:queue1", + "org:org1:env:env1:queue:queue3", + ]); + }); + + it("should handle multiple environments", async () => { + const inputQueues: EnvQueues[] = [ + { + envId: "env1", + queues: [ + "org:org1:env:env1:queue:queue1:priority:1", + "org:org1:env:env1:queue:queue2:priority:2", + ], + }, + { + envId: "env2", + queues: [ + "org:org1:env:env2:queue:queue3:priority:3", + "org:org1:env:env2:queue:queue4:priority:1", + ], + }, + ]; + + const delegate = new TestDelegate(inputQueues); + const strategy = new EnvPriorityDequeuingStrategy({ + delegate, + keys: keyProducer, + }); + + const result = await strategy.distributeFairQueuesFromParentQueue("parentQueue", "consumer1"); + + expect(result).toHaveLength(2); + expect(result[0].queues).toEqual([ + "org:org1:env:env1:queue:queue2:priority:2", + "org:org1:env:env1:queue:queue1:priority:1", + ]); + expect(result[1].queues).toEqual([ + "org:org1:env:env2:queue:queue3:priority:3", + "org:org1:env:env2:queue:queue4:priority:1", + ]); + }); + + it("should handle negative priorities correctly", async () => { + const inputQueues: EnvQueues[] = [ + { + envId: "env1", + queues: [ + "org:org1:env:env1:queue:queue1:priority:-1", + "org:org1:env:env1:queue:queue2:priority:1", + "org:org1:env:env1:queue:queue3:priority:-2", + ], + }, + ]; + + const delegate = new TestDelegate(inputQueues); + const strategy = new EnvPriorityDequeuingStrategy({ + delegate, + keys: keyProducer, + }); + + const result = await strategy.distributeFairQueuesFromParentQueue("parentQueue", "consumer1"); + + expect(result[0].queues).toEqual([ + "org:org1:env:env1:queue:queue2:priority:1", + "org:org1:env:env1:queue:queue1:priority:-1", + "org:org1:env:env1:queue:queue3:priority:-2", + ]); + }); + + it("should maintain stable sort for mixed priority and non-priority queues", async () => { + const inputQueues: EnvQueues[] = [ + { + envId: "env1", + queues: [ + "org:org1:env:env1:queue:queue1", + "org:org1:env:env1:queue:queue2:priority:1", + "org:org1:env:env1:queue:queue3", + "org:org1:env:env1:queue:queue4:priority:1", + "org:org1:env:env1:queue:queue5", + ], + }, + ]; + + const delegate = new TestDelegate(inputQueues); + const strategy = new EnvPriorityDequeuingStrategy({ + delegate, + keys: keyProducer, + }); + + const result = await strategy.distributeFairQueuesFromParentQueue("parentQueue", "consumer1"); + + // Check that queue2 and queue4 (priority 1) maintain their relative order + // and queue1, queue3, and queue5 (priority 0) maintain their relative order + expect(result[0].queues).toEqual([ + "org:org1:env:env1:queue:queue2:priority:1", + "org:org1:env:env1:queue:queue4:priority:1", + "org:org1:env:env1:queue:queue1", + "org:org1:env:env1:queue:queue3", + "org:org1:env:env1:queue:queue5", + ]); + }); + + it("should handle empty queue arrays", async () => { + const inputQueues: EnvQueues[] = [ + { + envId: "env1", + queues: [], + }, + ]; + + const delegate = new TestDelegate(inputQueues); + const strategy = new EnvPriorityDequeuingStrategy({ + delegate, + keys: keyProducer, + }); + + const result = await strategy.distributeFairQueuesFromParentQueue("parentQueue", "consumer1"); + + expect(result).toEqual(inputQueues); + expect(result[0].queues).toEqual([]); + }); + + it("should handle empty environments array", async () => { + const inputQueues: EnvQueues[] = []; + + const delegate = new TestDelegate(inputQueues); + const strategy = new EnvPriorityDequeuingStrategy({ + delegate, + keys: keyProducer, + }); + + const result = await strategy.distributeFairQueuesFromParentQueue("parentQueue", "consumer1"); + + expect(result).toEqual([]); + }); + + it("should handle large priority differences", async () => { + const inputQueues: EnvQueues[] = [ + { + envId: "env1", + queues: [ + "org:org1:env:env1:queue:queue1:priority:1", + "org:org1:env:env1:queue:queue2:priority:1000", + "org:org1:env:env1:queue:queue3:priority:500", + ], + }, + ]; + + const delegate = new TestDelegate(inputQueues); + const strategy = new EnvPriorityDequeuingStrategy({ + delegate, + keys: keyProducer, + }); + + const result = await strategy.distributeFairQueuesFromParentQueue("parentQueue", "consumer1"); + + expect(result[0].queues).toEqual([ + "org:org1:env:env1:queue:queue2:priority:1000", + "org:org1:env:env1:queue:queue3:priority:500", + "org:org1:env:env1:queue:queue1:priority:1", + ]); + }); + + it("should handle multiple environments with mixed priority patterns", async () => { + const inputQueues: EnvQueues[] = [ + { + envId: "env1", + queues: [ + "org:org1:env:env1:queue:queue1", // priority 0 + "org:org1:env:env1:queue:queue2:priority:2", + ], + }, + { + envId: "env2", + queues: [ + "org:org1:env:env2:queue:queue3:priority:1", + "org:org1:env:env2:queue:queue4", // priority 0 + ], + }, + { + envId: "env3", + queues: [ + "org:org1:env:env3:queue:queue5:priority:1", + "org:org1:env:env3:queue:queue6:priority:1", + ], + }, + ]; + + const delegate = new TestDelegate(inputQueues); + const strategy = new EnvPriorityDequeuingStrategy({ + delegate, + keys: keyProducer, + }); + + const result = await strategy.distributeFairQueuesFromParentQueue("parentQueue", "consumer1"); + + expect(result).toHaveLength(3); + expect(result[0].queues).toEqual([ + "org:org1:env:env1:queue:queue2:priority:2", + "org:org1:env:env1:queue:queue1", + ]); + expect(result[1].queues).toEqual([ + "org:org1:env:env2:queue:queue3:priority:1", + "org:org1:env:env2:queue:queue4", + ]); + expect(result[2].queues).toEqual([ + "org:org1:env:env3:queue:queue5:priority:1", + "org:org1:env:env3:queue:queue6:priority:1", + ]); + }); + + it("should sort queues with concurrency keys while maintaining priority order", async () => { + const inputQueues: EnvQueues[] = [ + { + envId: "env1", + queues: [ + "org:org1:env:env1:queue:queue1:ck:key1:priority:1", + "org:org1:env:env1:queue:queue2:ck:key1:priority:3", + "org:org1:env:env1:queue:queue3:ck:key2:priority:2", + ], + }, + ]; + + const delegate = new TestDelegate(inputQueues); + const strategy = new EnvPriorityDequeuingStrategy({ + delegate, + keys: keyProducer, + }); + + const result = await strategy.distributeFairQueuesFromParentQueue("parentQueue", "consumer1"); + + expect(result[0].queues).toEqual([ + "org:org1:env:env1:queue:queue2:ck:key1:priority:3", + "org:org1:env:env1:queue:queue3:ck:key2:priority:2", + "org:org1:env:env1:queue:queue1:ck:key1:priority:1", + ]); + }); + + it("should handle mixed queues with and without concurrency keys", async () => { + const inputQueues: EnvQueues[] = [ + { + envId: "env1", + queues: [ + "org:org1:env:env1:queue:queue1:priority:1", + "org:org1:env:env1:queue:queue2:ck:shared-key:priority:2", + "org:org1:env:env1:queue:queue3:ck:shared-key:priority:1", + "org:org1:env:env1:queue:queue4:priority:3", + "org:org1:env:env1:queue:queue5:ck:other-key:priority:2", + ], + }, + ]; + + const delegate = new TestDelegate(inputQueues); + const strategy = new EnvPriorityDequeuingStrategy({ + delegate, + keys: keyProducer, + }); + + const result = await strategy.distributeFairQueuesFromParentQueue("parentQueue", "consumer1"); + + expect(result[0].queues).toEqual([ + "org:org1:env:env1:queue:queue4:priority:3", + "org:org1:env:env1:queue:queue2:ck:shared-key:priority:2", + "org:org1:env:env1:queue:queue5:ck:other-key:priority:2", + "org:org1:env:env1:queue:queue1:priority:1", + "org:org1:env:env1:queue:queue3:ck:shared-key:priority:1", + ]); + }); + + it("should only return the highest priority queue of the same queue", async () => { + const inputQueues: EnvQueues[] = [ + { + envId: "env1", + queues: [ + "org:org1:env:env1:queue:queue1", + "org:org1:env:env1:queue:queue1:priority:1", + "org:org1:env:env1:queue:queue1:priority:2", + "org:org1:env:env1:queue:queue1:priority:3", + "org:org1:env:env1:queue:queue2", + ], + }, + ]; + + const delegate = new TestDelegate(inputQueues); + const strategy = new EnvPriorityDequeuingStrategy({ + delegate, + keys: keyProducer, + }); + + const result = await strategy.distributeFairQueuesFromParentQueue("parentQueue", "consumer1"); + + expect(result[0].queues).toEqual([ + "org:org1:env:env1:queue:queue1:priority:3", + "org:org1:env:env1:queue:queue2", + ]); + }); + }); +}); diff --git a/apps/webapp/test/fairDequeuingStrategy.test.ts b/apps/webapp/test/fairDequeuingStrategy.test.ts index fa96d11bc4..5d0b8e4af8 100644 --- a/apps/webapp/test/fairDequeuingStrategy.test.ts +++ b/apps/webapp/test/fairDequeuingStrategy.test.ts @@ -8,25 +8,24 @@ import { setupQueue, } from "./utils/marqs.js"; import { trace } from "@opentelemetry/api"; +import { EnvQueues } from "~/v3/marqs/types.js"; const tracer = trace.getTracer("test"); vi.setConfig({ testTimeout: 30_000 }); // 30 seconds timeout describe("FairDequeuingStrategy", () => { - redisTest("should distribute a single queue from a single org/env", async ({ redis }) => { + redisTest("should distribute a single queue from a single env", async ({ redis }) => { const keyProducer = createKeyProducer("test"); const strategy = new FairDequeuingStrategy({ tracer, redis, keys: keyProducer, - defaultOrgConcurrency: 10, defaultEnvConcurrency: 5, parentQueueLimit: 100, seed: "test-seed-1", // for deterministic shuffling }); - // Setup a single queue await setupQueue({ redis, keyProducer, @@ -40,42 +39,10 @@ describe("FairDequeuingStrategy", () => { const result = await strategy.distributeFairQueuesFromParentQueue("parent-queue", "consumer-1"); expect(result).toHaveLength(1); - expect(result[0]).toBe("org:org-1:env:env-1:queue:queue-1"); - }); - - redisTest("should respect org concurrency limits", async ({ redis }) => { - const keyProducer = createKeyProducer("test"); - const strategy = new FairDequeuingStrategy({ - tracer, - redis, - keys: keyProducer, - defaultOrgConcurrency: 2, - defaultEnvConcurrency: 5, - parentQueueLimit: 100, - seed: "test-seed-2", - }); - - // Setup queue - await setupQueue({ - redis, - keyProducer, - parentQueue: "parent-queue", - score: Date.now() - 1000, - queueId: "queue-1", - orgId: "org-1", + expect(result[0]).toEqual({ envId: "env-1", + queues: ["org:org-1:env:env-1:queue:queue-1"], }); - - // Set org-1 to be at its concurrency limit - await setupConcurrency({ - redis, - keyProducer, - org: { id: "org-1", currentConcurrency: 2, limit: 2 }, - env: { id: "env-1", currentConcurrency: 0 }, - }); - - const result = await strategy.distributeFairQueuesFromParentQueue("parent-queue", "consumer-1"); - expect(result).toHaveLength(0); }); redisTest("should respect env concurrency limits", async ({ redis }) => { @@ -84,7 +51,6 @@ describe("FairDequeuingStrategy", () => { tracer, redis, keys: keyProducer, - defaultOrgConcurrency: 10, defaultEnvConcurrency: 2, parentQueueLimit: 100, seed: "test-seed-3", @@ -103,7 +69,6 @@ describe("FairDequeuingStrategy", () => { await setupConcurrency({ redis, keyProducer, - org: { id: "org-1", currentConcurrency: 0 }, env: { id: "env-1", currentConcurrency: 2, limit: 2 }, }); @@ -111,13 +76,53 @@ describe("FairDequeuingStrategy", () => { expect(result).toHaveLength(0); }); + redisTest( + "should give extra concurrency when the env has reserve concurrency", + async ({ redis }) => { + const keyProducer = createKeyProducer("test"); + const strategy = new FairDequeuingStrategy({ + tracer, + redis, + keys: keyProducer, + defaultEnvConcurrency: 2, + parentQueueLimit: 100, + seed: "test-seed-3", + }); + + await setupQueue({ + redis, + keyProducer, + parentQueue: "parent-queue", + score: Date.now() - 1000, + queueId: "queue-1", + orgId: "org-1", + envId: "env-1", + }); + + await setupConcurrency({ + redis, + keyProducer, + env: { id: "env-1", currentConcurrency: 2, limit: 2, reserveConcurrency: 1 }, + }); + + const result = await strategy.distributeFairQueuesFromParentQueue( + "parent-queue", + "consumer-1" + ); + expect(result).toHaveLength(1); + expect(result[0]).toEqual({ + envId: "env-1", + queues: ["org:org-1:env:env-1:queue:queue-1"], + }); + } + ); + redisTest("should respect parentQueueLimit", async ({ redis }) => { const keyProducer = createKeyProducer("test"); const strategy = new FairDequeuingStrategy({ tracer, redis, keys: keyProducer, - defaultOrgConcurrency: 10, defaultEnvConcurrency: 5, parentQueueLimit: 2, // Only take 2 queues seed: "test-seed-6", @@ -158,11 +163,13 @@ describe("FairDequeuingStrategy", () => { const result = await strategy.distributeFairQueuesFromParentQueue("parent-queue", "consumer-1"); - expect(result).toHaveLength(2); - // Should only get the two oldest queues + expect(result).toHaveLength(1); const queue1 = keyProducer.queueKey("org-1", "env-1", "queue-1"); const queue2 = keyProducer.queueKey("org-1", "env-1", "queue-2"); - expect(result).toEqual([queue1, queue2]); + expect(result[0]).toEqual({ + envId: "env-1", + queues: [queue1, queue2], + }); }); redisTest("should reuse snapshots across calls for the same consumer", async ({ redis }) => { @@ -171,7 +178,6 @@ describe("FairDequeuingStrategy", () => { tracer, redis, keys: keyProducer, - defaultOrgConcurrency: 10, defaultEnvConcurrency: 5, parentQueueLimit: 10, seed: "test-seed-reuse-1", @@ -212,7 +218,11 @@ describe("FairDequeuingStrategy", () => { const startDistribute1 = performance.now(); - const result = await strategy.distributeFairQueuesFromParentQueue("parent-queue", "consumer-1"); + const envResult = await strategy.distributeFairQueuesFromParentQueue( + "parent-queue", + "consumer-1" + ); + const result = flattenResults(envResult); const distribute1Duration = performance.now() - startDistribute1; @@ -236,8 +246,8 @@ describe("FairDequeuingStrategy", () => { console.log("Second distribution took", distribute2Duration, "ms"); - // Make sure the second call is more than 10 times faster than the first - expect(distribute2Duration).toBeLessThan(distribute1Duration / 10); + // Make sure the second call is more than 9 times faster than the first + expect(distribute2Duration).toBeLessThan(distribute1Duration / 9); const startDistribute3 = performance.now(); @@ -260,7 +270,6 @@ describe("FairDequeuingStrategy", () => { tracer, redis, keys: keyProducer, - defaultOrgConcurrency: 10, defaultEnvConcurrency: 5, parentQueueLimit: 100, seed: "test-seed-5", @@ -296,7 +305,6 @@ describe("FairDequeuingStrategy", () => { await setupConcurrency({ redis, keyProducer, - org: { id: orgId, currentConcurrency: 2, limit: 10 }, env: { id: envId, currentConcurrency: 1, limit: 5 }, }); } @@ -323,10 +331,11 @@ describe("FairDequeuingStrategy", () => { // Run multiple iterations for (let i = 0; i < iterations; i++) { - const result = await strategy.distributeFairQueuesFromParentQueue( + const envResult = await strategy.distributeFairQueuesFromParentQueue( "parent-queue", `consumer-${i % 3}` // Simulate 3 different consumers ); + const result = flattenResults(envResult); // Track positions of queues result.forEach((queueId, position) => { @@ -417,7 +426,6 @@ describe("FairDequeuingStrategy", () => { tracer, redis, keys: keyProducer, - defaultOrgConcurrency: 10, defaultEnvConcurrency: 5, parentQueueLimit: 100, seed: "fixed-seed", @@ -472,20 +480,19 @@ describe("FairDequeuingStrategy", () => { await setupConcurrency({ redis, keyProducer, - org: { id: "org-1", currentConcurrency: 0, limit: 10 }, env: { id: "env-1", currentConcurrency: 0, limit: 5 }, }); await setupConcurrency({ redis, keyProducer, - org: { id: "org-1", currentConcurrency: 0, limit: 10 }, env: { id: "env-2", currentConcurrency: 0, limit: 5 }, }); - const result = await strategy.distributeFairQueuesFromParentQueue( + const envResult = await strategy.distributeFairQueuesFromParentQueue( "parent-queue", "consumer-1" ); + const result = flattenResults(envResult); // Group queues by environment const queuesByEnv = result.reduce((acc, queueId) => { @@ -546,7 +553,6 @@ describe("FairDequeuingStrategy", () => { await setupConcurrency({ redis, keyProducer, - org: { id: "org-1", currentConcurrency: 0, limit: 200 }, env: { id: setup.envId, currentConcurrency: setup.current, @@ -576,7 +582,6 @@ describe("FairDequeuingStrategy", () => { tracer, redis, keys: keyProducer, - defaultOrgConcurrency: 10, defaultEnvConcurrency: 5, parentQueueLimit: 100, seed: `test-seed-${i}`, @@ -596,10 +601,11 @@ describe("FairDequeuingStrategy", () => { const firstPositionCounts: Record = {}; for (let i = 0; i < iterationsPerStrategy; i++) { - const result = await strategy.distributeFairQueuesFromParentQueue( + const envResult = await strategy.distributeFairQueuesFromParentQueue( "parent-queue", `consumer-${i % 3}` ); + const result = flattenResults(envResult); expect(result.length).toBeGreaterThan(0); @@ -660,7 +666,6 @@ describe("FairDequeuingStrategy", () => { tracer, redis, keys: keyProducer, - defaultOrgConcurrency: 10, defaultEnvConcurrency: 5, parentQueueLimit: 100, seed: "fixed-seed", @@ -679,10 +684,11 @@ describe("FairDequeuingStrategy", () => { const iterations = 1000; for (let i = 0; i < iterations; i++) { - const result = await strategy.distributeFairQueuesFromParentQueue( + const envResult = await strategy.distributeFairQueuesFromParentQueue( "parent-queue", "consumer-1" ); + const result = flattenResults(envResult); result.forEach((queueId, position) => { const baseQueueId = queueId.split(":").pop()!; @@ -709,7 +715,6 @@ describe("FairDequeuingStrategy", () => { await setupConcurrency({ redis, keyProducer, - org: { id: "org-1", currentConcurrency: 0, limit: 10 }, env: { id: "env-1", currentConcurrency: 0, limit: 5 }, }); @@ -738,46 +743,45 @@ describe("FairDequeuingStrategy", () => { }); redisTest( - "should respect maximumOrgCount and select orgs based on queue ages", + "should respect maximumEnvCount and select envs based on queue ages", async ({ redis }) => { const keyProducer = createKeyProducer("test"); const strategy = new FairDequeuingStrategy({ tracer, redis, keys: keyProducer, - defaultOrgConcurrency: 10, defaultEnvConcurrency: 5, parentQueueLimit: 100, seed: "test-seed-max-orgs", - maximumOrgCount: 2, // Only select top 2 orgs + maximumEnvCount: 2, // Only select top 2 orgs }); const now = Date.now(); - // Setup 4 orgs with different queue age profiles - const orgSetups = [ + // Setup 4 envs with different queue age profiles + const envSetups = [ { - orgId: "org-1", + envId: "env-1", queues: [ { age: 1000 }, // Average age: 1000 ], }, { - orgId: "org-2", + envId: "env-2", queues: [ { age: 5000 }, // Average age: 5000 { age: 5000 }, ], }, { - orgId: "org-3", + envId: "env-3", queues: [ { age: 2000 }, // Average age: 2000 { age: 2000 }, ], }, { - orgId: "org-4", + envId: "env-4", queues: [ { age: 500 }, // Average age: 500 { age: 500 }, @@ -786,12 +790,11 @@ describe("FairDequeuingStrategy", () => { ]; // Setup queues and concurrency for each org - for (const setup of orgSetups) { + for (const setup of envSetups) { await setupConcurrency({ redis, keyProducer, - org: { id: setup.orgId, currentConcurrency: 0, limit: 10 }, - env: { id: "env-1", currentConcurrency: 0, limit: 5 }, + env: { id: setup.envId, currentConcurrency: 0, limit: 5 }, }); for (let i = 0; i < setup.queues.length; i++) { @@ -800,56 +803,57 @@ describe("FairDequeuingStrategy", () => { keyProducer, parentQueue: "parent-queue", score: now - setup.queues[i].age, - queueId: `queue-${setup.orgId}-${i}`, - orgId: setup.orgId, - envId: "env-1", + queueId: `queue-${setup.envId}-${i}`, + orgId: `org-${setup.envId}`, + envId: setup.envId, }); } } // Run multiple iterations to verify consistent behavior const iterations = 100; - const selectedOrgCounts: Record = {}; + const selectedEnvCounts: Record = {}; for (let i = 0; i < iterations; i++) { - const result = await strategy.distributeFairQueuesFromParentQueue( + const envResult = await strategy.distributeFairQueuesFromParentQueue( "parent-queue", `consumer-${i}` ); + const result = flattenResults(envResult); // Track which orgs were included in the result - const selectedOrgs = new Set(result.map((queueId) => keyProducer.orgIdFromQueue(queueId))); + const selectedEnvs = new Set(result.map((queueId) => keyProducer.envIdFromQueue(queueId))); // Verify we never get more than maximumOrgCount orgs - expect(selectedOrgs.size).toBeLessThanOrEqual(2); + expect(selectedEnvs.size).toBeLessThanOrEqual(2); - for (const orgId of selectedOrgs) { - selectedOrgCounts[orgId] = (selectedOrgCounts[orgId] || 0) + 1; + for (const envId of selectedEnvs) { + selectedEnvCounts[envId] = (selectedEnvCounts[envId] || 0) + 1; } } - console.log("Organization selection counts:", selectedOrgCounts); + console.log("Environment selection counts:", selectedEnvCounts); // org-2 should be selected most often (highest average age) - expect(selectedOrgCounts["org-2"]).toBeGreaterThan(selectedOrgCounts["org-4"] || 0); + expect(selectedEnvCounts["env-2"]).toBeGreaterThan(selectedEnvCounts["env-4"] || 0); // org-4 should be selected least often (lowest average age) - const org4Count = selectedOrgCounts["org-4"] || 0; - expect(org4Count).toBeLessThan(selectedOrgCounts["org-2"]); + const env4Count = selectedEnvCounts["env-4"] || 0; + expect(env4Count).toBeLessThan(selectedEnvCounts["env-2"]); - // Verify that orgs with higher average queue age are selected more frequently - const sortedOrgs = Object.entries(selectedOrgCounts).sort((a, b) => b[1] - a[1]); - console.log("Sorted organization frequencies:", sortedOrgs); + // Verify that envs with higher average queue age are selected more frequently + const sortedEnvs = Object.entries(selectedEnvCounts).sort((a, b) => b[1] - a[1]); + console.log("Sorted environment frequencies:", sortedEnvs); - // The top 2 most frequently selected orgs should be org-2 and org-3 + // The top 2 most frequently selected orgs should be env-2 and env-3 // as they have the highest average queue ages - const topTwoOrgs = new Set([sortedOrgs[0][0], sortedOrgs[1][0]]); - expect(topTwoOrgs).toContain("org-2"); // Highest average age - expect(topTwoOrgs).toContain("org-3"); // Second highest average age + const topTwoEnvs = new Set([sortedEnvs[0][0], sortedEnvs[1][0]]); + expect(topTwoEnvs).toContain("env-2"); // Highest average age + expect(topTwoEnvs).toContain("env-3"); // Second highest average age // Calculate selection percentages - const totalSelections = Object.values(selectedOrgCounts).reduce((a, b) => a + b, 0); - const selectionPercentages = Object.entries(selectedOrgCounts).reduce( + const totalSelections = Object.values(selectedEnvCounts).reduce((a, b) => a + b, 0); + const selectionPercentages = Object.entries(selectedEnvCounts).reduce( (acc, [orgId, count]) => { acc[orgId] = (count / totalSelections) * 100; return acc; @@ -857,13 +861,18 @@ describe("FairDequeuingStrategy", () => { {} as Record ); - console.log("Organization selection percentages:", selectionPercentages); + console.log("Environment selection percentages:", selectionPercentages); - // Verify that org-2 (highest average age) gets selected in at least 40% of iterations - expect(selectionPercentages["org-2"]).toBeGreaterThan(40); + // Verify that env-2 (highest average age) gets selected in at least 40% of iterations + expect(selectionPercentages["env-2"]).toBeGreaterThan(40); - // Verify that org-4 (lowest average age) gets selected in less than 20% of iterations - expect(selectionPercentages["org-4"] || 0).toBeLessThan(20); + // Verify that env-4 (lowest average age) gets selected in less than 20% of iterations + expect(selectionPercentages["env-4"] || 0).toBeLessThan(20); } ); }); + +// Helper function to flatten results for counting +function flattenResults(results: Array): string[] { + return results.flatMap((envQueue) => envQueue.queues); +} diff --git a/apps/webapp/test/marqsKeyProducer.test.ts b/apps/webapp/test/marqsKeyProducer.test.ts new file mode 100644 index 0000000000..1f6af24545 --- /dev/null +++ b/apps/webapp/test/marqsKeyProducer.test.ts @@ -0,0 +1,315 @@ +import { describe, it, expect } from "vitest"; +import { MarQSShortKeyProducer } from "../app/v3/marqs/marqsKeyProducer.js"; +import { MarQSKeyProducerEnv } from "~/v3/marqs/types.js"; + +describe("MarQSShortKeyProducer", () => { + const prefix = "test:"; + const producer = new MarQSShortKeyProducer(prefix); + + // Sample test data + const sampleEnv: MarQSKeyProducerEnv = { + id: "123456789012345678901234", + organizationId: "987654321098765432109876", + type: "PRODUCTION", + }; + + const devEnv: MarQSKeyProducerEnv = { + id: "123456789012345678901234", + organizationId: "987654321098765432109876", + type: "DEVELOPMENT", + }; + + describe("sharedQueueScanPattern", () => { + it("should return correct shared queue scan pattern", () => { + expect(producer.sharedQueueScanPattern()).toBe("test:*sharedQueue"); + }); + }); + + describe("queueCurrentConcurrencyScanPattern", () => { + it("should return correct queue current concurrency scan pattern", () => { + expect(producer.queueCurrentConcurrencyScanPattern()).toBe( + "test:org:*:env:*:queue:*:currentConcurrency" + ); + }); + }); + + describe("stripKeyPrefix", () => { + it("should strip prefix from key if present", () => { + expect(producer.stripKeyPrefix("test:someKey")).toBe("someKey"); + }); + + it("should return original key if prefix not present", () => { + expect(producer.stripKeyPrefix("someKey")).toBe("someKey"); + }); + }); + + describe("queueKey", () => { + it("should generate queue key with environment object", () => { + expect(producer.queueKey(sampleEnv, "testQueue")).toBe( + "org:765432109876:env:345678901234:queue:testQueue" + ); + }); + + it("should generate queue key with separate parameters", () => { + expect(producer.queueKey("org123", "env456", "testQueue")).toBe( + "org:org123:env:env456:queue:testQueue" + ); + }); + + it("should include concurrency key when provided", () => { + expect(producer.queueKey(sampleEnv, "testQueue", "concKey")).toBe( + "org:765432109876:env:345678901234:queue:testQueue:ck:concKey" + ); + }); + + it("should include priority when provided", () => { + expect(producer.queueKey(sampleEnv, "testQueue", undefined, 1)).toBe( + "org:765432109876:env:345678901234:queue:testQueue:priority:1" + ); + }); + + it("should NOT include priority when provided with 0", () => { + expect(producer.queueKey(sampleEnv, "testQueue", undefined, 0)).toBe( + "org:765432109876:env:345678901234:queue:testQueue" + ); + }); + + it("should include priority when provided with overloaded call", () => { + expect( + producer.queueKey(sampleEnv.organizationId, sampleEnv.id, "testQueue", undefined, 1) + ).toBe("org:765432109876:env:345678901234:queue:testQueue:priority:1"); + }); + }); + + describe("queueKeyFromQueue", () => { + it("should generate queue key", () => { + expect(producer.queueKeyFromQueue("org:765432109876:env:345678901234:queue:testQueue")).toBe( + "org:765432109876:env:345678901234:queue:testQueue" + ); + }); + + it("should include concurrency key when provided", () => { + expect( + producer.queueKeyFromQueue("org:765432109876:env:345678901234:queue:testQueue:ck:concKey") + ).toBe("org:765432109876:env:345678901234:queue:testQueue:ck:concKey"); + }); + + it("should include priority when provided", () => { + expect( + producer.queueKeyFromQueue("org:765432109876:env:345678901234:queue:testQueue", 1) + ).toBe("org:765432109876:env:345678901234:queue:testQueue:priority:1"); + }); + + it("should NOT include priority when provided with 0", () => { + expect( + producer.queueKeyFromQueue("org:765432109876:env:345678901234:queue:testQueue", 0) + ).toBe("org:765432109876:env:345678901234:queue:testQueue"); + }); + + it("should NOT change the priority when provided", () => { + expect( + producer.queueKeyFromQueue( + "org:765432109876:env:345678901234:queue:testQueue:priority:1", + 10 + ) + ).toBe("org:765432109876:env:345678901234:queue:testQueue:priority:1"); + }); + }); + + describe("envSharedQueueKey", () => { + it("should return organization-specific shared queue for development environment", () => { + expect(producer.envSharedQueueKey(devEnv)).toBe( + "org:765432109876:env:345678901234:sharedQueue" + ); + }); + + it("should return global shared queue for production environment", () => { + expect(producer.envSharedQueueKey(sampleEnv)).toBe("sharedQueue"); + }); + }); + + describe("queueDescriptorFromQueue", () => { + it("should parse queue string into descriptor", () => { + const queueString = "org:123:env:456:queue:testQueue:ck:concKey:priority:5"; + const descriptor = producer.queueDescriptorFromQueue(queueString); + + expect(descriptor).toEqual({ + name: "testQueue", + environment: "456", + organization: "123", + concurrencyKey: "concKey", + priority: 5, + }); + }); + + it("should parse queue string without optional parameters", () => { + const queueString = "org:123:env:456:queue:testQueue"; + const descriptor = producer.queueDescriptorFromQueue(queueString); + + expect(descriptor).toEqual({ + name: "testQueue", + environment: "456", + organization: "123", + concurrencyKey: undefined, + priority: undefined, + }); + }); + + it("should throw error for invalid queue string", () => { + const invalidQueue = "invalid:queue:string"; + expect(() => producer.queueDescriptorFromQueue(invalidQueue)).toThrow("Invalid queue"); + }); + }); + + describe("messageKey", () => { + it("should generate correct message key", () => { + expect(producer.messageKey("msg123")).toBe("message:msg123"); + }); + }); + + describe("nackCounterKey", () => { + it("should generate correct nack counter key", () => { + expect(producer.nackCounterKey("msg123")).toBe("message:msg123:nacks"); + }); + }); + + describe("currentConcurrencyKey", () => { + it("should generate correct current concurrency key", () => { + expect(producer.queueCurrentConcurrencyKey(sampleEnv, "testQueue")).toBe( + "org:765432109876:env:345678901234:queue:testQueue:currentConcurrency" + ); + }); + + it("should include concurrency key when provided", () => { + expect(producer.queueCurrentConcurrencyKey(sampleEnv, "testQueue", "concKey")).toBe( + "org:765432109876:env:345678901234:queue:testQueue:ck:concKey:currentConcurrency" + ); + }); + }); + + describe("currentConcurrencyKeyFromQueue", () => { + it("should generate correct current concurrency key", () => { + expect( + producer.queueCurrentConcurrencyKeyFromQueue( + "org:765432109876:env:345678901234:queue:testQueue" + ) + ).toBe("org:765432109876:env:345678901234:queue:testQueue:currentConcurrency"); + }); + + it("should include concurrency key when provided", () => { + expect( + producer.queueCurrentConcurrencyKeyFromQueue( + "org:765432109876:env:345678901234:queue:testQueue:ck:concKey" + ) + ).toBe("org:765432109876:env:345678901234:queue:testQueue:ck:concKey:currentConcurrency"); + }); + + it("should remove the priority bit when provided", () => { + expect( + producer.queueCurrentConcurrencyKeyFromQueue( + "org:765432109876:env:345678901234:queue:testQueue:priority:1" + ) + ).toBe("org:765432109876:env:345678901234:queue:testQueue:currentConcurrency"); + }); + + it("should remove the priority bit when provided, but keep the concurrency key", () => { + expect( + producer.queueCurrentConcurrencyKeyFromQueue( + "org:765432109876:env:345678901234:queue:testQueue:ck:concKey:priority:1" + ) + ).toBe("org:765432109876:env:345678901234:queue:testQueue:ck:concKey:currentConcurrency"); + }); + }); + + describe("queueReserveConcurrencyKeyFromQueue", () => { + it("should generate correct queue reserve concurrency key", () => { + expect( + producer.queueReserveConcurrencyKeyFromQueue( + "org:765432109876:env:345678901234:queue:testQueue" + ) + ).toBe("org:765432109876:env:345678901234:queue:testQueue:reserveConcurrency"); + }); + + it("should NOT include the concurrency key when provided", () => { + expect( + producer.queueReserveConcurrencyKeyFromQueue( + "org:765432109876:env:345678901234:queue:testQueue:ck:concKey" + ) + ).toBe("org:765432109876:env:345678901234:queue:testQueue:reserveConcurrency"); + }); + + it("should remove the priority bit when provided", () => { + expect( + producer.queueReserveConcurrencyKeyFromQueue( + "org:765432109876:env:345678901234:queue:testQueue:priority:1" + ) + ).toBe("org:765432109876:env:345678901234:queue:testQueue:reserveConcurrency"); + }); + + it("should remove the priority bit when provided, AND remove the concurrency key", () => { + expect( + producer.queueReserveConcurrencyKeyFromQueue( + "org:765432109876:env:345678901234:queue:testQueue:ck:concKey:priority:1" + ) + ).toBe("org:765432109876:env:345678901234:queue:testQueue:reserveConcurrency"); + }); + }); + + describe("queueConcurrencyLimitKeyFromQueue", () => { + it("should generate correct queue concurrency limit key", () => { + expect( + producer.queueConcurrencyLimitKeyFromQueue( + "org:765432109876:env:345678901234:queue:testQueue" + ) + ).toBe("org:765432109876:env:345678901234:queue:testQueue:concurrency"); + }); + + it("should NOT include the concurrency key when provided", () => { + expect( + producer.queueConcurrencyLimitKeyFromQueue( + "org:765432109876:env:345678901234:queue:testQueue:ck:concKey" + ) + ).toBe("org:765432109876:env:345678901234:queue:testQueue:concurrency"); + }); + + it("should remove the priority bit when provided", () => { + expect( + producer.queueConcurrencyLimitKeyFromQueue( + "org:765432109876:env:345678901234:queue:testQueue:priority:1" + ) + ).toBe("org:765432109876:env:345678901234:queue:testQueue:concurrency"); + }); + + it("should remove the priority bit when provided, AND remove the concurrency key", () => { + expect( + producer.queueConcurrencyLimitKeyFromQueue( + "org:765432109876:env:345678901234:queue:testQueue:ck:concKey:priority:1" + ) + ).toBe("org:765432109876:env:345678901234:queue:testQueue:concurrency"); + }); + }); + + describe("envCurrentConcurrencyKey", () => { + it("should generate correct env current concurrency key with environment object", () => { + expect(producer.envCurrentConcurrencyKey(sampleEnv)).toBe( + "env:345678901234:currentConcurrency" + ); + }); + + it("should generate correct env current concurrency key with env id", () => { + expect(producer.envCurrentConcurrencyKey("env456")).toBe("env:env456:currentConcurrency"); + }); + }); + + describe("orgIdFromQueue and envIdFromQueue", () => { + it("should extract org id from queue string", () => { + const queue = "org:123:env:456:queue:testQueue"; + expect(producer.orgIdFromQueue(queue)).toBe("123"); + }); + + it("should extract env id from queue string", () => { + const queue = "org:123:env:456:queue:testQueue"; + expect(producer.envIdFromQueue(queue)).toBe("456"); + }); + }); +}); diff --git a/apps/webapp/test/utils/marqs.ts b/apps/webapp/test/utils/marqs.ts index 246ce413d3..4fc3884926 100644 --- a/apps/webapp/test/utils/marqs.ts +++ b/apps/webapp/test/utils/marqs.ts @@ -1,5 +1,5 @@ import { MarQSKeyProducer } from "~/v3/marqs/types"; -import { MarQSShortKeyProducer } from "~/v3/marqs/marqsKeyProducer.server.js"; +import { MarQSShortKeyProducer } from "~/v3/marqs/marqsKeyProducer.js"; import Redis from "ioredis"; export function createKeyProducer(prefix: string): MarQSKeyProducer { @@ -48,36 +48,13 @@ export async function setupQueue({ type SetupConcurrencyOptions = { redis: Redis; keyProducer: MarQSKeyProducer; - org: { id: string; currentConcurrency: number; limit?: number; isDisabled?: boolean }; - env: { id: string; currentConcurrency: number; limit?: number }; + env: { id: string; currentConcurrency: number; limit?: number; reserveConcurrency?: number }; }; /** * Sets up concurrency-related Redis keys for orgs and envs */ -export async function setupConcurrency({ redis, keyProducer, org, env }: SetupConcurrencyOptions) { - // Set org concurrency limit if provided - if (typeof org.limit === "number") { - await redis.set(keyProducer.orgConcurrencyLimitKey(org.id), org.limit.toString()); - } - - if (org.currentConcurrency > 0) { - // Set current concurrency by adding dummy members to the set - const orgCurrentKey = keyProducer.orgCurrentConcurrencyKey(org.id); - - // Add dummy running job IDs to simulate current concurrency - const dummyJobs = Array.from( - { length: org.currentConcurrency }, - (_, i) => `dummy-job-${i}-${Date.now()}` - ); - - await redis.sadd(orgCurrentKey, ...dummyJobs); - } - - if (org.isDisabled) { - await redis.set(keyProducer.disabledConcurrencyLimitKey(org.id), "1"); - } - +export async function setupConcurrency({ redis, keyProducer, env }: SetupConcurrencyOptions) { // Set env concurrency limit if (typeof env.limit === "number") { await redis.set(keyProducer.envConcurrencyLimitKey(env.id), env.limit.toString()); @@ -95,6 +72,19 @@ export async function setupConcurrency({ redis, keyProducer, org, env }: SetupCo await redis.sadd(envCurrentKey, ...dummyJobs); } + + if (env.reserveConcurrency && env.reserveConcurrency > 0) { + // Set reserved concurrency by adding dummy members to the set + const envReservedKey = keyProducer.envReserveConcurrencyKey(env.id); + + // Add dummy reserved job IDs to simulate reserved concurrency + const dummyJobs = Array.from( + { length: env.reserveConcurrency }, + (_, i) => `dummy-reserved-job-${i}-${Date.now()}` + ); + + await redis.sadd(envReservedKey, ...dummyJobs); + } } /** diff --git a/docs/images/recursive-task-deadlock-min.png b/docs/images/recursive-task-deadlock-min.png new file mode 100644 index 0000000000..423a539928 Binary files /dev/null and b/docs/images/recursive-task-deadlock-min.png differ diff --git a/docs/queue-concurrency.mdx b/docs/queue-concurrency.mdx index b629834eda..3529eaa44c 100644 --- a/docs/queue-concurrency.mdx +++ b/docs/queue-concurrency.mdx @@ -3,13 +3,26 @@ title: "Concurrency & Queues" description: "Configure what you want to happen when there is more than one run at a time." --- +When you trigger a task, it isn't executed immediately. Instead, the task [run](/runs) is placed into a queue for execution. By default, each task gets its own queue with unbounded concurrency—meaning the task runs as soon as resources are available, subject only to the overall concurrency limits of your environment. If you need more control (for example, to limit concurrency or share limits across multiple tasks), you can define a custom queue as described later in this document. + Controlling concurrency is useful when you have a task that can't be run concurrently, or when you want to limit the number of runs to avoid overloading a resource. -## One at a time +## Default concurrency + +By default, all tasks have an unbounded concurrency limit, limited only by the overall concurrency limits of your environment. This means that each task could possibly "fill up" the entire +concurrency limit of your environment. + + + Your environment has a maximum concurrency limit which depends on your plan. If you're a paying + customer you can request a higher limit by [contacting us](https://www.trigger.dev/contact). + + +## Setting task concurrency -This task will only ever have a single run executing at a time. All other runs will be queued until the current run is complete. +You can set the concurrency limit for a task by setting the `concurrencyLimit` property on the task's queue. This limits the number of runs that can be executing at any one time: ```ts /trigger/one-at-a-time.ts +// This task will only run one at a time export const oneAtATime = task({ id: "one-at-a-time", queue: { @@ -21,38 +34,14 @@ export const oneAtATime = task({ }); ``` -## Parallelism +This is useful if you need to control access to a shared resource, like a database or an API that has rate limits. -You can execute lots of tasks at once by combining high concurrency with [batch triggering](/triggering) (or just triggering in a loop). - -```ts /trigger/parallelism.ts -export const parallelism = task({ - id: "parallelism", - queue: { - concurrencyLimit: 100, - }, - run: async (payload) => { - //... - }, -}); -``` - - - Be careful with high concurrency. If you're doing API requests you might hit rate limits. If - you're hitting your database you might overload it. - - - - Your organization has a maximum concurrency limit which depends on your plan. If you're a paying - customer you can request a higher limit by [contacting us](https://www.trigger.dev/contact). - - -## Defining a queue +## Sharing concurrency between tasks As well as putting queue settings directly on a task, you can define a queue and reuse it across multiple tasks. This allows you to share the same concurrency limit: ```ts /trigger/queue.ts -const myQueue = queue({ +export const myQueue = queue({ name: "my-queue", concurrencyLimit: 1, }); @@ -74,6 +63,8 @@ export const task2 = task({ }); ``` +In this example, `task1` and `task2` share the same queue, so only one of them can run at a time. + ## Setting the concurrency when you trigger a run When you trigger a task you can override the concurrency limit. This is really useful if you sometimes have high priority runs. @@ -81,7 +72,7 @@ When you trigger a task you can override the concurrency limit. This is really u The task: ```ts /trigger/override-concurrency.ts -const generatePullRequest = task({ +export const generatePullRequest = task({ id: "generate-pull-request", queue: { //normally when triggering this task it will be limited to 1 run at a time @@ -107,7 +98,7 @@ export async function POST(request: Request) { queue: { //the "main-branch" queue will have a concurrency limit of 10 //this triggered run will use that queue - name: "main-branch", + name: "main-branch", // Make sure to change the queue name or the task concurrency limit will be updated concurrencyLimit: 10, }, }); @@ -123,7 +114,7 @@ export async function POST(request: Request) { ## Concurrency keys and per-tenant queuing -If you're building an application where you want to run tasks for your users, you might want a separate queue for each of your users. (It doesn't have to be users, it can be any entity you want to separately limit the concurrency for.) +If you're building an application where you want to run tasks for your users, you might want a separate queue for each of your users (or orgs, projects, etc.). You can do this by using `concurrencyKey`. It creates a separate queue for each value of the key. @@ -164,3 +155,146 @@ export async function POST(request: Request) { } } ``` + +## Concurrency and subtasks + +When you trigger a task that has subtasks, the subtasks will not inherit the concurrency settings of the parent task. Unless otherwise specified, subtasks will run on their own queue + +```ts /trigger/subtasks.ts +export const parentTask = task({ + id: "parent-task", + run: async (payload) => { + //trigger a subtask + await subtask.triggerAndWait(payload); + }, +}); + +// This subtask will run on its own queue +export const subtask = task({ + id: "subtask", + run: async (payload) => { + //... + }, +}); +``` + +## Waits and concurrency + +With our [task checkpoint system](/how-it-works#the-checkpoint-resume-system), a parent task can trigger and wait for a subtask to complete. The way this system interacts with the concurrency system is a little complicated but important to understand. There are two main scenarios that we handle slightly differently: + +- When a parent task waits for a subtask on a different queue. +- When a parent task waits for a subtask on the same queue. + +These scenarios are discussed in more detail below: + + + We sometimes refer to the parent task as the "parent" and the subtask as the "child". Subtask and + child task are used interchangeably. We apologize for the confusion. + + +### Waiting for a subtask on a different queue + +During the time when a parent task is waiting on a subtask, the "concurrency" slot of the parent task is still considered occupied on the parent task queue, but is temporarily "released" to the environment. An example will help illustrate this: + +```ts /trigger/waiting.ts +export const parentTask = task({ + id: "parent-task", + queue: { + concurrencyLimit: 1, + }, + run: async (payload) => { + //trigger a subtask + await subtask.triggerAndWait(payload); + }, +}); + +export const subtask = task({ + id: "subtask", + run: async (payload) => { + //... + }, +}); +``` + +For example purposes, let's say the environment concurrency limit is 1. When the parent task is triggered, it will occupy the only slot in the environment. When the parent task triggers the subtask, the subtask will be placed in the queue for the subtask. The parent task will then wait for the subtask to complete. During this time, the parent task slot is temporarily released to the environment, allowing another task to run. Once the subtask completes, the parent task slot is reoccupied. + +This system prevents "stuck" tasks. If the parent task were to wait on the subtask and not release the slot, the environment would be stuck with only one task running. + +And because only the environment slot is released, the parent task queue slot is still occupied. This means that if another task is triggered on the parent task queue, it will be placed in the queue and wait for the parent task to complete, respecting the concurrency limit. + +### Waiting for a subtask on the same queue + +Because tasks can trigger and wait recursively, or share the same queue, we've added special handling for when a parent task waits for a subtask on the same queue. + +Recall above that when waiting for a subtask on a different queue, the parent task slot is temporarily released to the environment. When the parent task and the subtask share a queue, we also release the parent task slot to the queue. Again, an example will help illustrate this: + +```ts /trigger/waiting-same-queue.ts +export const myQueue = queue({ + name: "my-queue", + concurrencyLimit: 1, +}); + +export const parentTask = task({ + id: "parent-task", + queue: myQueue, + run: async (payload) => { + //trigger a subtask + await subtask.triggerAndWait(payload); + }, +}); + +export const subtask = task({ + id: "subtask", + queue: myQueue, + run: async (payload) => { + //... + }, +}); +``` + +In this example, the parent task and the subtask share the same queue with a concurrency limit of 1. When the parent task triggers the subtask, the parent task slot is released to the queue, giving the subtask the opportunity to run. Once the subtask completes, the parent task slot is reoccupied. + +It's very important to note that we only release at-most X slots to the queue, where X is the concurrency limit of the queue. This means that you can only trigger and wait for X subtasks on the same queue. If you try to trigger and wait for more than X subtasks, you will receive a `RECURSIVE_WAIT_DEADLOCK` error. The following example will result in a deadlock: + +```ts /trigger/deadlock.ts +export const myQueue = queue({ + name: "my-queue", + concurrencyLimit: 1, +}); + +export const parentTask = task({ + id: "parent-task", + queue: myQueue, + run: async (payload) => { + //trigger a subtask + await subtask.triggerAndWait(payload); + }, +}); + +export const subtask = task({ + id: "subtask", + queue: myQueue, + run: async (payload) => { + //trigger a subtask + await subsubtask.triggerAndWait(payload); + }, +}); + +export const subsubtask = task({ + id: "subsubtask", + queue: myQueue, + run: async (payload) => { + //... + }, +}); +``` + +Now this will result in a `RECURSIVE_WAIT_DEADLOCK` error because the parent task is waiting for the subtask, and the subtask is waiting for the subsubtask, but there is no more concurrency available in the queue. It will look a bit like this in the logs: + +![Recursive task deadlock](/images/recursive-task-deadlock-min.png) + +### Mitigating recursive wait deadlocks + +If you are recursively triggering and waiting for tasks on the same queue, you can mitigate the risk of a deadlock by increasing the concurrency limit of the queue. This will allow you to trigger and wait for more subtasks. + +You can also use different queues for the parent task and the subtask. This will allow you to trigger and wait for more subtasks without the risk of a deadlock. diff --git a/internal-packages/database/prisma/migrations/20250210164232_add_queue_timestamp_to_run/migration.sql b/internal-packages/database/prisma/migrations/20250210164232_add_queue_timestamp_to_run/migration.sql new file mode 100644 index 0000000000..ef3add8f6d --- /dev/null +++ b/internal-packages/database/prisma/migrations/20250210164232_add_queue_timestamp_to_run/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable +ALTER TABLE + "TaskRun" +ADD + COLUMN "queueTimestamp" TIMESTAMP(3); \ No newline at end of file diff --git a/internal-packages/database/prisma/migrations/20250211150836_add_aborted_batch_task_run_status/migration.sql b/internal-packages/database/prisma/migrations/20250211150836_add_aborted_batch_task_run_status/migration.sql new file mode 100644 index 0000000000..4fcc2da707 --- /dev/null +++ b/internal-packages/database/prisma/migrations/20250211150836_add_aborted_batch_task_run_status/migration.sql @@ -0,0 +1,4 @@ +-- AlterEnum +ALTER TYPE "BatchTaskRunStatus" +ADD + VALUE 'ABORTED'; \ No newline at end of file diff --git a/internal-packages/database/prisma/schema.prisma b/internal-packages/database/prisma/schema.prisma index d676484d9b..7ff86d95e1 100644 --- a/internal-packages/database/prisma/schema.prisma +++ b/internal-packages/database/prisma/schema.prisma @@ -1731,6 +1731,8 @@ model TaskRun { taskEventStore String @default("taskEvent") + queueTimestamp DateTime? + batchItems BatchTaskRunItem[] dependency TaskRunDependency? CheckpointRestoreEvent CheckpointRestoreEvent[] @@ -2208,6 +2210,7 @@ model BatchTaskRun { enum BatchTaskRunStatus { PENDING COMPLETED + ABORTED } model BatchTaskRunItem { diff --git a/packages/core/src/v3/errors.ts b/packages/core/src/v3/errors.ts index a079d0d71c..46d1f1a3df 100644 --- a/packages/core/src/v3/errors.ts +++ b/packages/core/src/v3/errors.ts @@ -235,6 +235,7 @@ export function shouldRetryError(error: TaskRunError): boolean { case "TASK_RUN_HEARTBEAT_TIMEOUT": case "OUTDATED_SDK_VERSION": case "TASK_DID_CONCURRENT_WAIT": + case "RECURSIVE_WAIT_DEADLOCK": return false; case "GRACEFUL_EXIT_TIMEOUT": @@ -512,6 +513,14 @@ const prettyInternalErrors: Partial< href: links.docs.troubleshooting.concurrentWaits, }, }, + RECURSIVE_WAIT_DEADLOCK: { + message: + "This run will never execute because it was triggered recursively and the task has no remaining concurrency available.", + link: { + name: "See docs for help", + href: links.docs.concurrency.recursiveDeadlock, + }, + }, }; const getPrettyTaskRunError = (code: TaskRunInternalError["code"]): TaskRunInternalError => { @@ -672,6 +681,11 @@ export function exceptionEventEnhancer( default: return exception; } + } else if (exception.message?.includes(TaskRunErrorCodes.RECURSIVE_WAIT_DEADLOCK)) { + return { + ...exception, + ...prettyInternalErrors.RECURSIVE_WAIT_DEADLOCK, + }; } break; } @@ -898,3 +912,20 @@ function tryJsonParse(data: string | undefined): any { return; } } + +export function taskRunErrorToString(error: TaskRunError): string { + switch (error.type) { + case "INTERNAL_ERROR": { + return `Internal error [${error.code}]${error.message ? `: ${error.message}` : ""}`; + } + case "BUILT_IN_ERROR": { + return `${error.name}: ${error.message}`; + } + case "STRING_ERROR": { + return error.raw; + } + case "CUSTOM_ERROR": { + return error.raw; + } + } +} diff --git a/packages/core/src/v3/links.ts b/packages/core/src/v3/links.ts index 11a45fe023..b7936a5bc3 100644 --- a/packages/core/src/v3/links.ts +++ b/packages/core/src/v3/links.ts @@ -15,6 +15,10 @@ export const links = { troubleshooting: { concurrentWaits: "https://trigger.dev/docs/troubleshooting#parallel-waits-are-not-supported", }, + concurrency: { + recursiveDeadlock: + "https://trigger.dev/docs/queue-concurrency#waiting-for-a-subtask-on-the-same-queue", + }, }, site: { home: "https://trigger.dev", diff --git a/packages/core/src/v3/schemas/common.ts b/packages/core/src/v3/schemas/common.ts index 5b47d99c20..a8f9faaeed 100644 --- a/packages/core/src/v3/schemas/common.ts +++ b/packages/core/src/v3/schemas/common.ts @@ -172,6 +172,7 @@ export const TaskRunInternalError = z.object({ "POD_UNKNOWN_ERROR", "OUTDATED_SDK_VERSION", "TASK_DID_CONCURRENT_WAIT", + "RECURSIVE_WAIT_DEADLOCK", ]), message: z.string().optional(), stackTrace: z.string().optional(), diff --git a/packages/core/src/v3/utils/flattenAttributes.ts b/packages/core/src/v3/utils/flattenAttributes.ts index 545b0184e6..e2791f21ac 100644 --- a/packages/core/src/v3/utils/flattenAttributes.ts +++ b/packages/core/src/v3/utils/flattenAttributes.ts @@ -5,7 +5,7 @@ export const CIRCULAR_REFERENCE_SENTINEL = "$@circular(("; export function flattenAttributes( obj: Record | Array | string | boolean | number | null | undefined, - prefix?: string , + prefix?: string, seen: WeakSet = new WeakSet() ): Attributes { const result: Attributes = {}; @@ -51,14 +51,13 @@ export function flattenAttributes( seen.add(obj); } - for (const [key, value] of Object.entries(obj)) { const newPrefix = `${prefix ? `${prefix}.` : ""}${Array.isArray(obj) ? `[${key}]` : key}`; if (Array.isArray(value)) { for (let i = 0; i < value.length; i++) { if (typeof value[i] === "object" && value[i] !== null) { // update null check here as well - Object.assign(result, flattenAttributes(value[i], `${newPrefix}.[${i}]`,seen)); + Object.assign(result, flattenAttributes(value[i], `${newPrefix}.[${i}]`, seen)); } else { if (value[i] === null) { result[`${newPrefix}.[${i}]`] = NULL_SENTINEL; @@ -152,7 +151,6 @@ export function unflattenAttributes( if (lastPart !== undefined) { current[lastPart] = rehydrateNull(rehydrateCircular(value)); - } } diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 7d14d7f89d..af762410b5 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1744,6 +1744,25 @@ importers: specifier: ^5 version: 5.5.4 + references/test-tasks: + dependencies: + '@trigger.dev/sdk': + specifier: workspace:* + version: link:../../packages/trigger-sdk + zod: + specifier: 3.23.8 + version: 3.23.8 + devDependencies: + '@trigger.dev/build': + specifier: workspace:* + version: link:../../packages/build + trigger.dev: + specifier: workspace:* + version: link:../../packages/cli-v3 + typescript: + specifier: ^5.5.4 + version: 5.5.4 + references/v3-catalog: dependencies: '@effect/schema': diff --git a/references/test-tasks/package.json b/references/test-tasks/package.json new file mode 100644 index 0000000000..6739b44b30 --- /dev/null +++ b/references/test-tasks/package.json @@ -0,0 +1,18 @@ +{ + "name": "references-test-tasks", + "private": true, + "type": "module", + "scripts": { + "dev": "trigger dev", + "deploy": "trigger deploy --self-hosted --load-image" + }, + "dependencies": { + "@trigger.dev/sdk": "workspace:*", + "zod": "3.23.8" + }, + "devDependencies": { + "@trigger.dev/build": "workspace:*", + "trigger.dev": "workspace:*", + "typescript": "^5.5.4" + } +} \ No newline at end of file diff --git a/references/test-tasks/src/trigger/test-reserve-concurrency-system.ts b/references/test-tasks/src/trigger/test-reserve-concurrency-system.ts new file mode 100644 index 0000000000..bf33753e8d --- /dev/null +++ b/references/test-tasks/src/trigger/test-reserve-concurrency-system.ts @@ -0,0 +1,545 @@ +import { BatchResult, logger, queue, task, wait } from "@trigger.dev/sdk/v3"; +import assert from "assert"; +import { + updateEnvironmentConcurrencyLimit, + waitForRunStatus, + getEnvironmentStats, +} from "../utils.js"; + +export const describeReserveConcurrencySystem = task({ + id: "describe/reserve-concurrency-system", + retry: { + maxAttempts: 1, + }, + run: async (payload: any, { ctx }) => { + await testRetryPriority.triggerAndWait({ holdDelayMs: 10_000 }).unwrap(); + + logger.info("✅ Tested retry priority, now testing resume priority"); + + await testResumePriority.triggerAndWait({ initialDelayMs: 5_000, useBatch: false }).unwrap(); + await testResumePriority.triggerAndWait({ initialDelayMs: 30_000, useBatch: false }).unwrap(); + + logger.info("✅ Tested resume priority with triggerAndWait"); + + await testResumePriority.triggerAndWait({ initialDelayMs: 5_000, useBatch: true }).unwrap(); + await testResumePriority.triggerAndWait({ initialDelayMs: 30_000, useBatch: true }).unwrap(); + + logger.info("✅ Tested resume priority with batchTriggerAndWait"); + + await testResumeDurationPriority.triggerAndWait({ waitDurationInSeconds: 30 }).unwrap(); + await testResumeDurationPriority.triggerAndWait({ waitDurationInSeconds: 65 }).unwrap(); + + logger.info("✅ Tested resume duration priority with wait.for"); + + await testEnvReserveConcurrency + .triggerAndWait({ envConcurrencyLimit: 4, holdTaskCount: 1, useBatch: false }) + .unwrap(); + + logger.info("✅ Tested env reserve concurrency system with triggerAndWait"); + + await testEnvReserveConcurrency + .triggerAndWait({ envConcurrencyLimit: 4, holdTaskCount: 1, useBatch: true }) + .unwrap(); + + logger.info("✅ Tested env reserve concurrency system with batchTriggerAndWait"); + + await testQueueReserveConcurrency.triggerAndWait({ useBatch: false }).unwrap(); + + logger.info("✅ Tested queue reserve concurrency system with triggerAndWait"); + + await testQueueReserveConcurrency.triggerAndWait({ useBatch: true }).unwrap(); + + logger.info("✅ Tested queue reserve concurrency system with batchTriggerAndWait"); + }, +}); + +export const testRetryPriority = task({ + id: "test/retry-priority", + retry: { + maxAttempts: 1, + }, + run: async ({ holdDelayMs = 10_000 }: { holdDelayMs: number }, { ctx }) => { + const startEnvStats = await getEnvironmentStats(ctx.environment.id); + + // We need to test the reserve concurrency system + // 1. Retries are prioritized over new runs + // Setup: Trigger a run that fails and will re-attempt in 5 seconds + // Trigger another run that uses the same concurrency, and hits the max concurrency of that queue + // Trigger a run on that same queue before the retry is attempted + // The "hold" run will then complete and the retry should be dequeued + // Once the retry completes successfully, the 3rd run should be dequeued + + const failureRun = await retryTask.trigger( + { delayMs: 0, throwError: true, failureCount: 1 }, + { tags: ["failure"] } + ); + await waitForRunStatus(failureRun.id, ["EXECUTING", "REATTEMPTING"]); + + logger.info("Failure run is executing, triggering a run that will hit the concurrency limit"); + + const holdRun = await retryTask.trigger( + { delayMs: holdDelayMs, throwError: false, failureCount: 0 }, + { tags: ["hold"] } + ); + await waitForRunStatus(holdRun.id, ["EXECUTING"]); + + logger.info("Hold run is executing, triggering a run that will be queued"); + + const queuedRun = await retryTask.trigger( + { delayMs: 0, throwError: false, failureCount: 0 }, + { tags: ["queued"] } + ); + + logger.info("Queued run is queued, waiting for the hold run to complete"); + + const completedFailureRun = await waitForRunStatus(failureRun.id, ["COMPLETED"]); + const completedQueuedRun = await waitForRunStatus(queuedRun.id, ["COMPLETED"]); + + logger.info("Runs completed", { + completedFailureRun, + completedQueuedRun, + }); + + // Now we need to assert the completedFailureRun.completedAt is before completedQueuedRun.completedAt + assert( + completedFailureRun.finishedAt! < completedQueuedRun.finishedAt!, + "Failure run should complete before queued run" + ); + + // Now lets make sure all the runs are completed + await waitForRunStatus(holdRun.id, ["COMPLETED"]); + + const envStats = await getEnvironmentStats(ctx.environment.id); + + logger.info("Environment stats", envStats); + + assert( + startEnvStats.reserveConcurrency - envStats.reserveConcurrency === 0, + "Reserve concurrency should be 0" + ); + + logger.info("✅ Failure run completed before queued run"); + }, +}); + +export const testResumePriority = task({ + id: "test/resume-priority", + retry: { + maxAttempts: 1, + }, + run: async ( + { initialDelayMs = 5_000, useBatch = false }: { initialDelayMs: number; useBatch: boolean }, + { ctx } + ) => { + const startEnvStats = await getEnvironmentStats(ctx.environment.id); + + // 2. Resumed runs are prioritized over new runs + const resumeRun = await resumeParentTask.trigger( + { delayMs: initialDelayMs, triggerChildTask: true, useBatch }, + { tags: ["resume"] } + ); + await waitForRunStatus(resumeRun.id, ["EXECUTING", "FROZEN"]); + + logger.info("Resume run is executing, triggering a run that should be queued"); + const queuedRun = await resumeParentTask.trigger( + { delayMs: 1_000, triggerChildTask: false, useBatch }, + { tags: ["queued"] } + ); + await waitForRunStatus(queuedRun.id, ["QUEUED"]); + + const completedResumeRun = await waitForRunStatus(resumeRun.id, ["COMPLETED"]); + const completedQueuedRun = await waitForRunStatus(queuedRun.id, ["COMPLETED"]); + + logger.info("Runs completed", { + completedResumeRun, + completedQueuedRun, + }); + + // Now we need to assert the completedResumeRun.completedAt is before completedQueuedRun.completedAt + assert( + completedResumeRun.finishedAt! < completedQueuedRun.finishedAt!, + "Resume run should complete before queued run" + ); + + const envStats = await getEnvironmentStats(ctx.environment.id); + + assert( + startEnvStats.reserveConcurrency - envStats.reserveConcurrency === 0, + "Reserve concurrency should be 0" + ); + + logger.info("✅ Resume run completed before queued run"); + }, +}); + +export const testResumeDurationPriority = task({ + id: "test/resume-duration-priority", + retry: { + maxAttempts: 1, + }, + run: async ({ waitDurationInSeconds = 5 }: { waitDurationInSeconds: number }, { ctx }) => { + const startEnvStats = await getEnvironmentStats(ctx.environment.id); + + // 2. Resumed runs are prioritized over new runs + const resumeRun = await durationWaitTask.trigger( + { waitDurationInSeconds, doWait: true }, + { tags: ["resume"] } + ); + await waitForRunStatus(resumeRun.id, ["EXECUTING", "FROZEN"]); + + logger.info( + "Resume run is executing, triggering a run that will hold the concurrency until both the resume run and the queued run are in the queue" + ); + + if (ctx.environment.type !== "DEVELOPMENT") { + const holdRun = await durationWaitTask.trigger( + { waitDurationInSeconds: waitDurationInSeconds + 10, doWait: false }, + { tags: ["hold"] } + ); + await waitForRunStatus(holdRun.id, ["EXECUTING"]); + + logger.info("Hold run is executing, triggering a run that should be queued"); + } + + const queuedRun = await durationWaitTask.trigger( + { waitDurationInSeconds: 1, doWait: false }, + { tags: ["queued"] } + ); + await waitForRunStatus(queuedRun.id, ["QUEUED"]); + + const completedResumeRun = await waitForRunStatus(resumeRun.id, ["COMPLETED"]); + const completedQueuedRun = await waitForRunStatus(queuedRun.id, ["COMPLETED"]); + + logger.info("Runs completed", { + completedResumeRun, + completedQueuedRun, + }); + + // Now we need to assert the completedResumeRun.completedAt is before completedQueuedRun.completedAt + assert( + completedResumeRun.finishedAt! < completedQueuedRun.finishedAt!, + "Resume run should complete before queued run" + ); + + const envStats = await getEnvironmentStats(ctx.environment.id); + + assert( + startEnvStats.reserveConcurrency - envStats.reserveConcurrency === 0, + "Reserve concurrency should be 0" + ); + + logger.info("✅ Resume run completed before queued run"); + }, +}); + +export const testEnvReserveConcurrency = task({ + id: "test/env-reserve-concurrency", + retry: { + maxAttempts: 1, + }, + run: async ( + { + envConcurrencyLimit = 3, + holdTaskCount = 1, + useBatch = false, + }: { envConcurrencyLimit: number; holdTaskCount: number; useBatch: boolean }, + { ctx } + ) => { + const startEnvStats = await getEnvironmentStats(ctx.environment.id); + + // 3. When a task triggerAndWaits another task, the parent run should be added to the envs reserve concurrency + // Giving the environment "back" another concurrency slot. Another task (not the parent task) can then be dequeued + // We need to be able to "fill" the env concurrency (sans 1), then trigger the parent task. The parent task then triggerAndWaits + // a child task. We need to make sure the child task executes + await updateEnvironmentConcurrencyLimit(ctx.environment.id, envConcurrencyLimit); + + const holdBatch = await delayTask.batchTrigger( + Array.from({ length: holdTaskCount }, (_, i) => ({ + payload: { delayMs: 30_000 }, + options: { tags: ["hold"] }, + })) + ); + + // Wait for the hold tasks to be executing + await Promise.all(holdBatch.runs.map((run) => waitForRunStatus(run.id, ["EXECUTING"]))); + + // Now we will trigger a parent task that will trigger a child task + const parentRun = await genericParentTask.trigger( + { delayMs: 1_000, triggerChildTask: true, useBatch }, + { tags: ["parent"] } + ); + + // Once the parentRun starts executing, we will be at the max concurrency limit + await waitForRunStatus(parentRun.id, ["EXECUTING"], 5); // timeout after 5 seconds, to ensure the parent task is executing + + // But because the parent task triggers a child task, the env reserve concurrency will allow the child task to execute + logger.info("Parent task is executing, waiting for child task to complete"); + + await waitForRunStatus(parentRun.id, ["COMPLETED"], 10); // timeout after 10 seconds, to ensure the child task finished before the delay runs + + logger.info( + "Parent task completed, which means the child task completed. Now waiting for the hold tasks to complete" + ); + + const envStats = await getEnvironmentStats(ctx.environment.id, "task/generic-parent-task"); + + assert( + startEnvStats.reserveConcurrency - envStats.reserveConcurrency === 0, + "Reserve concurrency should be 0" + ); + assert( + envStats.queueCurrentConcurrency === 0, + "generic-parent-task current concurrency should be 0" + ); + assert( + envStats.queueReserveConcurrency === 0, + "generic-parent-task reserve concurrency should be 0" + ); + + const childStats = await getEnvironmentStats(ctx.environment.id, "task/generic-child-task"); + + assert( + childStats.queueReserveConcurrency === 0, + "generic-child-task reserve concurrency should be 0" + ); + assert( + childStats.queueCurrentConcurrency === 0, + "generic-child-task current concurrency should be 0" + ); + + // Wait for the hold tasks to be completed + await Promise.all(holdBatch.runs.map((run) => waitForRunStatus(run.id, ["COMPLETED"]))); + + await updateEnvironmentConcurrencyLimit(ctx.environment.id, 100); + + logger.info("✅ Environment reserve concurrency system is working as expected"); + }, +}); + +export const testQueueReserveConcurrency = task({ + id: "test/queue-reserve-concurrency", + retry: { + maxAttempts: 1, + }, + run: async ({ useBatch = false }: { useBatch: boolean }, { ctx }) => { + const startEnvStats = await getEnvironmentStats(ctx.environment.id); + // This test ensures that when triggerAndWait is called where the parent and the child share a queue, + // the queue reserve concurrency is used to allow the child to execute. + // We also want to test that the queue can only "reserve" at most up to the concurrency limit, and if + // the reservation fails, the child task will fail + const rootRecursiveRun = await recursiveTask.trigger( + { delayMs: 1_000, depth: 1, useBatch }, + { tags: ["root"] } + ); + + const completedRootRun = await waitForRunStatus(rootRecursiveRun.id, ["COMPLETED"], 20); + + assert(completedRootRun.status === "COMPLETED", "Root recursive run should be completed"); + + const failingRootRecursiveRun = await recursiveTask.trigger( + { delayMs: 1_000, depth: 2, useBatch }, + { tags: ["failing-root"] } + ); + + const failedRootRun = await waitForRunStatus(failingRootRecursiveRun.id, ["COMPLETED"], 20); + + assert(!failedRootRun.output?.ok, "Child of failing root run should fail"); + + const envStats = await getEnvironmentStats(ctx.environment.id, "task/recursive-task"); + + logger.info("Environment stats", envStats); + + assert( + startEnvStats.reserveConcurrency - envStats.reserveConcurrency === 0, + "Env reserve concurrency should be 0" + ); + assert( + envStats.queueCurrentConcurrency === 0, + "queue-reserve-concurrency current concurrency should be 0" + ); + assert( + envStats.queueReserveConcurrency === 0, + "queue-reserve-concurrency reserve concurrency should be 0" + ); + }, +}); + +export const recursiveTask = task({ + id: "recursive-task", + queue: { + concurrencyLimit: 1, + }, + retry: { + maxAttempts: 1, + }, + run: async ( + { delayMs, depth, useBatch = false }: { delayMs: number; depth: number; useBatch: boolean }, + { ctx } + ) => { + if (depth === 0) { + return; + } + + await new Promise((resolve) => setTimeout(resolve, delayMs)); + + if (useBatch) { + const batchResult = await recursiveTask.batchTriggerAndWait([ + { + payload: { delayMs, depth: depth - 1, useBatch }, + options: { tags: ["recursive"] }, + }, + ]); + + const firstRun = batchResult.runs[0] as any; + + return { + ok: firstRun.ok, + }; + } else { + const result = (await recursiveTask.triggerAndWait({ + delayMs, + depth: depth - 1, + useBatch, + })) as any; + + return { + ok: result.ok, + }; + } + }, +}); + +export const singleQueue = queue({ + name: "single-queue", + concurrencyLimit: 1, +}); + +export const delayTask = task({ + id: "delay-task", + retry: { + maxAttempts: 1, + }, + run: async (payload: { delayMs: number }, { ctx }) => { + await new Promise((resolve) => setTimeout(resolve, payload.delayMs)); + }, +}); + +export const retryTask = task({ + id: "retry-task", + queue: singleQueue, + retry: { + maxAttempts: 10, + minTimeoutInMs: 5_000, // Will retry in 5 seconds + maxTimeoutInMs: 5_000, + }, + run: async (payload: { delayMs: number; throwError: boolean; failureCount: number }, { ctx }) => { + await new Promise((resolve) => setTimeout(resolve, payload.delayMs)); + + if (payload.throwError && ctx.attempt.number <= payload.failureCount) { + throw new Error("Error"); + } + }, +}); + +export const durationWaitTask = task({ + id: "duration-wait-task", + queue: { + concurrencyLimit: 1, + }, + run: async ( + { + waitDurationInSeconds = 5, + doWait = true, + }: { waitDurationInSeconds: number; doWait: boolean }, + { ctx } + ) => { + if (doWait) { + await wait.for({ seconds: waitDurationInSeconds }); + } else { + await new Promise((resolve) => setTimeout(resolve, waitDurationInSeconds * 1000)); + } + }, +}); + +export const resumeParentTask = task({ + id: "resume-parent-task", + queue: { + concurrencyLimit: 1, + }, + run: async ( + { + delayMs = 5_000, + triggerChildTask, + useBatch = false, + }: { delayMs: number; triggerChildTask: boolean; useBatch: boolean }, + { ctx } + ) => { + if (triggerChildTask) { + if (useBatch) { + const batchResult = await resumeChildTask.batchTriggerAndWait([ + { + payload: { delayMs }, + options: { tags: ["resume-child"] }, + }, + ]); + + unwrapBatchResult(batchResult); + } else { + await resumeChildTask.triggerAndWait({ delayMs }, { tags: ["resume-child"] }).unwrap(); + } + } else { + await new Promise((resolve) => setTimeout(resolve, delayMs)); + } + }, +}); + +export const resumeChildTask = task({ + id: "resume-child-task", + run: async (payload: { delayMs: number }, { ctx }) => { + await new Promise((resolve) => setTimeout(resolve, payload.delayMs)); + }, +}); + +export const genericParentTask = task({ + id: "generic-parent-task", + run: async ( + { + delayMs = 5_000, + triggerChildTask, + useBatch = false, + }: { delayMs: number; triggerChildTask: boolean; useBatch: boolean }, + { ctx } + ) => { + if (triggerChildTask) { + if (useBatch) { + const batchResult = await genericChildTask.batchTriggerAndWait([ + { + payload: { delayMs }, + options: { tags: ["resume-child"] }, + }, + ]); + + return unwrapBatchResult(batchResult); + } else { + await genericChildTask.triggerAndWait({ delayMs }, { tags: ["resume-child"] }).unwrap(); + } + } else { + await new Promise((resolve) => setTimeout(resolve, delayMs)); + } + }, +}); + +function unwrapBatchResult(batchResult: BatchResult) { + if (batchResult.runs.some((run) => !run.ok)) { + throw new Error(`Child task failed: ${batchResult.runs.find((run) => !run.ok)?.error}`); + } + + return batchResult.runs; +} + +export const genericChildTask = task({ + id: "generic-child-task", + run: async (payload: { delayMs: number }, { ctx }) => { + await new Promise((resolve) => setTimeout(resolve, payload.delayMs)); + }, +}); diff --git a/references/test-tasks/src/utils.ts b/references/test-tasks/src/utils.ts new file mode 100644 index 0000000000..21a83e1037 --- /dev/null +++ b/references/test-tasks/src/utils.ts @@ -0,0 +1,109 @@ +import { runs } from "@trigger.dev/sdk/v3"; +import { z } from "zod"; + +export type RunStatus = Awaited>["status"]; + +export async function waitForRunStatus( + id: string, + statuses: RunStatus[], + timeoutInSeconds?: number +) { + const run = await runs.retrieve(id); + + if (statuses.includes(run.status)) { + return run; + } + + const start = Date.now(); + + while (Date.now() - start < (timeoutInSeconds ?? 300) * 1_000) { + const run = await runs.retrieve(id); + + if (statuses.includes(run.status)) { + return run; + } + + await new Promise((resolve) => setTimeout(resolve, 1_000)); + } + + throw new Error( + `Run did not reach status ${statuses.join(" or ")} within ${timeoutInSeconds ?? 300} seconds` + ); +} + +export async function updateEnvironmentConcurrencyLimit( + environmentId: string, + concurrencyLimit: number +) { + if (!process.env.TRIGGER_API_URL) { + throw new Error("TRIGGER_API_URL is not set"); + } + + if (!process.env.TRIGGER_ACCESS_TOKEN) { + throw new Error("TRIGGER_ACCESS_TOKEN is not set"); + } + + // We need to make a request to baseURL + `/admin/api/v1/environments/${environmentId}` with a POST request + // The body needs to be a JSON object with the key `envMaximumConcurrencyLimit` and the value `concurrencyLimit`, and the key `orgMaximumConcurrencyLimit` with the value `concurrencyLimit` + // we also need a Authorization header that has the personal access token from process.env.TRIGGER_ACCESS_TOKEN + + const response = await fetch( + `${process.env.TRIGGER_API_URL}/admin/api/v1/environments/${environmentId}`, + { + method: "POST", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${process.env.TRIGGER_ACCESS_TOKEN}`, + }, + body: JSON.stringify({ + envMaximumConcurrencyLimit: concurrencyLimit, + orgMaximumConcurrencyLimit: concurrencyLimit, + }), + } + ); + + if (!response.ok) { + throw new Error(`Failed to update environment concurrency limit: ${response.statusText}`); + } + + return response.json(); +} + +const EnvironmentStatsResponseBody = z.object({ + id: z.string(), + concurrencyLimit: z.number(), + currentConcurrency: z.number(), + reserveConcurrency: z.number(), + queueConcurrency: z.number().optional(), + queueReserveConcurrency: z.number().optional(), + queueCurrentConcurrency: z.number().optional(), +}); + +export type EnvironmentStatsResponseBody = z.infer; + +export async function getEnvironmentStats( + environmentId: string, + queue?: string +): Promise { + const url = new URL(`${process.env.TRIGGER_API_URL}/admin/api/v1/environments/${environmentId}`); + + if (queue) { + url.searchParams.set("queue", queue); + } + + const response = await fetch(url.toString(), { + method: "GET", + headers: { + Accept: "application/json", + Authorization: `Bearer ${process.env.TRIGGER_ACCESS_TOKEN}`, + }, + }); + + if (!response.ok) { + throw new Error(`Failed to fetch environment stats: ${response.statusText}`); + } + + const responseBody = await response.json(); + + return EnvironmentStatsResponseBody.parse(responseBody); +} diff --git a/references/test-tasks/trigger.config.ts b/references/test-tasks/trigger.config.ts new file mode 100644 index 0000000000..b515908f41 --- /dev/null +++ b/references/test-tasks/trigger.config.ts @@ -0,0 +1,19 @@ +import { defineConfig } from "@trigger.dev/sdk/v3"; + +export default defineConfig({ + runtime: "node", + project: "proj_qwdshjjnuoepcuupfuvo", + machine: "small-1x", + maxDuration: 3600, + dirs: ["./src/trigger"], + retries: { + enabledInDev: true, + default: { + maxAttempts: 10, + minTimeoutInMs: 5_000, + maxTimeoutInMs: 30_000, + factor: 2, + randomize: true, + }, + }, +}); diff --git a/references/test-tasks/tsconfig.json b/references/test-tasks/tsconfig.json new file mode 100644 index 0000000000..635ecdb766 --- /dev/null +++ b/references/test-tasks/tsconfig.json @@ -0,0 +1,20 @@ +{ + "compilerOptions": { + "target": "esnext", + "module": "NodeNext", + "moduleResolution": "NodeNext", + "esModuleInterop": true, + "strict": true, + "outDir": "dist", + "skipLibCheck": true, + "customConditions": ["@triggerdotdev/source"], + "jsx": "preserve", + "emitDecoratorMetadata": true, + "experimentalDecorators": true, + "lib": ["DOM", "DOM.Iterable"], + "paths": { + "@/*": ["./src/*"] + } + }, + "include": ["./src/**/*.ts", "trigger.config.ts"] +} diff --git a/references/v3-catalog/src/trigger/concurrency.ts b/references/v3-catalog/src/trigger/concurrency.ts index 693b8d4c19..8fbfa021c5 100644 --- a/references/v3-catalog/src/trigger/concurrency.ts +++ b/references/v3-catalog/src/trigger/concurrency.ts @@ -1,4 +1,4 @@ -import { logger, task, wait } from "@trigger.dev/sdk/v3"; +import { logger, queue, task, wait } from "@trigger.dev/sdk/v3"; export const oneAtATime = task({ id: "on-at-a-time", @@ -17,30 +17,230 @@ export const oneAtATime = task({ }); export const testConcurrency = task({ - id: "test-concurrency", - run: async ({ count = 10, delay = 5000 }: { count: number; delay: number }) => { - logger.info(`Running ${count} tasks`); + id: "test-concurrency-controller", + run: async ({ + count = 10, + delay = 5000, + childDelay = 1000, + }: { + count: number; + delay: number; + childDelay: number; + }) => { + logger.info(`Running ${count} tasks baby`); - await new Promise((resolve) => setTimeout(resolve, 3000)); - - await testConcurrencyChild.batchTrigger( + await testConcurrencyParent.batchTrigger( Array.from({ length: count }).map((_, index) => ({ payload: { delay, + childDelay, }, })) ); logger.info(`All ${count} tasks triggered`); + // wait for about 2 seconds + await new Promise((resolve) => setTimeout(resolve, 2000)); + + // Now trigger the parent task again + await testConcurrencyParent.trigger({ + delay, + childDelay, + }); + return { finished: new Date().toISOString(), }; }, }); +export const testConcurrencyParent = task({ + id: "test-concurrency-parent", + run: async ({ delay = 5000, childDelay = 1000 }: { delay: number; childDelay: number }) => { + logger.info(`Delaying for ${delay}ms`); + + await new Promise((resolve) => setTimeout(resolve, delay)); + + logger.info(`Delay of ${delay}ms completed`); + + return await testConcurrencyChild.triggerAndWait({ + delay: childDelay, + }); + }, +}); + export const testConcurrencyChild = task({ id: "test-concurrency-child", + queue: { + concurrencyLimit: 10, + }, + run: async ({ delay = 5000 }: { delay: number }) => { + logger.info(`Delaying for ${delay}ms`); + + await new Promise((resolve) => setTimeout(resolve, delay)); + + logger.info(`Delay of ${delay}ms completed`); + + return { + completedAt: new Date(), + }; + }, +}); + +export const testReserveConcurrencyRecursiveWaits = task({ + id: "test-reserve-concurrency-recursive-waits", + retry: { + maxAttempts: 1, + }, + run: async ({ + delay = 5000, + depth = 2, + currentDepth = 0, + batchSize = 1, + useBatch, + }: { + delay: number; + depth: number; + currentDepth?: number; + batchSize?: number; + useBatch?: boolean; + }) => { + logger.info(`Running task at depth ${currentDepth} 1`); + + logger.info(`Delaying for ${delay}ms`); + + await new Promise((resolve) => setTimeout(resolve, delay)); + + logger.info(`Delay of ${delay}ms completed`); + + if (currentDepth < depth) { + logger.info(`Triggering child task at depth ${currentDepth + 1}`); + + if (useBatch) { + await testReserveConcurrencyRecursiveWaits.batchTriggerAndWait( + Array.from({ length: batchSize }).map((_, index) => ({ + payload: { + delay, + depth, + currentDepth: currentDepth + 1, + batchSize, + useBatch, + }, + })) + ); + } else { + await testReserveConcurrencyRecursiveWaits.triggerAndWait({ + delay, + depth, + currentDepth: currentDepth + 1, + batchSize, + useBatch, + }); + } + + logger.info(`Child task at depth ${currentDepth + 1} completed`); + } + + logger.info(`Task at depth ${currentDepth} completed`); + + return { + completedAt: new Date(), + }; + }, +}); + +export const testChildTaskPriorityController = task({ + id: "test-child-task-priority-controller", + run: async ({ delay = 5000 }: { delay: number }) => { + logger.info(`Delaying for ${delay}ms`); + + await new Promise((resolve) => setTimeout(resolve, delay)); + + return { + completedAt: new Date(), + }; + }, +}); + +export const testChildTaskPriorityParent = task({ + id: "test-child-task-priority-parent", + run: async ({ delay = 5000 }: { delay: number }) => { + logger.info(`Delaying for ${delay}ms`); + + await new Promise((resolve) => setTimeout(resolve, delay)); + + await testChildTaskPriorityChild.triggerAndWait({ + delay, + }); + + logger.info(`Delay of ${delay}ms completed`); + + return { + completedAt: new Date(), + }; + }, +}); + +export const testChildTaskPriorityChildCreator = task({ + id: "test-child-task-priority-child-creator", + run: async ({ delay = 5000 }: { delay: number }) => { + await testChildTaskPriorityChild.batchTrigger([ + { payload: { delay, propagate: false } }, + { payload: { delay, propagate: false } }, + { payload: { delay, propagate: false } }, + { payload: { delay, propagate: false } }, + ]); + + return { + completedAt: new Date(), + }; + }, +}); + +export const testChildTaskPriorityChild = task({ + id: "test-child-task-priority-child", + queue: { + concurrencyLimit: 1, + }, + run: async ({ delay = 5000, propagate }: { delay: number; propagate?: boolean }) => { + logger.info(`Delaying for ${delay}ms`); + + await new Promise((resolve) => setTimeout(resolve, delay)); + + if (typeof propagate === "undefined" || propagate) { + await testChildTaskPriorityGrandChild.triggerAndWait({ + delay, + }); + } + + logger.info(`Delay of ${delay}ms completed`); + + return { + completedAt: new Date(), + }; + }, +}); + +export const testChildTaskPriorityGrandChildCreator = task({ + id: "test-child-task-priority-grand-child-creator", + run: async ({ delay = 5000 }: { delay: number }) => { + await testChildTaskPriorityGrandChild.batchTrigger([ + { payload: { delay } }, + { payload: { delay } }, + { payload: { delay } }, + ]); + + logger.info(`Delay of ${delay}ms completed`); + + return { + completedAt: new Date(), + }; + }, +}); + +export const testChildTaskPriorityGrandChild = task({ + id: "test-child-task-priority-grandchild", queue: { concurrencyLimit: 1, }, @@ -56,3 +256,34 @@ export const testConcurrencyChild = task({ }; }, }); + +export const myQueue = queue({ + name: "my-queue", + concurrencyLimit: 1, +}); + +export const parentTask = task({ + id: "parent-task", + queue: myQueue, + run: async (payload) => { + //trigger a subtask + await subtask.triggerAndWait(payload); + }, +}); + +export const subtask = task({ + id: "subtask", + queue: myQueue, + run: async (payload) => { + //trigger a subtask + await subsubtask.triggerAndWait(payload); + }, +}); + +export const subsubtask = task({ + id: "subsubtask", + queue: myQueue, + run: async (payload) => { + //... + }, +}); diff --git a/references/v3-catalog/src/trigger/simple.ts b/references/v3-catalog/src/trigger/simple.ts index 8dcf96f46e..9172317e7a 100644 --- a/references/v3-catalog/src/trigger/simple.ts +++ b/references/v3-catalog/src/trigger/simple.ts @@ -141,92 +141,92 @@ function thisFunctionWillThrow() { throw new Error("This function will throw"); } -export const parentTask = task({ - id: "parent-task", - run: async (payload: { message: string }, { ctx }) => { - logger.info("Parent task payload", { payload }); - - console.info("This is an info message"); - logger.info("This is an info message from logger.info"); - console.log(JSON.stringify({ ctx, message: "This is the parent task contexts" })); - logger.log(JSON.stringify({ ctx, message: "This is the parent task context from logger.log" })); - console.warn("You've been warned buddy"); - logger.warn("You've been warned buddy from logger.warn"); - console.error("This is an error message"); - logger.error("This is an error message from logger.error"); - - await wait.for({ seconds: 5 }); - - const childTaskResponse = await childTask - .triggerAndWait({ - message: payload.message, - forceError: false, - }) - .unwrap(); - - logger.info("Child task response", { childTaskResponse }); - - await childTask.trigger({ - message: `${payload.message} - 2.a`, - forceError: true, - }); - - await new Promise((resolve) => setTimeout(resolve, 1000)); - - return { - message: payload.message, - childTaskResponse, - }; - }, -}); - -export const childTask = task({ - id: "child-task", - run: async ( - payload: { message: string; forceError: boolean; delayInSeconds?: number }, - { ctx } - ) => { - logger.info("Child task payload", { payload }); - logger.info("Child task payload 2", { payload }); - logger.info("Child task payload 3", { payload }); - logger.info("Child task payload 4", { payload }); - logger.info("Child task payload 5", { payload }); - - await wait.for({ seconds: payload.delayInSeconds ?? 5 }); - - logger.info("Child task payload 6", { payload }); - logger.info("Child task payload 7", { payload }); - logger.info("Child task payload 8", { payload }); - - const response = await fetch("https://jsonhero.io/api/create.json", { - method: "POST", - headers: { - "Content-Type": "application/json", - }, - body: JSON.stringify({ - title: "childTask payload and ctxr", - content: { - payload, - ctx, - }, - readOnly: true, - }), - }); - - const json: any = await response.json(); - - logger.info("JSONHero response", { json }); - - if (payload.forceError) { - throw new Error(`Forced error: ${payload.message}`); - } - - return { - message: "This is the child task", - parentMessage: payload.message, - }; - }, -}); +// export const parentTask = task({ +// id: "parent-task", +// run: async (payload: { message: string }, { ctx }) => { +// logger.info("Parent task payload", { payload }); + +// console.info("This is an info message"); +// logger.info("This is an info message from logger.info"); +// console.log(JSON.stringify({ ctx, message: "This is the parent task contexts" })); +// logger.log(JSON.stringify({ ctx, message: "This is the parent task context from logger.log" })); +// console.warn("You've been warned buddy"); +// logger.warn("You've been warned buddy from logger.warn"); +// console.error("This is an error message"); +// logger.error("This is an error message from logger.error"); + +// await wait.for({ seconds: 5 }); + +// const childTaskResponse = await childTask +// .triggerAndWait({ +// message: payload.message, +// forceError: false, +// }) +// .unwrap(); + +// logger.info("Child task response", { childTaskResponse }); + +// await childTask.trigger({ +// message: `${payload.message} - 2.a`, +// forceError: true, +// }); + +// await new Promise((resolve) => setTimeout(resolve, 1000)); + +// return { +// message: payload.message, +// childTaskResponse, +// }; +// }, +// }); + +// export const childTask = task({ +// id: "child-task", +// run: async ( +// payload: { message: string; forceError: boolean; delayInSeconds?: number }, +// { ctx } +// ) => { +// logger.info("Child task payload", { payload }); +// logger.info("Child task payload 2", { payload }); +// logger.info("Child task payload 3", { payload }); +// logger.info("Child task payload 4", { payload }); +// logger.info("Child task payload 5", { payload }); + +// await wait.for({ seconds: payload.delayInSeconds ?? 5 }); + +// logger.info("Child task payload 6", { payload }); +// logger.info("Child task payload 7", { payload }); +// logger.info("Child task payload 8", { payload }); + +// const response = await fetch("https://jsonhero.io/api/create.json", { +// method: "POST", +// headers: { +// "Content-Type": "application/json", +// }, +// body: JSON.stringify({ +// title: "childTask payload and ctxr", +// content: { +// payload, +// ctx, +// }, +// readOnly: true, +// }), +// }); + +// const json: any = await response.json(); + +// logger.info("JSONHero response", { json }); + +// if (payload.forceError) { +// throw new Error(`Forced error: ${payload.message}`); +// } + +// return { +// message: "This is the child task", +// parentMessage: payload.message, +// }; +// }, +// }); export const retryTask = task({ id: "retry-task", diff --git a/references/v3-catalog/trigger.config.ts b/references/v3-catalog/trigger.config.ts index 830ac38999..9a24160445 100644 --- a/references/v3-catalog/trigger.config.ts +++ b/references/v3-catalog/trigger.config.ts @@ -20,7 +20,7 @@ export default defineConfig({ maxDuration: 3600, dirs: ["./src/trigger"], retries: { - enabledInDev: true, + enabledInDev: false, default: { maxAttempts: 10, minTimeoutInMs: 5_000,