diff --git a/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.wait.duration.ts b/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.wait.duration.ts index 5ffdd138b..24aa18140 100644 --- a/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.wait.duration.ts +++ b/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.wait.duration.ts @@ -48,10 +48,9 @@ const { action } = createActionApiRoute( const waitResult = await engine.blockRunWithWaitpoint({ runId: run.id, waitpoints: waitpoint.id, - environmentId: authentication.environment.id, projectId: authentication.environment.project.id, organizationId: authentication.environment.organization.id, - releaseConcurrency: true, + releaseConcurrency: body.releaseConcurrency, }); return json({ diff --git a/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.waitpoints.tokens.$waitpointFriendlyId.wait.ts b/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.waitpoints.tokens.$waitpointFriendlyId.wait.ts index c0e19c1f4..e9bd27d69 100644 --- a/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.waitpoints.tokens.$waitpointFriendlyId.wait.ts +++ b/apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.waitpoints.tokens.$waitpointFriendlyId.wait.ts @@ -34,13 +34,12 @@ const { action } = createActionApiRoute( throw json({ error: "Waitpoint not found" }, { status: 404 }); } + // TODO: Add releaseConcurrency from the body const result = await engine.blockRunWithWaitpoint({ runId, waitpoints: [waitpointId], - environmentId: authentication.environment.id, projectId: authentication.environment.project.id, organizationId: authentication.environment.organization.id, - releaseConcurrency: true, }); return json( diff --git a/apps/webapp/app/v3/services/triggerTaskV2.server.ts b/apps/webapp/app/v3/services/triggerTaskV2.server.ts index a1d431b65..4a3c5efe5 100644 --- a/apps/webapp/app/v3/services/triggerTaskV2.server.ts +++ b/apps/webapp/app/v3/services/triggerTaskV2.server.ts @@ -167,10 +167,10 @@ export class TriggerTaskServiceV2 extends WithRunEngine { index: options.batchIndex ?? 0, } : undefined, - environmentId: environment.id, projectId: environment.projectId, organizationId: environment.organizationId, tx: this._prisma, + releaseConcurrency: body.options?.releaseConcurrency, }); } ); @@ -373,6 +373,7 @@ export class TriggerTaskServiceV2 extends WithRunEngine { : undefined, machine: body.options?.machine, priorityMs: body.options?.priority ? body.options.priority * 1_000 : undefined, + releaseConcurrency: body.options?.releaseConcurrency, }, this._prisma ); diff --git a/internal-packages/database/prisma/migrations/20250319103257_add_release_concurrency_on_waitpoint_to_task_queue/migration.sql b/internal-packages/database/prisma/migrations/20250319103257_add_release_concurrency_on_waitpoint_to_task_queue/migration.sql new file mode 100644 index 000000000..66cea8acd --- /dev/null +++ b/internal-packages/database/prisma/migrations/20250319103257_add_release_concurrency_on_waitpoint_to_task_queue/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable +ALTER TABLE + "TaskQueue" +ADD + COLUMN "releaseConcurrencyOnWaitpoint" BOOLEAN NOT NULL DEFAULT false; \ No newline at end of file diff --git a/internal-packages/database/prisma/migrations/20250319110754_add_org_and_project_to_execution_snapshots/migration.sql b/internal-packages/database/prisma/migrations/20250319110754_add_org_and_project_to_execution_snapshots/migration.sql new file mode 100644 index 000000000..afdb979e8 --- /dev/null +++ b/internal-packages/database/prisma/migrations/20250319110754_add_org_and_project_to_execution_snapshots/migration.sql @@ -0,0 +1,26 @@ +/* + Warnings: + + - Added the required column `organizationId` to the `TaskRunExecutionSnapshot` table without a default value. This is not possible if the table is not empty. + - Added the required column `projectId` to the `TaskRunExecutionSnapshot` table without a default value. This is not possible if the table is not empty. + + */ +-- AlterTable +ALTER TABLE + "TaskRunExecutionSnapshot" +ADD + COLUMN "organizationId" TEXT NOT NULL, +ADD + COLUMN "projectId" TEXT NOT NULL; + +-- AddForeignKey +ALTER TABLE + "TaskRunExecutionSnapshot" +ADD + CONSTRAINT "TaskRunExecutionSnapshot_projectId_fkey" FOREIGN KEY ("projectId") REFERENCES "Project"("id") ON DELETE RESTRICT ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE + "TaskRunExecutionSnapshot" +ADD + CONSTRAINT "TaskRunExecutionSnapshot_organizationId_fkey" FOREIGN KEY ("organizationId") REFERENCES "Organization"("id") ON DELETE RESTRICT ON UPDATE CASCADE; \ No newline at end of file diff --git a/internal-packages/database/prisma/migrations/20250319114436_add_metadata_to_task_run_execution_snapshots/migration.sql b/internal-packages/database/prisma/migrations/20250319114436_add_metadata_to_task_run_execution_snapshots/migration.sql new file mode 100644 index 000000000..d4121ed92 --- /dev/null +++ b/internal-packages/database/prisma/migrations/20250319114436_add_metadata_to_task_run_execution_snapshots/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable +ALTER TABLE + "TaskRunExecutionSnapshot" +ADD + COLUMN "metadata" JSONB; \ No newline at end of file diff --git a/internal-packages/database/prisma/migrations/20250319131807_add_locked_queue_id_to_task_run/migration.sql b/internal-packages/database/prisma/migrations/20250319131807_add_locked_queue_id_to_task_run/migration.sql new file mode 100644 index 000000000..b1bf829b7 --- /dev/null +++ b/internal-packages/database/prisma/migrations/20250319131807_add_locked_queue_id_to_task_run/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable +ALTER TABLE + "TaskRun" +ADD + COLUMN "lockedQueueId" TEXT; \ No newline at end of file diff --git a/internal-packages/database/prisma/schema.prisma b/internal-packages/database/prisma/schema.prisma index d4b9cb011..8b92c6448 100644 --- a/internal-packages/database/prisma/schema.prisma +++ b/internal-packages/database/prisma/schema.prisma @@ -164,6 +164,7 @@ model Organization { organizationIntegrations OrganizationIntegration[] workerGroups WorkerInstanceGroup[] workerInstances WorkerInstance[] + executionSnapshots TaskRunExecutionSnapshot[] } model ExternalAccount { @@ -504,6 +505,7 @@ model Project { waitpoints Waitpoint[] taskRunWaitpoints TaskRunWaitpoint[] taskRunCheckpoints TaskRunCheckpoint[] + executionSnapshots TaskRunExecutionSnapshot[] } enum ProjectVersion { @@ -1724,7 +1726,9 @@ model TaskRun { projectId String // The specific queue this run is in - queue String + queue String + // The queueId is set when the run is locked to a specific queue + lockedQueueId String? /// The main queue that this run is part of masterQueue String @default("main") @@ -1985,6 +1989,12 @@ model TaskRunExecutionSnapshot { environment RuntimeEnvironment @relation(fields: [environmentId], references: [id]) environmentType RuntimeEnvironmentType + projectId String + project Project @relation(fields: [projectId], references: [id]) + + organizationId String + organization Organization @relation(fields: [organizationId], references: [id]) + /// Waitpoints that have been completed for this execution completedWaitpoints Waitpoint[] @relation("completedWaitpoints") @@ -2006,6 +2016,9 @@ model TaskRunExecutionSnapshot { lastHeartbeatAt DateTime? + /// Metadata used by various systems in the run engine + metadata Json? + /// Used to get the latest valid snapshot quickly @@index([runId, isValid, createdAt(sort: Desc)]) } @@ -2531,6 +2544,9 @@ model TaskQueue { paused Boolean @default(false) + /// If true, when a run is paused and waiting for waitpoints to be completed, the run will release the concurrency capacity. + releaseConcurrencyOnWaitpoint Boolean @default(false) + createdAt DateTime @default(now()) updatedAt DateTime @updatedAt diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts index 97be60793..210a3b766 100644 --- a/internal-packages/run-engine/src/engine/index.ts +++ b/internal-packages/run-engine/src/engine/index.ts @@ -36,7 +36,6 @@ import { import { EventBus, EventBusEvents } from "./eventBus.js"; import { RunLocker } from "./locking.js"; import { ReleaseConcurrencyTokenBucketQueue } from "./releaseConcurrencyTokenBucketQueue.js"; -import { canReleaseConcurrency } from "./statuses.js"; import { BatchSystem } from "./systems/batchSystem.js"; import { CheckpointSystem } from "./systems/checkpointSystem.js"; import { DelayedRunSystem } from "./systems/delayedRunSystem.js"; @@ -46,13 +45,14 @@ import { ExecutionSnapshotSystem, getLatestExecutionSnapshot, } from "./systems/executionSnapshotSystem.js"; +import { ReleaseConcurrencySystem } from "./systems/releaseConcurrencySystem.js"; import { RunAttemptSystem } from "./systems/runAttemptSystem.js"; import { SystemResources } from "./systems/systems.js"; import { TtlSystem } from "./systems/ttlSystem.js"; +import { WaitingForWorkerSystem } from "./systems/waitingForWorkerSystem.js"; import { WaitpointSystem } from "./systems/waitpointSystem.js"; import { EngineWorker, HeartbeatTimeouts, RunEngineOptions, TriggerParams } from "./types.js"; import { workerCatalog } from "./workerCatalog.js"; -import { WaitingForWorkerSystem } from "./systems/waitingForWorkerSystem.js"; export class RunEngine { private runLockRedis: Redis; @@ -63,7 +63,7 @@ export class RunEngine { private logger = new Logger("RunEngine", "debug"); private tracer: Tracer; private heartbeatTimeouts: HeartbeatTimeouts; - private releaseConcurrencyQueue: ReleaseConcurrencyTokenBucketQueue<{ + releaseConcurrencyQueue: ReleaseConcurrencyTokenBucketQueue<{ orgId: string; projectId: string; envId: string; @@ -79,6 +79,7 @@ export class RunEngine { delayedRunSystem: DelayedRunSystem; ttlSystem: TtlSystem; waitingForWorkerSystem: WaitingForWorkerSystem; + releaseConcurrencySystem: ReleaseConcurrencySystem; constructor(private readonly options: RunEngineOptions) { this.prisma = options.prisma; @@ -188,7 +189,7 @@ export class RunEngine { redis: { ...options.queue.redis, // Use base queue redis options ...options.releaseConcurrency?.redis, // Allow overrides - keyPrefix: `${options.queue.redis.keyPrefix}release-concurrency:`, + keyPrefix: `${options.queue.redis.keyPrefix ?? ""}release-concurrency:`, }, retry: { maxRetries: options.releaseConcurrency?.maxRetries ?? 5, @@ -201,8 +202,8 @@ export class RunEngine { consumersCount: options.releaseConcurrency?.consumersCount ?? 1, pollInterval: options.releaseConcurrency?.pollInterval ?? 1000, batchSize: options.releaseConcurrency?.batchSize ?? 10, - executor: async (descriptor, runId) => { - await this.#executeReleasedConcurrencyFromQueue(descriptor, runId); + executor: async (descriptor, snapshotId) => { + await this.releaseConcurrencySystem.executeReleaseConcurrencyForSnapshot(snapshotId); }, maxTokens: async (descriptor) => { const environment = await this.prisma.runtimeEnvironment.findFirstOrThrow({ @@ -239,6 +240,10 @@ export class RunEngine { releaseConcurrencyQueue: this.releaseConcurrencyQueue, }; + this.releaseConcurrencySystem = new ReleaseConcurrencySystem({ + resources, + }); + this.executionSnapshotSystem = new ExecutionSnapshotSystem({ resources, heartbeatTimeouts: this.heartbeatTimeouts, @@ -251,6 +256,7 @@ export class RunEngine { this.checkpointSystem = new CheckpointSystem({ resources, + releaseConcurrencySystem: this.releaseConcurrencySystem, executionSnapshotSystem: this.executionSnapshotSystem, enqueueSystem: this.enqueueSystem, }); @@ -269,6 +275,7 @@ export class RunEngine { resources, executionSnapshotSystem: this.executionSnapshotSystem, enqueueSystem: this.enqueueSystem, + releaseConcurrencySystem: this.releaseConcurrencySystem, }); this.ttlSystem = new TtlSystem({ @@ -344,6 +351,7 @@ export class RunEngine { machine, workerId, runnerId, + releaseConcurrency, }: TriggerParams, tx?: PrismaClientOrTransaction ): Promise { @@ -435,6 +443,8 @@ export class RunEngine { runStatus: status, environmentId: environment.id, environmentType: environment.type, + projectId: environment.project.id, + organizationId: environment.organization.id, workerId, runnerId, }, @@ -490,12 +500,11 @@ export class RunEngine { runId: parentTaskRunId, waitpoints: associatedWaitpoint.id, projectId: associatedWaitpoint.projectId, - organizationId: environment.organization.id, batch, workerId, runnerId, tx: prisma, - releaseConcurrency: true, // TODO: This needs to use the release concurrency system + releaseConcurrency, }); } @@ -1015,7 +1024,6 @@ export class RunEngine { runId, waitpoints, projectId, - organizationId, releaseConcurrency, timeout, spanIdToComplete, @@ -1040,7 +1048,6 @@ export class RunEngine { runId, waitpoints, projectId, - organizationId, releaseConcurrency, timeout, spanIdToComplete, @@ -1051,35 +1058,6 @@ export class RunEngine { }); } - async #executeReleasedConcurrencyFromQueue( - descriptor: { orgId: string; projectId: string; envId: string }, - runId: string - ) { - this.logger.debug("Executing released concurrency", { - descriptor, - runId, - }); - - // - Runlock the run - // - Get latest snapshot - // - If the run is non suspended or going to be, then bail - // - If the run is suspended or going to be, then release the concurrency - await this.runLock.lock([runId], 5_000, async () => { - const snapshot = await getLatestExecutionSnapshot(this.prisma, runId); - - if (!canReleaseConcurrency(snapshot.executionStatus)) { - this.logger.debug("Run is not in a state to release concurrency", { - runId, - snapshot, - }); - - return; - } - - return await this.runQueue.releaseConcurrency(descriptor.orgId, snapshot.runId); - }); - } - /** This completes a waitpoint and updates all entries so the run isn't blocked, * if they're no longer blocked. This doesn't suffer from race conditions. */ async completeWaitpoint({ @@ -1340,7 +1318,8 @@ export class RunEngine { id: latestSnapshot.environmentId, type: latestSnapshot.environmentType, }, - orgId: run.runtimeEnvironment.organizationId, + orgId: latestSnapshot.organizationId, + projectId: latestSnapshot.projectId, error: { type: "INTERNAL_ERROR", code: "TASK_RUN_DEQUEUED_MAX_RETRIES", diff --git a/internal-packages/run-engine/src/engine/releaseConcurrencyTokenBucketQueue.ts b/internal-packages/run-engine/src/engine/releaseConcurrencyTokenBucketQueue.ts index 212188fb5..fcdfb774e 100644 --- a/internal-packages/run-engine/src/engine/releaseConcurrencyTokenBucketQueue.ts +++ b/internal-packages/run-engine/src/engine/releaseConcurrencyTokenBucketQueue.ts @@ -114,9 +114,59 @@ export class ReleaseConcurrencyTokenBucketQueue { retryCount: 0, lastAttempt: Date.now(), }); + } else { + this.logger.info("No token available, adding to queue", { + releaseQueueDescriptor, + releaserId, + maxTokens, + }); } } + /** + * Consume a token from the token bucket for a release queue. + * + * This is mainly used for testing purposes + */ + public async consumeToken(releaseQueueDescriptor: T, releaserId: string) { + const maxTokens = await this.#callMaxTokens(releaseQueueDescriptor); + + if (maxTokens === 0) { + return; + } + + const releaseQueue = this.keys.fromDescriptor(releaseQueueDescriptor); + + await this.redis.consumeToken( + this.masterQueuesKey, + this.#bucketKey(releaseQueue), + this.#queueKey(releaseQueue), + this.#metadataKey(releaseQueue), + releaseQueue, + releaserId, + String(maxTokens), + String(Date.now()) + ); + } + + /** + * Return a token to the token bucket for a release queue. + * + * This is mainly used for testing purposes + */ + public async returnToken(releaseQueueDescriptor: T, releaserId: string) { + const releaseQueue = this.keys.fromDescriptor(releaseQueueDescriptor); + + await this.redis.returnTokenOnly( + this.masterQueuesKey, + this.#bucketKey(releaseQueue), + this.#queueKey(releaseQueue), + this.#metadataKey(releaseQueue), + releaseQueue, + releaserId + ); + } + /** * Refill the token bucket for a release queue. * @@ -384,7 +434,7 @@ local queueKey = keyPrefix .. queueName .. ":queue" local metadataKey = keyPrefix .. queueName .. ":metadata" -- Get the oldest item from the queue -local items = redis.call("ZRANGEBYSCORE", queueKey, 0, currentTime, "LIMIT", 0, batchSize - 1) +local items = redis.call("ZRANGEBYSCORE", queueKey, 0, currentTime, "LIMIT", 0, batchSize) if #items == 0 then -- No items ready to be processed yet return nil diff --git a/internal-packages/run-engine/src/engine/retrying.ts b/internal-packages/run-engine/src/engine/retrying.ts index f214738ad..a621552e9 100644 --- a/internal-packages/run-engine/src/engine/retrying.ts +++ b/internal-packages/run-engine/src/engine/retrying.ts @@ -10,7 +10,7 @@ import { } from "@trigger.dev/core/v3"; import { PrismaClientOrTransaction } from "@trigger.dev/database"; import { MAX_TASK_RUN_ATTEMPTS } from "./consts.js"; -import { ServiceValidationError } from "./index.js"; +import { ServiceValidationError } from "./errors.js"; type Params = { runId: string; diff --git a/internal-packages/run-engine/src/engine/systems/batchSystem.ts b/internal-packages/run-engine/src/engine/systems/batchSystem.ts index 0c09256eb..5f1948a83 100644 --- a/internal-packages/run-engine/src/engine/systems/batchSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/batchSystem.ts @@ -1,8 +1,5 @@ -import { Tracer, startSpan } from "@internal/tracing"; -import { Logger } from "@trigger.dev/core/logger"; -import { PrismaClient } from "@trigger.dev/database"; +import { startSpan } from "@internal/tracing"; import { isFinalRunStatus } from "../statuses.js"; -import { EngineWorker } from "../types.js"; import { SystemResources } from "./systems.js"; export type BatchSystemOptions = { diff --git a/internal-packages/run-engine/src/engine/systems/checkpointSystem.ts b/internal-packages/run-engine/src/engine/systems/checkpointSystem.ts index ae38de434..de06fca52 100644 --- a/internal-packages/run-engine/src/engine/systems/checkpointSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/checkpointSystem.ts @@ -11,22 +11,25 @@ import { import { SystemResources } from "./systems.js"; import { ServiceValidationError } from "../errors.js"; import { EnqueueSystem } from "./enqueueSystem.js"; - +import { ReleaseConcurrencySystem } from "./releaseConcurrencySystem.js"; export type CheckpointSystemOptions = { resources: SystemResources; executionSnapshotSystem: ExecutionSnapshotSystem; enqueueSystem: EnqueueSystem; + releaseConcurrencySystem: ReleaseConcurrencySystem; }; export class CheckpointSystem { private readonly $: SystemResources; private readonly executionSnapshotSystem: ExecutionSnapshotSystem; private readonly enqueueSystem: EnqueueSystem; + private readonly releaseConcurrencySystem: ReleaseConcurrencySystem; constructor(private readonly options: CheckpointSystemOptions) { this.$ = options.resources; this.executionSnapshotSystem = options.executionSnapshotSystem; this.enqueueSystem = options.enqueueSystem; + this.releaseConcurrencySystem = options.releaseConcurrencySystem; } /** @@ -163,6 +166,7 @@ export class CheckpointSystem { status: "QUEUED", description: "Run was QUEUED, because it was queued and executing and a checkpoint was created", + metadata: snapshot.metadata, }, previousSnapshotId: snapshot.id, batchId: snapshot.batchId ?? undefined, @@ -174,14 +178,7 @@ export class CheckpointSystem { }); // Refill the token bucket for the release concurrency queue - await this.$.releaseConcurrencyQueue.refillTokens( - { - orgId: run.runtimeEnvironment.organizationId, - projectId: run.runtimeEnvironment.projectId, - envId: run.runtimeEnvironment.id, - }, - 1 - ); + await this.releaseConcurrencySystem.checkpointCreatedOnEnvironment(run.runtimeEnvironment); return { ok: true as const, @@ -195,24 +192,20 @@ export class CheckpointSystem { snapshot: { executionStatus: "SUSPENDED", description: "Run was suspended after creating a checkpoint.", + metadata: snapshot.metadata, }, previousSnapshotId: snapshot.id, environmentId: snapshot.environmentId, environmentType: snapshot.environmentType, + projectId: snapshot.projectId, + organizationId: snapshot.organizationId, checkpointId: taskRunCheckpoint.id, workerId, runnerId, }); // Refill the token bucket for the release concurrency queue - await this.$.releaseConcurrencyQueue.refillTokens( - { - orgId: run.runtimeEnvironment.organizationId, - projectId: run.runtimeEnvironment.projectId, - envId: run.runtimeEnvironment.id, - }, - 1 - ); + await this.releaseConcurrencySystem.checkpointCreatedOnEnvironment(run.runtimeEnvironment); return { ok: true as const, @@ -284,6 +277,8 @@ export class CheckpointSystem { previousSnapshotId: snapshot.id, environmentId: snapshot.environmentId, environmentType: snapshot.environmentType, + projectId: snapshot.projectId, + organizationId: snapshot.organizationId, completedWaitpoints: snapshot.completedWaitpoints, workerId, runnerId, diff --git a/internal-packages/run-engine/src/engine/systems/delayedRunSystem.ts b/internal-packages/run-engine/src/engine/systems/delayedRunSystem.ts index bb2aaf308..c954a8d7e 100644 --- a/internal-packages/run-engine/src/engine/systems/delayedRunSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/delayedRunSystem.ts @@ -59,6 +59,8 @@ export class DelayedRunSystem { runStatus: "EXPIRED", environmentId: snapshot.environmentId, environmentType: snapshot.environmentType, + projectId: snapshot.projectId, + organizationId: snapshot.organizationId, }, }, }, diff --git a/internal-packages/run-engine/src/engine/systems/dequeueSystem.ts b/internal-packages/run-engine/src/engine/systems/dequeueSystem.ts index 3a976d897..33bdb5656 100644 --- a/internal-packages/run-engine/src/engine/systems/dequeueSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/dequeueSystem.ts @@ -105,6 +105,8 @@ export class DequeueSystem { previousSnapshotId: snapshot.id, environmentId: snapshot.environmentId, environmentType: snapshot.environmentType, + projectId: snapshot.projectId, + organizationId: snapshot.organizationId, checkpointId: snapshot.checkpointId ?? undefined, completedWaitpoints: snapshot.completedWaitpoints, error: `Tried to dequeue a run that is not in a valid state to be dequeued.`, @@ -146,6 +148,8 @@ export class DequeueSystem { previousSnapshotId: snapshot.id, environmentId: snapshot.environmentId, environmentType: snapshot.environmentType, + projectId: snapshot.projectId, + organizationId: snapshot.organizationId, batchId: snapshot.batchId ?? undefined, completedWaitpoints: snapshot.completedWaitpoints.map((waitpoint) => ({ id: waitpoint.id, @@ -337,6 +341,42 @@ export class DequeueSystem { maxAttempts = parsedConfig.data.maxAttempts; } + const queue = await prisma.taskQueue.findUnique({ + where: { + runtimeEnvironmentId_name: { + runtimeEnvironmentId: result.run.runtimeEnvironmentId, + name: sanitizeQueueName(result.run.queue), + }, + }, + }); + + if (!queue) { + this.$.logger.debug( + "RunEngine.dequeueFromMasterQueue(): queue not found, so nacking message", + { + queueMessage: message, + taskRunQueue: result.run.queue, + runtimeEnvironmentId: result.run.runtimeEnvironmentId, + } + ); + + //will auto-retry + const gotRequeued = await this.$.runQueue.nackMessage({ orgId, messageId: runId }); + if (!gotRequeued) { + await this.runAttemptSystem.systemFailure({ + runId, + error: { + type: "INTERNAL_ERROR", + code: "TASK_DEQUEUED_QUEUE_NOT_FOUND", + message: `Tried to dequeue the run but the queue doesn't exist: ${result.run.queue}`, + }, + tx: prisma, + }); + } + + return null; + } + //update the run const lockedTaskRun = await prisma.taskRun.update({ where: { @@ -346,6 +386,7 @@ export class DequeueSystem { lockedAt: new Date(), lockedById: result.task.id, lockedToVersionId: result.worker.id, + lockedQueueId: queue.id, startedAt: result.run.startedAt ?? new Date(), baseCostInCents: this.options.machines.baseCostInCents, machinePreset: machinePreset.name, @@ -378,42 +419,6 @@ export class DequeueSystem { return null; } - const queue = await prisma.taskQueue.findUnique({ - where: { - runtimeEnvironmentId_name: { - runtimeEnvironmentId: lockedTaskRun.runtimeEnvironmentId, - name: sanitizeQueueName(lockedTaskRun.queue), - }, - }, - }); - - if (!queue) { - this.$.logger.debug( - "RunEngine.dequeueFromMasterQueue(): queue not found, so nacking message", - { - queueMessage: message, - taskRunQueue: lockedTaskRun.queue, - runtimeEnvironmentId: lockedTaskRun.runtimeEnvironmentId, - } - ); - - //will auto-retry - const gotRequeued = await this.$.runQueue.nackMessage({ orgId, messageId: runId }); - if (!gotRequeued) { - await this.runAttemptSystem.systemFailure({ - runId, - error: { - type: "INTERNAL_ERROR", - code: "TASK_DEQUEUED_QUEUE_NOT_FOUND", - message: `Tried to dequeue the run but the queue doesn't exist: ${lockedTaskRun.queue}`, - }, - tx: prisma, - }); - } - - return null; - } - const currentAttemptNumber = lockedTaskRun.attemptNumber ?? 0; const nextAttemptNumber = currentAttemptNumber + 1; @@ -432,6 +437,8 @@ export class DequeueSystem { previousSnapshotId: snapshot.id, environmentId: snapshot.environmentId, environmentType: snapshot.environmentType, + projectId: snapshot.projectId, + organizationId: snapshot.organizationId, checkpointId: snapshot.checkpointId ?? undefined, completedWaitpoints: snapshot.completedWaitpoints, workerId, @@ -519,6 +526,7 @@ export class DequeueSystem { run, environment: run.runtimeEnvironment, orgId, + projectId: run.runtimeEnvironment.projectId, error: { type: "INTERNAL_ERROR", code: "TASK_RUN_DEQUEUED_MAX_RETRIES", @@ -572,7 +580,12 @@ export class DequeueSystem { status: true, attemptNumber: true, runtimeEnvironment: { - select: { id: true, type: true }, + select: { + id: true, + type: true, + projectId: true, + project: { select: { id: true, organizationId: true } }, + }, }, }, }); @@ -587,6 +600,8 @@ export class DequeueSystem { }, environmentId: run.runtimeEnvironment.id, environmentType: run.runtimeEnvironment.type, + projectId: run.runtimeEnvironment.projectId, + organizationId: run.runtimeEnvironment.project.organizationId, workerId, runnerId, }); diff --git a/internal-packages/run-engine/src/engine/systems/enqueueSystem.ts b/internal-packages/run-engine/src/engine/systems/enqueueSystem.ts index 842bbeed8..0ed309792 100644 --- a/internal-packages/run-engine/src/engine/systems/enqueueSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/enqueueSystem.ts @@ -1,4 +1,9 @@ -import { PrismaClientOrTransaction, TaskRun, TaskRunExecutionStatus } from "@trigger.dev/database"; +import { + Prisma, + PrismaClientOrTransaction, + TaskRun, + TaskRunExecutionStatus, +} from "@trigger.dev/database"; import { MinimalAuthenticatedEnvironment } from "../../shared/index.js"; import { ExecutionSnapshotSystem } from "./executionSnapshotSystem.js"; import { SystemResources } from "./systems.js"; @@ -37,6 +42,7 @@ export class EnqueueSystem { snapshot?: { status?: Extract; description?: string; + metadata?: Prisma.JsonValue; }; previousSnapshotId?: string; batchId?: string; @@ -56,11 +62,14 @@ export class EnqueueSystem { snapshot: { executionStatus: snapshot?.status ?? "QUEUED", description: snapshot?.description ?? "Run was QUEUED", + metadata: snapshot?.metadata ?? undefined, }, previousSnapshotId, batchId, environmentId: env.id, environmentType: env.type, + projectId: env.project.id, + organizationId: env.organization.id, checkpointId, completedWaitpoints, workerId, diff --git a/internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts b/internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts index 2abc64051..25320697b 100644 --- a/internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/executionSnapshotSystem.ts @@ -1,6 +1,7 @@ import { CompletedWaitpoint, ExecutionResult } from "@trigger.dev/core/v3"; import { BatchId, RunId, SnapshotId } from "@trigger.dev/core/v3/isomorphic"; import { + Prisma, PrismaClientOrTransaction, RuntimeEnvironmentType, TaskRunCheckpoint, @@ -158,6 +159,8 @@ export class ExecutionSnapshotSystem { batchId, environmentId, environmentType, + projectId, + organizationId, checkpointId, workerId, runnerId, @@ -168,11 +171,14 @@ export class ExecutionSnapshotSystem { snapshot: { executionStatus: TaskRunExecutionStatus; description: string; + metadata?: Prisma.JsonValue; }; previousSnapshotId?: string; batchId?: string; environmentId: string; environmentType: RuntimeEnvironmentType; + projectId: string; + organizationId: string; checkpointId?: string; workerId?: string; runnerId?: string; @@ -195,9 +201,12 @@ export class ExecutionSnapshotSystem { batchId, environmentId, environmentType, + projectId, + organizationId, checkpointId, workerId, runnerId, + metadata: snapshot.metadata ?? undefined, completedWaitpoints: { connect: completedWaitpoints?.map((w) => ({ id: w.id })), }, diff --git a/internal-packages/run-engine/src/engine/systems/releaseConcurrencySystem.ts b/internal-packages/run-engine/src/engine/systems/releaseConcurrencySystem.ts new file mode 100644 index 000000000..cf29e115c --- /dev/null +++ b/internal-packages/run-engine/src/engine/systems/releaseConcurrencySystem.ts @@ -0,0 +1,161 @@ +import { RuntimeEnvironment, TaskRunExecutionSnapshot } from "@trigger.dev/database"; +import { SystemResources } from "./systems.js"; +import { getLatestExecutionSnapshot } from "./executionSnapshotSystem.js"; +import { canReleaseConcurrency } from "../statuses.js"; +import { z } from "zod"; + +const ReleaseConcurrencyMetadata = z.object({ + releaseConcurrency: z.boolean().optional(), +}); + +type ReleaseConcurrencyMetadata = z.infer; + +export type ReleaseConcurrencySystemOptions = { + resources: SystemResources; +}; + +export class ReleaseConcurrencySystem { + private readonly $: SystemResources; + + constructor(private readonly options: ReleaseConcurrencySystemOptions) { + this.$ = options.resources; + } + + public async checkpointCreatedOnEnvironment(environment: RuntimeEnvironment) { + await this.$.releaseConcurrencyQueue.refillTokens( + { + orgId: environment.organizationId, + projectId: environment.projectId, + envId: environment.id, + }, + 1 + ); + } + + public async releaseConcurrencyForSnapshot(snapshot: TaskRunExecutionSnapshot) { + // Go ahead and release concurrency immediately if the run is in a development environment + if (snapshot.environmentType === "DEVELOPMENT") { + return await this.executeReleaseConcurrencyForSnapshot(snapshot.id); + } + + await this.$.releaseConcurrencyQueue.attemptToRelease( + { + orgId: snapshot.organizationId, + projectId: snapshot.projectId, + envId: snapshot.environmentId, + }, + snapshot.id + ); + } + + public async executeReleaseConcurrencyForSnapshot(snapshotId: string) { + this.$.logger.debug("Executing released concurrency", { + snapshotId, + }); + + // Fetch the snapshot + const snapshot = await this.$.prisma.taskRunExecutionSnapshot.findFirst({ + where: { id: snapshotId }, + select: { + id: true, + previousSnapshotId: true, + executionStatus: true, + organizationId: true, + metadata: true, + runId: true, + run: { + select: { + lockedQueueId: true, + }, + }, + }, + }); + + if (!snapshot) { + this.$.logger.error("Snapshot not found", { + snapshotId, + }); + + return; + } + + // - Runlock the run + // - Get latest snapshot + // - If the run is non suspended or going to be, then bail + // - If the run is suspended or going to be, then release the concurrency + await this.$.runLock.lock([snapshot.runId], 5_000, async () => { + const latestSnapshot = await getLatestExecutionSnapshot(this.$.prisma, snapshot.runId); + + const isValidSnapshot = + latestSnapshot.id === snapshot.id || + // Case 2: The provided snapshotId matches the previous snapshot + // AND we're in SUSPENDED state (which is valid) + (latestSnapshot.previousSnapshotId === snapshot.id && + latestSnapshot.executionStatus === "SUSPENDED"); + + if (!isValidSnapshot) { + this.$.logger.error("Tried to release concurrency on an invalid snapshot", { + latestSnapshot, + snapshot, + }); + + return; + } + + if (!canReleaseConcurrency(latestSnapshot.executionStatus)) { + this.$.logger.debug("Run is not in a state to release concurrency", { + runId: snapshot.runId, + snapshot: latestSnapshot, + }); + + return; + } + + const metadata = this.#parseMetadata(snapshot.metadata); + + if (typeof metadata.releaseConcurrency === "boolean") { + if (metadata.releaseConcurrency) { + return await this.$.runQueue.releaseAllConcurrency( + snapshot.organizationId, + snapshot.runId + ); + } + + return await this.$.runQueue.releaseEnvConcurrency(snapshot.organizationId, snapshot.runId); + } + + // Get the locked queue + const taskQueue = snapshot.run.lockedQueueId + ? await this.$.prisma.taskQueue.findFirst({ + where: { + id: snapshot.run.lockedQueueId, + }, + }) + : undefined; + + if ( + taskQueue && + (typeof taskQueue.concurrencyLimit === "undefined" || + taskQueue.releaseConcurrencyOnWaitpoint) + ) { + return await this.$.runQueue.releaseAllConcurrency(snapshot.organizationId, snapshot.runId); + } + + return await this.$.runQueue.releaseEnvConcurrency(snapshot.organizationId, snapshot.runId); + }); + } + + #parseMetadata(metadata?: unknown): ReleaseConcurrencyMetadata { + if (!metadata) { + return {}; + } + + const result = ReleaseConcurrencyMetadata.safeParse(metadata); + + if (!result.success) { + return {}; + } + + return result.data; + } +} diff --git a/internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts b/internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts index d6e6fd6ec..0ba498d77 100644 --- a/internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts @@ -216,6 +216,8 @@ export class RunAttemptSystem { previousSnapshotId: latestSnapshot.id, environmentId: latestSnapshot.environmentId, environmentType: latestSnapshot.environmentType, + projectId: latestSnapshot.projectId, + organizationId: latestSnapshot.organizationId, workerId, runnerId, }); @@ -436,6 +438,8 @@ export class RunAttemptSystem { attemptNumber: latestSnapshot.attemptNumber, environmentId: latestSnapshot.environmentId, environmentType: latestSnapshot.environmentType, + projectId: latestSnapshot.projectId, + organizationId: latestSnapshot.organizationId, workerId, runnerId, }, @@ -706,6 +710,7 @@ export class RunAttemptSystem { run, environment: run.runtimeEnvironment, orgId: run.runtimeEnvironment.organizationId, + projectId: run.runtimeEnvironment.project.id, timestamp: retryAt.getTime(), error: { type: "INTERNAL_ERROR", @@ -737,6 +742,8 @@ export class RunAttemptSystem { previousSnapshotId: latestSnapshot.id, environmentId: latestSnapshot.environmentId, environmentType: latestSnapshot.environmentType, + projectId: latestSnapshot.projectId, + organizationId: latestSnapshot.organizationId, workerId, runnerId, } @@ -820,6 +827,7 @@ export class RunAttemptSystem { run, environment, orgId, + projectId, timestamp, error, workerId, @@ -832,6 +840,7 @@ export class RunAttemptSystem { type: RuntimeEnvironmentType; }; orgId: string; + projectId: string; timestamp?: number; error: TaskRunInternalError; workerId?: string; @@ -865,6 +874,8 @@ export class RunAttemptSystem { }, environmentId: environment.id, environmentType: environment.type, + projectId: projectId, + organizationId: orgId, workerId, runnerId, }); @@ -988,6 +999,8 @@ export class RunAttemptSystem { previousSnapshotId: latestSnapshot.id, environmentId: latestSnapshot.environmentId, environmentType: latestSnapshot.environmentType, + projectId: latestSnapshot.projectId, + organizationId: latestSnapshot.organizationId, workerId, runnerId, }); @@ -1011,6 +1024,8 @@ export class RunAttemptSystem { previousSnapshotId: latestSnapshot.id, environmentId: latestSnapshot.environmentId, environmentType: latestSnapshot.environmentType, + projectId: latestSnapshot.projectId, + organizationId: latestSnapshot.organizationId, workerId, runnerId, }); @@ -1104,6 +1119,12 @@ export class RunAttemptSystem { id: true, type: true, organizationId: true, + project: { + select: { + id: true, + organizationId: true, + }, + }, }, }, taskEventStore: true, @@ -1121,6 +1142,8 @@ export class RunAttemptSystem { previousSnapshotId: snapshotId, environmentId: run.runtimeEnvironment.id, environmentType: run.runtimeEnvironment.type, + projectId: run.runtimeEnvironment.project.id, + organizationId: run.runtimeEnvironment.project.organizationId, workerId, runnerId, }); diff --git a/internal-packages/run-engine/src/engine/systems/ttlSystem.ts b/internal-packages/run-engine/src/engine/systems/ttlSystem.ts index da6a44b82..12910f463 100644 --- a/internal-packages/run-engine/src/engine/systems/ttlSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/ttlSystem.ts @@ -76,6 +76,8 @@ export class TtlSystem { runStatus: "EXPIRED", environmentId: snapshot.environmentId, environmentType: snapshot.environmentType, + projectId: snapshot.projectId, + organizationId: snapshot.organizationId, }, }, }, diff --git a/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts b/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts index 3d1a59dca..ee5d79895 100644 --- a/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/waitpointSystem.ts @@ -14,23 +14,26 @@ import { isExecuting } from "../statuses.js"; import { EnqueueSystem } from "./enqueueSystem.js"; import { ExecutionSnapshotSystem, getLatestExecutionSnapshot } from "./executionSnapshotSystem.js"; import { SystemResources } from "./systems.js"; +import { ReleaseConcurrencySystem } from "./releaseConcurrencySystem.js"; export type WaitpointSystemOptions = { resources: SystemResources; executionSnapshotSystem: ExecutionSnapshotSystem; enqueueSystem: EnqueueSystem; + releaseConcurrencySystem: ReleaseConcurrencySystem; }; export class WaitpointSystem { private readonly $: SystemResources; private readonly executionSnapshotSystem: ExecutionSnapshotSystem; - + private readonly releaseConcurrencySystem: ReleaseConcurrencySystem; private readonly enqueueSystem: EnqueueSystem; constructor(private readonly options: WaitpointSystemOptions) { this.$ = options.resources; this.executionSnapshotSystem = options.executionSnapshotSystem; this.enqueueSystem = options.enqueueSystem; + this.releaseConcurrencySystem = options.releaseConcurrencySystem; } public async clearBlockingWaitpoints({ @@ -326,7 +329,6 @@ export class WaitpointSystem { runId, waitpoints, projectId, - organizationId, releaseConcurrency, timeout, spanIdToComplete, @@ -338,7 +340,6 @@ export class WaitpointSystem { runId: string; waitpoints: string | string[]; projectId: string; - organizationId: string; releaseConcurrency?: boolean; timeout?: Date; spanIdToComplete?: string; @@ -378,7 +379,7 @@ export class WaitpointSystem { JOIN "Waitpoint" w ON w.id = i."waitpointId" WHERE w.status = 'PENDING';`; - const pendingCount = Number(insert.at(0)?.pending_count ?? 0); + const isRunBlocked = Number(insert.at(0)?.pending_count ?? 0) > 0; let newStatus: TaskRunExecutionStatus = "SUSPENDED"; if ( @@ -399,10 +400,15 @@ export class WaitpointSystem { snapshot: { executionStatus: newStatus, description: "Run was blocked by a waitpoint.", + metadata: { + releaseConcurrency, + }, }, previousSnapshotId: snapshot.id, environmentId: snapshot.environmentId, environmentType: snapshot.environmentType, + projectId: snapshot.projectId, + organizationId: snapshot.organizationId, batchId: batch?.id ?? snapshot.batchId ?? undefined, workerId, runnerId, @@ -428,7 +434,10 @@ export class WaitpointSystem { //no pending waitpoint, schedule unblocking the run //debounce if we're rapidly adding waitpoints - if (pendingCount === 0) { + if (isRunBlocked) { + //release concurrency + await this.releaseConcurrencySystem.releaseConcurrencyForSnapshot(snapshot); + } else { await this.$.worker.enqueue({ //this will debounce the call id: `continueRunIfUnblocked:${runId}`, @@ -437,11 +446,6 @@ export class WaitpointSystem { //in the near future availableAt: new Date(Date.now() + 50), }); - } else { - if (releaseConcurrency) { - //release concurrency - await this.#attemptToReleaseConcurrency(organizationId, snapshot); - } } return snapshot; @@ -515,6 +519,8 @@ export class WaitpointSystem { previousSnapshotId: snapshot.id, environmentId: snapshot.environmentId, environmentType: snapshot.environmentType, + projectId: snapshot.projectId, + organizationId: snapshot.organizationId, batchId: snapshot.batchId ?? undefined, completedWaitpoints: blockingWaitpoints.map((b) => ({ id: b.waitpoint.id, @@ -601,45 +607,4 @@ export class WaitpointSystem { }, }); } - - async #attemptToReleaseConcurrency(orgId: string, snapshot: TaskRunExecutionSnapshot) { - // Go ahead and release concurrency immediately if the run is in a development environment - if (snapshot.environmentType === "DEVELOPMENT") { - return await this.$.runQueue.releaseConcurrency(orgId, snapshot.runId); - } - - const run = await this.$.prisma.taskRun.findFirst({ - where: { - id: snapshot.runId, - }, - select: { - runtimeEnvironment: { - select: { - id: true, - projectId: true, - organizationId: true, - }, - }, - }, - }); - - if (!run) { - this.$.logger.error("Run not found for attemptToReleaseConcurrency", { - runId: snapshot.runId, - }); - - return; - } - - await this.$.releaseConcurrencyQueue.attemptToRelease( - { - orgId: run.runtimeEnvironment.organizationId, - projectId: run.runtimeEnvironment.projectId, - envId: run.runtimeEnvironment.id, - }, - snapshot.runId - ); - - return; - } } diff --git a/internal-packages/run-engine/src/engine/tests/releaseConcurrency.test.ts b/internal-packages/run-engine/src/engine/tests/releaseConcurrency.test.ts new file mode 100644 index 000000000..d97535145 --- /dev/null +++ b/internal-packages/run-engine/src/engine/tests/releaseConcurrency.test.ts @@ -0,0 +1,1094 @@ +import { + assertNonNullable, + containerTest, + setupAuthenticatedEnvironment, + setupBackgroundWorker, +} from "@internal/testcontainers"; +import { trace } from "@internal/tracing"; +import { RunEngine } from "../index.js"; +import { setTimeout } from "node:timers/promises"; + +vi.setConfig({ testTimeout: 60_000 }); + +describe("RunEngine Releasing Concurrency", () => { + containerTest("defaults to releasing env concurrency only", async ({ prisma, redisOptions }) => { + //create environment + const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + + const engine = new RunEngine({ + prisma, + worker: { + redis: redisOptions, + workers: 1, + tasksPerWorker: 10, + pollIntervalMs: 100, + }, + queue: { + redis: redisOptions, + }, + runLock: { + redis: redisOptions, + }, + machines: { + defaultMachine: "small-1x", + machines: { + "small-1x": { + name: "small-1x" as const, + cpu: 0.5, + memory: 0.5, + centsPerMs: 0.0001, + }, + }, + baseCostInCents: 0.0001, + }, + releaseConcurrency: { + maxTokensRatio: 1, + maxRetries: 3, + consumersCount: 1, + pollInterval: 500, + batchSize: 1, + }, + tracer: trace.getTracer("test", "0.0.0"), + }); + const taskIdentifier = "test-task"; + + await setupBackgroundWorker(prisma, authenticatedEnvironment, taskIdentifier); + + const run = await engine.trigger( + { + number: 1, + friendlyId: "run_p1234", + environment: authenticatedEnvironment, + taskIdentifier, + payload: "{}", + payloadType: "application/json", + context: {}, + traceContext: {}, + traceId: "t12345", + spanId: "s12345", + masterQueue: "main", + queueName: `task/${taskIdentifier}`, + isTest: false, + tags: [], + }, + prisma + ); + + const dequeued = await engine.dequeueFromMasterQueue({ + consumerId: "test_12345", + masterQueue: run.masterQueue, + maxRunCount: 10, + }); + + const queueConcurrency = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrency).toBe(1); + + const envConcurrency = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrency).toBe(1); + + // create an attempt + const attemptResult = await engine.startRunAttempt({ + runId: dequeued[0].run.id, + snapshotId: dequeued[0].snapshot.id, + }); + + expect(attemptResult.snapshot.executionStatus).toBe("EXECUTING"); + + // create a manual waitpoint + const result = await engine.createManualWaitpoint({ + environmentId: authenticatedEnvironment.id, + projectId: authenticatedEnvironment.projectId, + }); + + // Block the run, not specifying any release concurrency option + const executingWithWaitpointSnapshot = await engine.blockRunWithWaitpoint({ + runId: run.id, + waitpoints: result.waitpoint.id, + projectId: authenticatedEnvironment.projectId, + organizationId: authenticatedEnvironment.organizationId, + }); + + expect(executingWithWaitpointSnapshot.executionStatus).toBe("EXECUTING_WITH_WAITPOINTS"); + + // Now confirm the queue has the same concurrency as before + const queueConcurrencyAfter = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrencyAfter).toBe(1); + + // Now confirm the environment has a concurrency of 0 + const envConcurrencyAfter = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrencyAfter).toBe(0); + + await engine.completeWaitpoint({ + id: result.waitpoint.id, + }); + + await setTimeout(500); + + // Test that we've reacquired the queue concurrency + const queueConcurrencyAfterWaitpoint = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrencyAfterWaitpoint).toBe(1); + + // Test that we've reacquired the environment concurrency + const envConcurrencyAfterWaitpoint = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrencyAfterWaitpoint).toBe(1); + + // Now we are going to block with another waitpoint, this time specifiying we want to release the concurrency in the waitpoint + const result2 = await engine.createManualWaitpoint({ + environmentId: authenticatedEnvironment.id, + projectId: authenticatedEnvironment.projectId, + }); + + const executingWithWaitpointSnapshot2 = await engine.blockRunWithWaitpoint({ + runId: run.id, + waitpoints: result2.waitpoint.id, + projectId: authenticatedEnvironment.projectId, + organizationId: authenticatedEnvironment.organizationId, + releaseConcurrency: true, + }); + + expect(executingWithWaitpointSnapshot2.executionStatus).toBe("EXECUTING_WITH_WAITPOINTS"); + + // Test that we've released the queue concurrency + const queueConcurrencyAfterWaitpoint2 = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrencyAfterWaitpoint2).toBe(0); + + // Test that we've released the environment concurrency + const envConcurrencyAfterWaitpoint2 = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrencyAfterWaitpoint2).toBe(0); + + // Complete the waitpoint and make sure the run reacquires the queue and environment concurrency + await engine.completeWaitpoint({ + id: result2.waitpoint.id, + }); + + await setTimeout(500); + + // Test that we've reacquired the queue concurrency + const queueConcurrencyAfterWaitpoint3 = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrencyAfterWaitpoint3).toBe(1); + }); + + containerTest( + "releases all concurrency when configured on queue", + async ({ prisma, redisOptions }) => { + //create environment + const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + + const engine = new RunEngine({ + prisma, + worker: { + redis: redisOptions, + workers: 1, + tasksPerWorker: 10, + pollIntervalMs: 100, + }, + queue: { + redis: redisOptions, + }, + runLock: { + redis: redisOptions, + }, + machines: { + defaultMachine: "small-1x", + machines: { + "small-1x": { + name: "small-1x" as const, + cpu: 0.5, + memory: 0.5, + centsPerMs: 0.0001, + }, + }, + baseCostInCents: 0.0001, + }, + releaseConcurrency: { + maxTokensRatio: 1, + maxRetries: 3, + consumersCount: 1, + pollInterval: 500, + batchSize: 1, + }, + tracer: trace.getTracer("test", "0.0.0"), + }); + const taskIdentifier = "test-task"; + + await setupBackgroundWorker( + prisma, + authenticatedEnvironment, + taskIdentifier, + undefined, + undefined, + { + releaseConcurrencyOnWaitpoint: true, + } + ); + + const run = await engine.trigger( + { + number: 1, + friendlyId: "run_p1234", + environment: authenticatedEnvironment, + taskIdentifier, + payload: "{}", + payloadType: "application/json", + context: {}, + traceContext: {}, + traceId: "t12345", + spanId: "s12345", + masterQueue: "main", + queueName: `task/${taskIdentifier}`, + isTest: false, + tags: [], + }, + prisma + ); + + const dequeued = await engine.dequeueFromMasterQueue({ + consumerId: "test_12345", + masterQueue: run.masterQueue, + maxRunCount: 10, + }); + + const queueConcurrency = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrency).toBe(1); + + const envConcurrency = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrency).toBe(1); + + // create an attempt + const attemptResult = await engine.startRunAttempt({ + runId: dequeued[0].run.id, + snapshotId: dequeued[0].snapshot.id, + }); + + expect(attemptResult.snapshot.executionStatus).toBe("EXECUTING"); + + // create a manual waitpoint + const result = await engine.createManualWaitpoint({ + environmentId: authenticatedEnvironment.id, + projectId: authenticatedEnvironment.projectId, + }); + + // Block the run, not specifying any release concurrency option + const executingWithWaitpointSnapshot = await engine.blockRunWithWaitpoint({ + runId: run.id, + waitpoints: result.waitpoint.id, + projectId: authenticatedEnvironment.projectId, + organizationId: authenticatedEnvironment.organizationId, + }); + + expect(executingWithWaitpointSnapshot.executionStatus).toBe("EXECUTING_WITH_WAITPOINTS"); + + // Now confirm the queue has the same concurrency as before + const queueConcurrencyAfter = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrencyAfter).toBe(0); + + // Now confirm the environment has a concurrency of 0 + const envConcurrencyAfter = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrencyAfter).toBe(0); + + // Complete the waitpoint and make sure the run reacquires the queue and environment concurrency + await engine.completeWaitpoint({ + id: result.waitpoint.id, + }); + + await setTimeout(500); + + // Test that we've reacquired the queue concurrency + const queueConcurrencyAfterWaitpoint = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrencyAfterWaitpoint).toBe(1); + + // Test that we've reacquired the environment concurrency + const envConcurrencyAfterWaitpoint = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrencyAfterWaitpoint).toBe(1); + + // Now we are going to block with another waitpoint, this time specifiying we dont want to release the concurrency in the waitpoint + const result2 = await engine.createManualWaitpoint({ + environmentId: authenticatedEnvironment.id, + projectId: authenticatedEnvironment.projectId, + }); + + const executingWithWaitpointSnapshot2 = await engine.blockRunWithWaitpoint({ + runId: run.id, + waitpoints: result2.waitpoint.id, + projectId: authenticatedEnvironment.projectId, + organizationId: authenticatedEnvironment.organizationId, + releaseConcurrency: false, + }); + + expect(executingWithWaitpointSnapshot2.executionStatus).toBe("EXECUTING_WITH_WAITPOINTS"); + + // Test that we've not released the queue concurrency + const queueConcurrencyAfterWaitpoint2 = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrencyAfterWaitpoint2).toBe(1); + + // Test that we've still released the environment concurrency since we always release env concurrency + const envConcurrencyAfterWaitpoint2 = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrencyAfterWaitpoint2).toBe(0); + } + ); + + containerTest( + "releases all concurrency for unlimited queues", + async ({ prisma, redisOptions }) => { + //create environment + const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + + const engine = new RunEngine({ + prisma, + worker: { + redis: redisOptions, + workers: 1, + tasksPerWorker: 10, + pollIntervalMs: 100, + }, + queue: { + redis: redisOptions, + }, + runLock: { + redis: redisOptions, + }, + machines: { + defaultMachine: "small-1x", + machines: { + "small-1x": { + name: "small-1x" as const, + cpu: 0.5, + memory: 0.5, + centsPerMs: 0.0001, + }, + }, + baseCostInCents: 0.0001, + }, + releaseConcurrency: { + maxTokensRatio: 1, + maxRetries: 3, + consumersCount: 1, + pollInterval: 500, + batchSize: 1, + }, + tracer: trace.getTracer("test", "0.0.0"), + }); + const taskIdentifier = "test-task"; + + await setupBackgroundWorker( + prisma, + authenticatedEnvironment, + taskIdentifier, + undefined, + undefined, + { + releaseConcurrencyOnWaitpoint: true, + concurrencyLimit: null, + } + ); + + const run = await engine.trigger( + { + number: 1, + friendlyId: "run_p1234", + environment: authenticatedEnvironment, + taskIdentifier, + payload: "{}", + payloadType: "application/json", + context: {}, + traceContext: {}, + traceId: "t12345", + spanId: "s12345", + masterQueue: "main", + queueName: `task/${taskIdentifier}`, + isTest: false, + tags: [], + }, + prisma + ); + + const dequeued = await engine.dequeueFromMasterQueue({ + consumerId: "test_12345", + masterQueue: run.masterQueue, + maxRunCount: 10, + }); + + const queueConcurrency = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrency).toBe(1); + + const envConcurrency = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrency).toBe(1); + + // create an attempt + const attemptResult = await engine.startRunAttempt({ + runId: dequeued[0].run.id, + snapshotId: dequeued[0].snapshot.id, + }); + + expect(attemptResult.snapshot.executionStatus).toBe("EXECUTING"); + + // create a manual waitpoint + const result = await engine.createManualWaitpoint({ + environmentId: authenticatedEnvironment.id, + projectId: authenticatedEnvironment.projectId, + }); + + // Block the run, not specifying any release concurrency option + const executingWithWaitpointSnapshot = await engine.blockRunWithWaitpoint({ + runId: run.id, + waitpoints: result.waitpoint.id, + projectId: authenticatedEnvironment.projectId, + organizationId: authenticatedEnvironment.organizationId, + }); + + expect(executingWithWaitpointSnapshot.executionStatus).toBe("EXECUTING_WITH_WAITPOINTS"); + + // Now confirm the queue has the same concurrency as before + const queueConcurrencyAfter = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrencyAfter).toBe(0); + + // Now confirm the environment has a concurrency of 0 + const envConcurrencyAfter = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrencyAfter).toBe(0); + + // Complete the waitpoint and make sure the run reacquires the queue and environment concurrency + await engine.completeWaitpoint({ + id: result.waitpoint.id, + }); + + await setTimeout(500); + + // Test that we've reacquired the queue concurrency + const queueConcurrencyAfterWaitpoint = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrencyAfterWaitpoint).toBe(1); + + // Test that we've reacquired the environment concurrency + const envConcurrencyAfterWaitpoint = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrencyAfterWaitpoint).toBe(1); + + // Now we are going to block with another waitpoint, this time specifiying we dont want to release the concurrency in the waitpoint + const result2 = await engine.createManualWaitpoint({ + environmentId: authenticatedEnvironment.id, + projectId: authenticatedEnvironment.projectId, + }); + + const executingWithWaitpointSnapshot2 = await engine.blockRunWithWaitpoint({ + runId: run.id, + waitpoints: result2.waitpoint.id, + projectId: authenticatedEnvironment.projectId, + organizationId: authenticatedEnvironment.organizationId, + releaseConcurrency: false, + }); + + expect(executingWithWaitpointSnapshot2.executionStatus).toBe("EXECUTING_WITH_WAITPOINTS"); + + // Test that we've not released the queue concurrency + const queueConcurrencyAfterWaitpoint2 = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrencyAfterWaitpoint2).toBe(1); + + // Test that we've still released the environment concurrency since we always release env concurrency + const envConcurrencyAfterWaitpoint2 = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrencyAfterWaitpoint2).toBe(0); + } + ); + + containerTest( + "delays env concurrency release when token unavailable", + async ({ prisma, redisOptions }) => { + //create environment + const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + + const engine = new RunEngine({ + prisma, + worker: { + redis: redisOptions, + workers: 1, + tasksPerWorker: 10, + pollIntervalMs: 100, + }, + queue: { + redis: redisOptions, + }, + runLock: { + redis: redisOptions, + }, + machines: { + defaultMachine: "small-1x", + machines: { + "small-1x": { + name: "small-1x" as const, + cpu: 0.5, + memory: 0.5, + centsPerMs: 0.0001, + }, + }, + baseCostInCents: 0.0001, + }, + releaseConcurrency: { + maxTokensRatio: 0.1, // 10% of the concurrency limit = 1 token + maxRetries: 3, + consumersCount: 1, + pollInterval: 500, + batchSize: 1, + }, + tracer: trace.getTracer("test", "0.0.0"), + }); + const taskIdentifier = "test-task"; + + await setupBackgroundWorker(prisma, authenticatedEnvironment, taskIdentifier); + + const run = await engine.trigger( + { + number: 1, + friendlyId: "run_p1234", + environment: authenticatedEnvironment, + taskIdentifier, + payload: "{}", + payloadType: "application/json", + context: {}, + traceContext: {}, + traceId: "t12345", + spanId: "s12345", + masterQueue: "main", + queueName: `task/${taskIdentifier}`, + isTest: false, + tags: [], + }, + prisma + ); + + const dequeued = await engine.dequeueFromMasterQueue({ + consumerId: "test_12345", + masterQueue: run.masterQueue, + maxRunCount: 10, + }); + + const queueConcurrency = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrency).toBe(1); + + const envConcurrency = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrency).toBe(1); + + // create an attempt + const attemptResult = await engine.startRunAttempt({ + runId: dequeued[0].run.id, + snapshotId: dequeued[0].snapshot.id, + }); + + expect(attemptResult.snapshot.executionStatus).toBe("EXECUTING"); + + // create a manual waitpoint + const result = await engine.createManualWaitpoint({ + environmentId: authenticatedEnvironment.id, + projectId: authenticatedEnvironment.projectId, + }); + + await engine.releaseConcurrencyQueue.consumeToken( + { + orgId: authenticatedEnvironment.organizationId, + projectId: authenticatedEnvironment.projectId, + envId: authenticatedEnvironment.id, + }, + "test_12345" + ); + + // Block the run, not specifying any release concurrency option + const executingWithWaitpointSnapshot = await engine.blockRunWithWaitpoint({ + runId: run.id, + waitpoints: result.waitpoint.id, + projectId: authenticatedEnvironment.projectId, + organizationId: authenticatedEnvironment.organizationId, + }); + + expect(executingWithWaitpointSnapshot.executionStatus).toBe("EXECUTING_WITH_WAITPOINTS"); + + // Now confirm the queue has the same concurrency as before + const queueConcurrencyAfter = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrencyAfter).toBe(1); + + // Now confirm the environment is the same as before + const envConcurrencyAfter = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrencyAfter).toBe(1); + + // Now we return the token so the concurrency can be released + await engine.releaseConcurrencyQueue.returnToken( + { + orgId: authenticatedEnvironment.organizationId, + projectId: authenticatedEnvironment.projectId, + envId: authenticatedEnvironment.id, + }, + "test_12345" + ); + + // Wait until the token is released + await setTimeout(1_000); + + // Now the environment should have a concurrency of 0 + const envConcurrencyAfterReturn = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrencyAfterReturn).toBe(0); + + // and the queue should have a concurrency of 1 + const queueConcurrencyAfterReturn = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrencyAfterReturn).toBe(1); + } + ); + + containerTest( + "delays env concurrency release after checkpoint", + async ({ prisma, redisOptions }) => { + //create environment + const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + + const engine = new RunEngine({ + prisma, + worker: { + redis: redisOptions, + workers: 1, + tasksPerWorker: 10, + pollIntervalMs: 100, + }, + queue: { + redis: redisOptions, + }, + runLock: { + redis: redisOptions, + }, + machines: { + defaultMachine: "small-1x", + machines: { + "small-1x": { + name: "small-1x" as const, + cpu: 0.5, + memory: 0.5, + centsPerMs: 0.0001, + }, + }, + baseCostInCents: 0.0001, + }, + releaseConcurrency: { + maxTokensRatio: 0.1, // 10% of the concurrency limit = 1 token + maxRetries: 3, + consumersCount: 1, + pollInterval: 500, + batchSize: 1, + }, + tracer: trace.getTracer("test", "0.0.0"), + }); + const taskIdentifier = "test-task"; + + await setupBackgroundWorker(prisma, authenticatedEnvironment, taskIdentifier); + + const run = await engine.trigger( + { + number: 1, + friendlyId: "run_p1234", + environment: authenticatedEnvironment, + taskIdentifier, + payload: "{}", + payloadType: "application/json", + context: {}, + traceContext: {}, + traceId: "t12345", + spanId: "s12345", + masterQueue: "main", + queueName: `task/${taskIdentifier}`, + isTest: false, + tags: [], + }, + prisma + ); + + const dequeued = await engine.dequeueFromMasterQueue({ + consumerId: "test_12345", + masterQueue: run.masterQueue, + maxRunCount: 10, + }); + + const queueConcurrency = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrency).toBe(1); + + const envConcurrency = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrency).toBe(1); + + // create an attempt + const attemptResult = await engine.startRunAttempt({ + runId: dequeued[0].run.id, + snapshotId: dequeued[0].snapshot.id, + }); + + expect(attemptResult.snapshot.executionStatus).toBe("EXECUTING"); + + // create a manual waitpoint + const result = await engine.createManualWaitpoint({ + environmentId: authenticatedEnvironment.id, + projectId: authenticatedEnvironment.projectId, + }); + + await engine.releaseConcurrencyQueue.consumeToken( + { + orgId: authenticatedEnvironment.organizationId, + projectId: authenticatedEnvironment.projectId, + envId: authenticatedEnvironment.id, + }, + "test_12345" + ); + + // Block the run, not specifying any release concurrency option + const executingWithWaitpointSnapshot = await engine.blockRunWithWaitpoint({ + runId: run.id, + waitpoints: result.waitpoint.id, + projectId: authenticatedEnvironment.projectId, + organizationId: authenticatedEnvironment.organizationId, + }); + + expect(executingWithWaitpointSnapshot.executionStatus).toBe("EXECUTING_WITH_WAITPOINTS"); + + // Now confirm the queue has the same concurrency as before + const queueConcurrencyAfter = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrencyAfter).toBe(1); + + // Now confirm the environment is the same as before + const envConcurrencyAfter = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrencyAfter).toBe(1); + + const checkpointResult = await engine.createCheckpoint({ + runId: run.id, + snapshotId: executingWithWaitpointSnapshot.id, + checkpoint: { + type: "DOCKER", + reason: "TEST_CHECKPOINT", + location: "test-location", + imageRef: "test-image-ref", + }, + }); + + expect(checkpointResult.ok).toBe(true); + + const snapshot = checkpointResult.ok ? checkpointResult.snapshot : null; + assertNonNullable(snapshot); + + console.log("Snapshot", snapshot); + + expect(snapshot.executionStatus).toBe("SUSPENDED"); + + // Now we return the token so the concurrency can be released + await engine.releaseConcurrencyQueue.returnToken( + { + orgId: authenticatedEnvironment.organizationId, + projectId: authenticatedEnvironment.projectId, + envId: authenticatedEnvironment.id, + }, + "test_12345" + ); + + // Wait until the token is released + await setTimeout(1_000); + + // Now the environment should have a concurrency of 0 + const envConcurrencyAfterReturn = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrencyAfterReturn).toBe(0); + + // and the queue should have a concurrency of 1 + const queueConcurrencyAfterReturn = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrencyAfterReturn).toBe(1); + } + ); + + containerTest( + "maintains concurrency after waitpoint completion", + async ({ prisma, redisOptions }) => { + //create environment + const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + + const engine = new RunEngine({ + prisma, + worker: { + redis: redisOptions, + workers: 1, + tasksPerWorker: 10, + pollIntervalMs: 100, + }, + queue: { + redis: redisOptions, + }, + runLock: { + redis: redisOptions, + }, + machines: { + defaultMachine: "small-1x", + machines: { + "small-1x": { + name: "small-1x" as const, + cpu: 0.5, + memory: 0.5, + centsPerMs: 0.0001, + }, + }, + baseCostInCents: 0.0001, + }, + releaseConcurrency: { + maxTokensRatio: 0.1, // 10% of the concurrency limit = 1 token + maxRetries: 3, + consumersCount: 1, + pollInterval: 500, + batchSize: 1, + }, + tracer: trace.getTracer("test", "0.0.0"), + }); + const taskIdentifier = "test-task"; + + await setupBackgroundWorker(prisma, authenticatedEnvironment, taskIdentifier); + + const run = await engine.trigger( + { + number: 1, + friendlyId: "run_p1234", + environment: authenticatedEnvironment, + taskIdentifier, + payload: "{}", + payloadType: "application/json", + context: {}, + traceContext: {}, + traceId: "t12345", + spanId: "s12345", + masterQueue: "main", + queueName: `task/${taskIdentifier}`, + isTest: false, + tags: [], + }, + prisma + ); + + const dequeued = await engine.dequeueFromMasterQueue({ + consumerId: "test_12345", + masterQueue: run.masterQueue, + maxRunCount: 10, + }); + + const queueConcurrency = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrency).toBe(1); + + const envConcurrency = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrency).toBe(1); + + // create an attempt + const attemptResult = await engine.startRunAttempt({ + runId: dequeued[0].run.id, + snapshotId: dequeued[0].snapshot.id, + }); + + expect(attemptResult.snapshot.executionStatus).toBe("EXECUTING"); + + // create a manual waitpoint + const result = await engine.createManualWaitpoint({ + environmentId: authenticatedEnvironment.id, + projectId: authenticatedEnvironment.projectId, + }); + + await engine.releaseConcurrencyQueue.consumeToken( + { + orgId: authenticatedEnvironment.organizationId, + projectId: authenticatedEnvironment.projectId, + envId: authenticatedEnvironment.id, + }, + "test_12345" + ); + + // Block the run, not specifying any release concurrency option + const executingWithWaitpointSnapshot = await engine.blockRunWithWaitpoint({ + runId: run.id, + waitpoints: result.waitpoint.id, + projectId: authenticatedEnvironment.projectId, + organizationId: authenticatedEnvironment.organizationId, + }); + + expect(executingWithWaitpointSnapshot.executionStatus).toBe("EXECUTING_WITH_WAITPOINTS"); + + // Now confirm the queue has the same concurrency as before + const queueConcurrencyAfter = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrencyAfter).toBe(1); + + // Now confirm the environment is the same as before + const envConcurrencyAfter = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrencyAfter).toBe(1); + + // Complete the waitpoint + await engine.completeWaitpoint({ + id: result.waitpoint.id, + }); + + await setTimeout(1_000); + + // Verify the first run is now in EXECUTING state + const executionDataAfter = await engine.getRunExecutionData({ runId: run.id }); + expect(executionDataAfter?.snapshot.executionStatus).toBe("EXECUTING"); + + // Now we return the token so the concurrency can be released + await engine.releaseConcurrencyQueue.returnToken( + { + orgId: authenticatedEnvironment.organizationId, + projectId: authenticatedEnvironment.projectId, + envId: authenticatedEnvironment.id, + }, + "test_12345" + ); + + // give the release concurrency system time to run + await setTimeout(1_000); + + // Now the environment should have a concurrency of 1 + const envConcurrencyAfterReturn = await engine.runQueue.currentConcurrencyOfEnvironment( + authenticatedEnvironment + ); + + expect(envConcurrencyAfterReturn).toBe(1); + + // and the queue should have a concurrency of 1 + const queueConcurrencyAfterReturn = await engine.runQueue.currentConcurrencyOfQueue( + authenticatedEnvironment, + `task/${taskIdentifier}` + ); + + expect(queueConcurrencyAfterReturn).toBe(1); + } + ); +}); diff --git a/internal-packages/run-engine/src/engine/tests/releasingConcurrency.test.ts b/internal-packages/run-engine/src/engine/tests/releasingConcurrency.test.ts deleted file mode 100644 index 7d5c2e89e..000000000 --- a/internal-packages/run-engine/src/engine/tests/releasingConcurrency.test.ts +++ /dev/null @@ -1,10 +0,0 @@ -import { containerTest } from "@internal/testcontainers"; - -vi.setConfig({ testTimeout: 60_000 }); - -describe("RunEngine Releasing Concurrency", () => { - containerTest( - "blocking a run with a waitpoint with releasing concurrency", - async ({ prisma, redisOptions }) => {} - ); -}); diff --git a/internal-packages/run-engine/src/engine/tests/waitpoints.test.ts b/internal-packages/run-engine/src/engine/tests/waitpoints.test.ts index 3c06cd600..519212aed 100644 --- a/internal-packages/run-engine/src/engine/tests/waitpoints.test.ts +++ b/internal-packages/run-engine/src/engine/tests/waitpoints.test.ts @@ -7,7 +7,7 @@ import { import { trace } from "@internal/tracing"; import { expect } from "vitest"; import { RunEngine } from "../index.js"; -import { setTimeout } from "timers/promises"; +import { setTimeout } from "node:timers/promises"; import { EventBusEventArgs } from "../eventBus.js"; import { isWaitpointOutputTimeout } from "@trigger.dev/core/v3"; diff --git a/internal-packages/run-engine/src/engine/types.ts b/internal-packages/run-engine/src/engine/types.ts index 2d548f1ba..f823c80b1 100644 --- a/internal-packages/run-engine/src/engine/types.ts +++ b/internal-packages/run-engine/src/engine/types.ts @@ -105,6 +105,7 @@ export type TriggerParams = { machine?: MachinePresetName; workerId?: string; runnerId?: string; + releaseConcurrency?: boolean; }; export type EngineWorker = Worker; diff --git a/internal-packages/run-engine/src/index.ts b/internal-packages/run-engine/src/index.ts index 89bd08196..8d77c66a0 100644 --- a/internal-packages/run-engine/src/index.ts +++ b/internal-packages/run-engine/src/index.ts @@ -1,2 +1,3 @@ -export { RunEngine, RunDuplicateIdempotencyKeyError } from "./engine/index.js"; +export { RunEngine } from "./engine/index.js"; +export { RunDuplicateIdempotencyKeyError } from "./engine/errors.js"; export type { EventBusEventArgs } from "./engine/eventBus.js"; diff --git a/internal-packages/run-engine/src/run-queue/index.test.ts b/internal-packages/run-engine/src/run-queue/index.test.ts index 4085d87df..dbbb574bf 100644 --- a/internal-packages/run-engine/src/run-queue/index.test.ts +++ b/internal-packages/run-engine/src/run-queue/index.test.ts @@ -690,7 +690,10 @@ describe("RunQueue", () => { expect(await queue.currentConcurrencyOfEnvironment(authenticatedEnvProd)).toBe(1); //release the concurrency - await queue.releaseConcurrency(authenticatedEnvProd.organization.id, messages[0].messageId); + await queue.releaseAllConcurrency( + authenticatedEnvProd.organization.id, + messages[0].messageId + ); //concurrencies expect(await queue.currentConcurrencyOfQueue(authenticatedEnvProd, messageProd.queue)).toBe( @@ -708,7 +711,10 @@ describe("RunQueue", () => { expect(await queue.currentConcurrencyOfEnvironment(authenticatedEnvProd)).toBe(1); //release the concurrency (with the queue this time) - await queue.releaseConcurrency(authenticatedEnvProd.organization.id, messages[0].messageId); + await queue.releaseAllConcurrency( + authenticatedEnvProd.organization.id, + messages[0].messageId + ); //concurrencies expect(await queue.currentConcurrencyOfQueue(authenticatedEnvProd, messageProd.queue)).toBe( diff --git a/internal-packages/run-engine/src/run-queue/index.ts b/internal-packages/run-engine/src/run-queue/index.ts index cd4a571da..a5aacd957 100644 --- a/internal-packages/run-engine/src/run-queue/index.ts +++ b/internal-packages/run-engine/src/run-queue/index.ts @@ -561,14 +561,17 @@ export class RunQueue { ); } - public async releaseConcurrency(orgId: string, messageId: string) { + /** + * Release all concurrency for a message, including environment and queue concurrency + */ + public async releaseAllConcurrency(orgId: string, messageId: string) { return this.#trace( - "releaseConcurrency", + "releaseAllConcurrency", async (span) => { const message = await this.readMessage(orgId, messageId); if (!message) { - this.logger.log(`[${this.name}].acknowledgeMessage() message not found`, { + this.logger.log(`[${this.name}].releaseAllConcurrency() message not found`, { messageId, service: this.name, }); @@ -591,7 +594,44 @@ export class RunQueue { { kind: SpanKind.CONSUMER, attributes: { - [SEMATTRS_MESSAGING_OPERATION]: "releaseConcurrency", + [SEMATTRS_MESSAGING_OPERATION]: "releaseAllConcurrency", + [SEMATTRS_MESSAGE_ID]: messageId, + [SEMATTRS_MESSAGING_SYSTEM]: "runqueue", + }, + } + ); + } + + public async releaseEnvConcurrency(orgId: string, messageId: string) { + return this.#trace( + "releaseEnvConcurrency", + async (span) => { + const message = await this.readMessage(orgId, messageId); + + if (!message) { + this.logger.log(`[${this.name}].releaseEnvConcurrency() message not found`, { + messageId, + service: this.name, + }); + return; + } + + span.setAttributes({ + [SemanticAttributes.QUEUE]: message.queue, + [SemanticAttributes.ORG_ID]: message.orgId, + [SemanticAttributes.RUN_ID]: messageId, + [SemanticAttributes.CONCURRENCY_KEY]: message.concurrencyKey, + }); + + return this.redis.releaseEnvConcurrency( + this.keys.envCurrentConcurrencyKeyFromQueue(message.queue), + messageId + ); + }, + { + kind: SpanKind.CONSUMER, + attributes: { + [SEMATTRS_MESSAGING_OPERATION]: "releaseEnvConcurrency", [SEMATTRS_MESSAGE_ID]: messageId, [SEMATTRS_MESSAGING_SYSTEM]: "runqueue", }, @@ -1242,6 +1282,20 @@ redis.call('SREM', envCurrentConcurrencyKey, messageId) `, }); + this.redis.defineCommand("releaseEnvConcurrency", { + numberOfKeys: 1, + lua: ` +-- Keys: +local envCurrentConcurrencyKey = KEYS[1] + +-- Args: +local messageId = ARGV[1] + +-- Update the concurrency keys +redis.call('SREM', envCurrentConcurrencyKey, messageId) +`, + }); + this.redis.defineCommand("reacquireConcurrency", { numberOfKeys: 4, lua: ` @@ -1274,12 +1328,14 @@ if envCurrentConcurrency >= totalEnvConcurrencyLimit then end -- Check current queue concurrency against the limit -local queueCurrentConcurrency = tonumber(redis.call('SCARD', queueCurrentConcurrencyKey) or '0') -local queueConcurrencyLimit = math.min(tonumber(redis.call('GET', queueConcurrencyLimitKey) or '1000000'), envConcurrencyLimit) -local totalQueueConcurrencyLimit = queueConcurrencyLimit +if not isInQueueConcurrency then + local queueCurrentConcurrency = tonumber(redis.call('SCARD', queueCurrentConcurrencyKey) or '0') + local queueConcurrencyLimit = math.min(tonumber(redis.call('GET', queueConcurrencyLimitKey) or '1000000'), envConcurrencyLimit) + local totalQueueConcurrencyLimit = queueConcurrencyLimit -if queueCurrentConcurrency >= totalQueueConcurrencyLimit then - return false + if queueCurrentConcurrency >= totalQueueConcurrencyLimit then + return false + end end -- Update the concurrency keys @@ -1390,6 +1446,12 @@ declare module "@internal/redis" { callback?: Callback ): Result; + releaseEnvConcurrency( + envCurrentConcurrencyKey: string, + messageId: string, + callback?: Callback + ): Result; + reacquireConcurrency( queueCurrentConcurrencyKey: string, envCurrentConcurrencyKey: string, diff --git a/internal-packages/run-engine/src/run-queue/tests/reacquireConcurrency.test.ts b/internal-packages/run-engine/src/run-queue/tests/reacquireConcurrency.test.ts index 6261e4878..a9c0386ca 100644 --- a/internal-packages/run-engine/src/run-queue/tests/reacquireConcurrency.test.ts +++ b/internal-packages/run-engine/src/run-queue/tests/reacquireConcurrency.test.ts @@ -87,7 +87,7 @@ describe("RunQueue.reacquireConcurrency", () => { expect(await queue.currentConcurrencyOfEnvironment(authenticatedEnvProd)).toBe(1); // First, release the concurrency - await queue.releaseConcurrency(authenticatedEnvProd.organization.id, messageProd.runId); + await queue.releaseAllConcurrency(authenticatedEnvProd.organization.id, messageProd.runId); //reacquire the concurrency const result = await queue.reacquireConcurrency( diff --git a/internal-packages/run-engine/src/run-queue/tests/releaseConcurrency.test.ts b/internal-packages/run-engine/src/run-queue/tests/releaseConcurrency.test.ts index 47be728c6..63873a54b 100644 --- a/internal-packages/run-engine/src/run-queue/tests/releaseConcurrency.test.ts +++ b/internal-packages/run-engine/src/run-queue/tests/releaseConcurrency.test.ts @@ -81,7 +81,7 @@ describe("RunQueue.releaseConcurrency", () => { expect(await queue.currentConcurrencyOfEnvironment(authenticatedEnvProd)).toBe(1); //release the concurrency - await queue.releaseConcurrency(authenticatedEnvProd.organization.id, messageProd.runId); + await queue.releaseAllConcurrency(authenticatedEnvProd.organization.id, messageProd.runId); //concurrencies expect(await queue.currentConcurrencyOfQueue(authenticatedEnvProd, messageProd.queue)).toBe( @@ -137,7 +137,7 @@ describe("RunQueue.releaseConcurrency", () => { expect(await queue.currentConcurrencyOfEnvironment(authenticatedEnvProd)).toBe(1); //release the concurrency - await queue.releaseConcurrency(authenticatedEnvProd.organization.id, "r1235"); + await queue.releaseAllConcurrency(authenticatedEnvProd.organization.id, "r1235"); //concurrencies expect(await queue.currentConcurrencyOfQueue(authenticatedEnvProd, messageProd.queue)).toBe( diff --git a/internal-packages/testcontainers/src/setup.ts b/internal-packages/testcontainers/src/setup.ts index a51e24ead..b77663f06 100644 --- a/internal-packages/testcontainers/src/setup.ts +++ b/internal-packages/testcontainers/src/setup.ts @@ -69,7 +69,11 @@ export async function setupBackgroundWorker( environment: AuthenticatedEnvironment, taskIdentifier: string | string[], machineConfig?: MachineConfig, - retryOptions?: RetryOptions + retryOptions?: RetryOptions, + queueOptions?: { + releaseConcurrencyOnWaitpoint?: boolean; + concurrencyLimit?: number | null; + } ) { const worker = await prisma.backgroundWorker.create({ data: { @@ -115,10 +119,17 @@ export async function setupBackgroundWorker( data: { friendlyId: generateFriendlyId("queue"), name: queueName, - concurrencyLimit: 10, + concurrencyLimit: + typeof queueOptions?.concurrencyLimit === "undefined" + ? 10 + : queueOptions.concurrencyLimit, runtimeEnvironmentId: worker.runtimeEnvironmentId, projectId: worker.projectId, type: "VIRTUAL", + releaseConcurrencyOnWaitpoint: + typeof queueOptions?.releaseConcurrencyOnWaitpoint === "boolean" + ? queueOptions.releaseConcurrencyOnWaitpoint + : undefined, }, }); } diff --git a/packages/core/src/v3/schemas/api.ts b/packages/core/src/v3/schemas/api.ts index ff2028f7c..48be14e60 100644 --- a/packages/core/src/v3/schemas/api.ts +++ b/packages/core/src/v3/schemas/api.ts @@ -119,6 +119,7 @@ export const TriggerTaskRequestBody = z.object({ test: z.boolean().optional(), ttl: z.string().or(z.number().nonnegative().int()).optional(), priority: z.number().optional(), + releaseConcurrency: z.boolean().optional(), }) .optional(), }); @@ -956,6 +957,9 @@ export const WaitForDurationRequestBody = z.object({ * This means after that time if you pass the same idempotency key again, you will get a new waitpoint. */ idempotencyKeyTTL: z.string().optional(), + + releaseConcurrency: z.boolean().optional(), + date: z.coerce.date(), }); export type WaitForDurationRequestBody = z.infer; diff --git a/packages/core/src/v3/types/tasks.ts b/packages/core/src/v3/types/tasks.ts index 921712ec3..7a4f27915 100644 --- a/packages/core/src/v3/types/tasks.ts +++ b/packages/core/src/v3/types/tasks.ts @@ -824,8 +824,16 @@ export type TriggerOptions = { version?: string; }; -export type TriggerAndWaitOptions = Omit; - +export type TriggerAndWaitOptions = Omit & { + /** + * If set to true, this will cause the waitpoint to release the current run from the queue's concurrency. + * + * This is useful if you want to allow other runs to execute while the child task is executing + * + * @default false + */ + releaseConcurrency?: boolean; +}; export type BatchTriggerOptions = { /** * If no idempotencyKey is set on an individual item in the batch, it will use this key on each item + the array index. diff --git a/packages/trigger-sdk/src/v3/shared.ts b/packages/trigger-sdk/src/v3/shared.ts index d2e083aba..6d4b4606d 100644 --- a/packages/trigger-sdk/src/v3/shared.ts +++ b/packages/trigger-sdk/src/v3/shared.ts @@ -1345,6 +1345,7 @@ async function triggerAndWait_internal { logger.log("Hello, world from the parent", { payload }); - await childTask.triggerAndWait({ message: "Hello, world!" }); + await childTask.triggerAndWait( + { message: "Hello, world!" }, + { + releaseConcurrency: true, + } + ); }, }); diff --git a/references/hello-world/src/trigger/waits.ts b/references/hello-world/src/trigger/waits.ts index 0749cde17..f3a84a7a5 100644 --- a/references/hello-world/src/trigger/waits.ts +++ b/references/hello-world/src/trigger/waits.ts @@ -73,7 +73,12 @@ export const waitForDuration = task({ }) => { const idempotency = idempotencyKey ? await idempotencyKeys.create(idempotencyKey) : undefined; - await wait.for({ seconds: duration, idempotencyKey: idempotency, idempotencyKeyTTL }); + await wait.for({ + seconds: duration, + idempotencyKey: idempotency, + idempotencyKeyTTL, + releaseConcurrency: true, + }); await wait.until({ date: new Date(Date.now() + duration * 1000) }); await retry.fetch("https://example.com/404", { method: "GET" });