fix(run-engine,webapp): return unclaimed runs to the queue when pausing

Pausing a queue or environment set its concurrency limit to 0, which stops the
queue admitting new runs but does nothing about runs it had already admitted.
Those runs sit in a per-region worker queue that supervisors drain with a plain
BLPOP, with no pause or concurrency check on the pop path, so they went on to
execute during the pause. An idle environment with headroom can have its whole
concurrency limit's worth of runs in that state, since the trigger-time fast
path pushes straight to the worker queue.

Adds RunQueue.returnUnclaimedMessagesToQueue, which moves runs the queue has
admitted but no worker has claimed yet back into the pending queue. Candidates
come from the difference between currentConcurrency and currentDequeued, which
identifies worker queue residents without scanning the region-wide list. Each
move runs in a Lua script whose LREM doubles as the claim: if it removes
nothing then a worker already owns that run and every other key is left alone,
so a concurrent BLPOP can never produce a duplicate dispatch.

Ordering is preserved by re-adding at the message's original timestamp rather
than a fresh score the way nack does, so a returned run keeps its place in
line. Callers must set the concurrency limit to 0 first: returning a run costs
it its position in the worker queue FIFO and it can only re-enter at the back,
so a queue that can still admit work could let a newer run overtake one on its
way back. All three pause callers already do this ordering.

Covers the manual queue pause, the manual environment pause, and the
billing-limit environment pause.

Author: matt-aitken

.server-changes/pause-returns-waiting-runs.md

@@ -0,0 +1,6 @@
+---
+area: webapp
+type: fix
+---
+
+Pausing a queue or an environment now also holds back runs that were already waiting to start. Previously those runs would still go ahead, so a pause could take effect a little later than expected.

apps/webapp/app/v3/runQueue.server.ts

@@ -30,3 +30,19 @@ export async function removeQueueConcurrencyLimits(
 ) {
   await engine.runQueue.removeQueueConcurrencyLimits(environment, queueName);
 }
+
+/**
+ * Returns runs waiting to be picked up by a worker back into their queue.
+ *
+ * Only safe once the concurrency limit has been set to 0, otherwise a newer run can be
+ * admitted and overtake one on its way back.
+ */
+export async function returnUnclaimedMessagesToQueue({
+  environment,
+  queue,
+}: {
+  environment: AuthenticatedEnvironment;
+  queue?: string;
+}) {
+  return engine.returnUnclaimedMessagesToQueue({ environment, queue });
+}

apps/webapp/app/v3/services/billingLimit/billingLimitConvergeEnvironments.server.ts

@@ -29,13 +29,16 @@ type UpdateEnvConcurrency = (
   maximumConcurrencyLimit?: number
 ) => Promise<void>;
 
+type ReturnUnclaimedMessages = (environment: EnvironmentWithRelations) => Promise<void>;
+
 export async function convergeBillingLimitEnvironmentsForOrg(
   organizationId: string,
   targetState: BillingLimitConvergeTargetState,
   options?: {
     batchSize?: number;
     prismaClient?: PrismaClient;
     updateConcurrency?: UpdateEnvConcurrency;
+    returnUnclaimed?: ReturnUnclaimedMessages;
   }
 ): Promise<ConvergeOrgResult> {
   const db = options?.prismaClient ?? prisma;
@@ -51,18 +54,32 @@ export async function convergeBillingLimitEnvironmentsForOrg(
       return updateEnvConcurrencyLimits(environment, maximumConcurrencyLimit);
     });
 
+  const returnUnclaimed =
+    options?.returnUnclaimed ??
+    (async (environment) => {
+      const { returnUnclaimedMessagesToQueue } = await import("~/v3/runQueue.server");
+      await returnUnclaimedMessagesToQueue({ environment });
+    });
+
   if (targetState === "ok") {
     return unpauseBillingLimitEnvironments(organizationId, db, batchSize, updateConcurrency);
   }
 
-  return pauseBillingLimitEnvironments(organizationId, db, batchSize, updateConcurrency);
+  return pauseBillingLimitEnvironments(
+    organizationId,
+    db,
+    batchSize,
+    updateConcurrency,
+    returnUnclaimed
+  );
 }
 
 async function pauseBillingLimitEnvironments(
   organizationId: string,
   db: PrismaClient,
   batchSize: number,
-  updateConcurrency: UpdateEnvConcurrency
+  updateConcurrency: UpdateEnvConcurrency,
+  returnUnclaimed: ReturnUnclaimedMessages
 ): Promise<ConvergeOrgResult> {
   let paused = 0;
   let cursor: string | undefined;
@@ -88,7 +105,7 @@ async function pauseBillingLimitEnvironments(
     }
 
     for (const environment of environments) {
-      await pauseEnvironmentForBillingLimit(environment, db, updateConcurrency);
+      await pauseEnvironmentForBillingLimit(environment, db, updateConcurrency, returnUnclaimed);
       paused++;
     }
 
@@ -156,7 +173,8 @@ async function unpauseBillingLimitEnvironments(
 async function pauseEnvironmentForBillingLimit(
   environment: EnvironmentWithRelations,
   db: PrismaClient,
-  updateConcurrency: UpdateEnvConcurrency
+  updateConcurrency: UpdateEnvConcurrency,
+  returnUnclaimed: ReturnUnclaimedMessages
 ) {
   const updated = await db.runtimeEnvironment.update({
     where: { id: environment.id },
@@ -172,6 +190,7 @@ async function pauseEnvironmentForBillingLimit(
 
   try {
     await updateConcurrency(updated, 0);
+    await returnUnclaimed(updated);
   } catch (error) {
     await db.runtimeEnvironment.update({
       where: { id: environment.id },

apps/webapp/app/v3/services/pauseEnvironment.server.ts

@@ -2,7 +2,7 @@ import { EnvironmentPauseSource, type PrismaClientOrTransaction } from "@trigger
 import { prisma } from "~/db.server";
 import { logger } from "~/services/logger.server";
 import { getManualPauseEnvironmentResult } from "~/v3/services/billingLimit/manualPauseEnvironmentGuard.server";
-import { updateEnvConcurrencyLimits } from "../runQueue.server";
+import { returnUnclaimedMessagesToQueue, updateEnvConcurrencyLimits } from "../runQueue.server";
 import { WithRunEngine } from "./baseService.server";
 import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server";
@@ -114,6 +114,7 @@ export class PauseEnvironmentService extends WithRunEngine {
             environmentId: environment.id,
           });
           await updateEnvConcurrencyLimits(environment, 0);
+          await returnUnclaimedMessagesToQueue({ environment });
         } else {
           logger.debug("PauseEnvironmentService: resuming environment", {
             environmentId: environment.id,

apps/webapp/app/v3/services/pauseQueue.server.ts

@@ -4,7 +4,11 @@ import { type AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
 import { BaseService } from "./baseService.server";
 import { determineEngineVersion } from "../engineVersion.server";
-import { removeQueueConcurrencyLimits, updateQueueConcurrencyLimits } from "../runQueue.server";
+import {
+  removeQueueConcurrencyLimits,
+  returnUnclaimedMessagesToQueue,
+  updateQueueConcurrencyLimits,
+} from "../runQueue.server";
 import { engine } from "../runEngine.server";
 
 export type PauseStatus = "paused" | "resumed";
@@ -59,6 +63,7 @@ export class PauseQueueService extends BaseService {
 
       if (action === "paused") {
         await updateQueueConcurrencyLimits(environment, queue.name, 0);
+        await returnUnclaimedMessagesToQueue({ environment, queue: queue.name });
       } else {
         if (queue.concurrencyLimit) {
           await updateQueueConcurrencyLimits(environment, queue.name, queue.concurrencyLimit);

apps/webapp/test/billingLimitConvergeEnvironments.test.ts

@@ -49,6 +49,78 @@ describe("convergeBillingLimitEnvironmentsForOrg", () => {
     expect(envAfter.pauseSource).toBeNull();
   });
 
+  postgresTest("returns unclaimed runs after pausing for a billing limit", async ({ prisma }) => {
+    const { organization, project } = await createTestOrgProjectWithMember(prisma);
+    const environment = await createRuntimeEnvironment(prisma, {
+      projectId: project.id,
+      organizationId: organization.id,
+      type: "PRODUCTION",
+      slug: uniqueId("prod"),
+    });
+
+    const calls: Array<{ concurrency?: number; returnedFor?: string }> = [];
+
+    const result = await convergeBillingLimitEnvironmentsForOrg(organization.id, "grace", {
+      prismaClient: prisma,
+      updateConcurrency: async (_env, maximumConcurrencyLimit) => {
+        calls.push({ concurrency: maximumConcurrencyLimit });
+      },
+      returnUnclaimed: async (env) => {
+        calls.push({ returnedFor: env.id });
+      },
+    });
+
+    expect(result).toEqual({ paused: 1, unpaused: 0 });
+
+    expect(calls).toEqual([{ concurrency: 0 }, { returnedFor: environment.id }]);
+
+    const envAfter = await prisma.runtimeEnvironment.findUniqueOrThrow({
+      where: { id: environment.id },
+    });
+    expect(envAfter.paused).toBe(true);
+    expect(envAfter.pauseSource).toBe(EnvironmentPauseSource.BILLING_LIMIT);
+  });
+
+  postgresTest("does not return unclaimed runs when unpausing", async ({ prisma }) => {
+    const { organization } = await createBillingPausedProductionEnv(prisma);
+
+    const returnUnclaimed = vi.fn(async () => undefined);
+
+    await convergeBillingLimitEnvironmentsForOrg(organization.id, "ok", {
+      prismaClient: prisma,
+      updateConcurrency: async () => undefined,
+      returnUnclaimed,
+    });
+
+    expect(returnUnclaimed).not.toHaveBeenCalled();
+  });
+
+  postgresTest("rolls back pause when returning unclaimed runs fails", async ({ prisma }) => {
+    const { organization, project } = await createTestOrgProjectWithMember(prisma);
+    const environment = await createRuntimeEnvironment(prisma, {
+      projectId: project.id,
+      organizationId: organization.id,
+      type: "PRODUCTION",
+      slug: uniqueId("prod"),
+    });
+
+    await expect(
+      convergeBillingLimitEnvironmentsForOrg(organization.id, "grace", {
+        prismaClient: prisma,
+        updateConcurrency: async () => undefined,
+        returnUnclaimed: async () => {
+          throw new Error("run queue unavailable");
+        },
+      })
+    ).rejects.toThrow("run queue unavailable");
+
+    const envAfter = await prisma.runtimeEnvironment.findUniqueOrThrow({
+      where: { id: environment.id },
+    });
+    expect(envAfter.paused).toBe(false);
+    expect(envAfter.pauseSource).toBeNull();
+  });
+
   postgresTest("rolls back pause when concurrency update fails", async ({ prisma }) => {
     const { organization, project } = await createTestOrgProjectWithMember(prisma);
     const environment = await createRuntimeEnvironment(prisma, {

internal-packages/run-engine/src/engine/index.ts

@@ -1634,6 +1634,24 @@ export class RunEngine {
     return this.runQueue.currentConcurrencyOfQueues(environment, queues);
   }
 
+  /**
+   * Returns runs that have been admitted for execution but that no worker has claimed yet
+   * back into their queue, at the position they held before.
+   *
+   * Pausing only stops runs being admitted; anything already handed to a worker queue would
+   * otherwise still execute. Call this after the queue has been made ineligible for
+   * dequeuing, never before — see `RunQueue.returnUnclaimedMessagesToQueue`.
+   */
+  async returnUnclaimedMessagesToQueue({
+    environment,
+    queue,
+  }: {
+    environment: MinimalAuthenticatedEnvironment;
+    queue?: string;
+  }): Promise<{ returned: number; skipped: number }> {
+    return this.runQueue.returnUnclaimedMessagesToQueue({ env: environment, queue });
+  }
+
   async removeEnvironmentQueuesFromMasterQueue({
     runtimeEnvironmentId,
     organizationId,