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

.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

@@ -1,4 +1,5 @@
 import { type AuthenticatedEnvironment } from "~/services/apiAuth.server";
+import { logger } from "~/services/logger.server";
 import { engine } from "./runEngine.server";
 
 /** Updates the RunQueue env concurrency limits */
@@ -30,3 +31,47 @@ 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 });
+}
+
+/**
+ * Best-effort sweep for the pause paths.
+ *
+ * By the time this runs the pause is already durable and in force, so a failure here must
+ * not roll it back or surface as a failed pause — the worst case is that some already-admitted
+ * runs still execute, which is the behaviour that existed before the sweep.
+ */
+export async function sweepUnclaimedRuns(
+  environment: AuthenticatedEnvironment,
+  queue?: string
+): Promise<void> {
+  try {
+    const result = await returnUnclaimedMessagesToQueue({ environment, queue });
+
+    logger.debug("sweepUnclaimedRuns", {
+      environmentId: environment.id,
+      queue,
+      ...result,
+    });
+  } catch (error) {
+    logger.error("sweepUnclaimedRuns failed", {
+      environmentId: environment.id,
+      queue,
+      error,
+    });
+  }
+}

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 },
@@ -182,6 +200,15 @@ async function pauseEnvironmentForBillingLimit(
     // The env's paused state changed (or was rolled back); drop any cached copy either way.
     controlPlaneResolver.invalidateEnvironment(environment.id);
   }
+
+  try {
+    await returnUnclaimed(updated);
+  } catch (error) {
+    logger.error("Billing limit converge failed to return unclaimed runs", {
+      environmentId: environment.id,
+      error,
+    });
+  }
 }
 
 async function resumeEnvironmentFromBillingLimit(

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 { sweepUnclaimedRuns, updateEnvConcurrencyLimits } from "../runQueue.server";
 import { WithRunEngine } from "./baseService.server";
 import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server";
@@ -136,6 +136,10 @@ export class PauseEnvironmentService extends WithRunEngine {
       // The env's `paused` state changed in the control-plane; drop any cached copy.
       controlPlaneResolver.invalidateEnvironment(environment.id);
 
+      if (action === "paused") {
+        await sweepUnclaimedRuns(environment);
+      }
+
       return {
         success: true,
         state: action,

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,
+  sweepUnclaimedRuns,
+  updateQueueConcurrencyLimits,
+} from "../runQueue.server";
 import { engine } from "../runEngine.server";
 
 export type PauseStatus = "paused" | "resumed";
@@ -73,6 +77,10 @@ export class PauseQueueService extends BaseService {
         environmentId: environment.id,
       });
 
+      if (action === "paused") {
+        await sweepUnclaimedRuns(environment, queue.name);
+      }
+
       const results = await Promise.all([
         engine.lengthOfQueues(environment, [queue.name]),
         engine.currentConcurrencyOfQueues(environment, [queue.name]),

apps/webapp/test/billingLimitConvergeEnvironments.test.ts

@@ -49,6 +49,86 @@ 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("keeps the 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"),
+    });
+    const second = await createRuntimeEnvironment(prisma, {
+      projectId: project.id,
+      organizationId: organization.id,
+      type: "PRODUCTION",
+      slug: uniqueId("prod"),
+    });
+
+    const result = await convergeBillingLimitEnvironmentsForOrg(organization.id, "grace", {
+      prismaClient: prisma,
+      updateConcurrency: async () => undefined,
+      returnUnclaimed: async () => {
+        throw new Error("run queue unavailable");
+      },
+    });
+
+    expect(result).toEqual({ paused: 2, unpaused: 0 });
+
+    for (const env of [environment, second]) {
+      const envAfter = await prisma.runtimeEnvironment.findUniqueOrThrow({
+        where: { id: env.id },
+      });
+      expect(envAfter.paused).toBe(true);
+      expect(envAfter.pauseSource).toBe(EnvironmentPauseSource.BILLING_LIMIT);
+    }
+  });
+
   postgresTest("rolls back pause when concurrency update fails", async ({ prisma }) => {
     const { organization, project } = await createTestOrgProjectWithMember(prisma);
     const environment = await createRuntimeEnvironment(prisma, {

apps/webapp/test/pauseSweepWiring.test.ts

@@ -0,0 +1,94 @@
+import { type PrismaClient } from "@trigger.dev/database";
+import { describe, expect, vi } from "vitest";
+import { postgresTest } from "@internal/testcontainers";
+import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
+import {
+  createRuntimeEnvironment,
+  createTestOrgProjectWithMember,
+  uniqueId,
+} from "./fixtures/environmentVariablesFixtures";
+
+vi.setConfig({ testTimeout: 60_000 });
+
+const calls: string[] = [];
+
+vi.mock("~/v3/runQueue.server", () => ({
+  updateEnvConcurrencyLimits: vi.fn(async (_env: unknown, limit?: number) => {
+    calls.push(`updateEnvConcurrencyLimits:${limit}`);
+  }),
+  updateQueueConcurrencyLimits: vi.fn(async (_env: unknown, name: string, limit: number) => {
+    calls.push(`updateQueueConcurrencyLimits:${name}:${limit}`);
+  }),
+  removeQueueConcurrencyLimits: vi.fn(async () => {
+    calls.push("removeQueueConcurrencyLimits");
+  }),
+  returnUnclaimedMessagesToQueue: vi.fn(async () => ({
+    returned: 0,
+    skippedLastPass: 0,
+    errors: 0,
+    passes: 1,
+  })),
+  sweepUnclaimedRuns: vi.fn(async (_env: unknown, queue?: string) => {
+    calls.push(`sweepUnclaimedRuns:${queue ?? "*"}`);
+  }),
+}));
+
+async function loadServices() {
+  const [{ PauseEnvironmentService }, { authIncludeBase, toAuthenticated }] = await Promise.all([
+    import("~/v3/services/pauseEnvironment.server"),
+    import("~/models/runtimeEnvironment.server"),
+  ]);
+  return { PauseEnvironmentService, authIncludeBase, toAuthenticated };
+}
+
+type Loaded = Awaited<ReturnType<typeof loadServices>>;
+
+async function seedEnv(
+  loaded: Loaded,
+  prisma: PrismaClient
+): Promise<{ environment: AuthenticatedEnvironment; environmentId: string }> {
+  const { organization, project } = await createTestOrgProjectWithMember(prisma);
+  const created = await createRuntimeEnvironment(prisma, {
+    projectId: project.id,
+    organizationId: organization.id,
+    type: "PRODUCTION",
+    slug: uniqueId("prod"),
+  });
+
+  const row = await prisma.runtimeEnvironment.findFirstOrThrow({
+    where: { id: created.id },
+    include: loaded.authIncludeBase,
+  });
+
+  return { environment: loaded.toAuthenticated(row), environmentId: created.id };
+}
+
+describe("pause sweep wiring", () => {
+  postgresTest("pausing an environment sweeps after the limit is zeroed", async ({ prisma }) => {
+    calls.length = 0;
+    const loaded = await loadServices();
+    const { environment } = await seedEnv(loaded, prisma);
+
+    const result = await new loaded.PauseEnvironmentService(prisma).call(environment, "paused");
+
+    expect(result).toEqual({ success: true, state: "paused" });
+    expect(calls).toEqual(["updateEnvConcurrencyLimits:0", "sweepUnclaimedRuns:*"]);
+  });
+
+  postgresTest("resuming an environment does not sweep", async ({ prisma }) => {
+    const loaded = await loadServices();
+    const { environment, environmentId } = await seedEnv(loaded, prisma);
+
+    await prisma.runtimeEnvironment.update({
+      where: { id: environmentId },
+      data: { paused: true },
+    });
+
+    calls.length = 0;
+
+    const result = await new loaded.PauseEnvironmentService(prisma).call(environment, "resumed");
+
+    expect(result).toEqual({ success: true, state: "resumed" });
+    expect(calls).toEqual(["updateEnvConcurrencyLimits:undefined"]);
+  });
+});