fix(redis-worker): bound DRAINING blast radius with worker-pool drain

d-cs · claude · d-cs · commit 3bceb549fd70 · 2026-06-01T20:37:58.000+01:00
Addresses CodeRabbit review: the prefetched-pop tick design moved every
popped entry into DRAINING before any of them got a pLimit slot, so a
process crash mid-tick stranded up to maxOrgsPerTick × drainBatchSize
entries for stale-sweep to recover (~25k at defaults).

Replaces it with a worker-pool: spawn min(concurrency, totalBudget)
workers; each worker round-robin-picks an env with budget remaining,
pops one entry, processes it, releases its slot. At any moment, the
count of popped-but-not-acked entries is bounded by `concurrency` —
identical safety to the pre-batch one-pop-per-env path — while a
single-env burst still uses the full concurrency budget (all workers
can pull from the same env).

Adds a regression test pinning the bound: never has more than
`concurrency` entries popped-but-not-acked at any moment.

Two existing batch tests now use concurrency=1 to isolate the
break-on-empty/error semantic from the worker-pool's parallel-pick
race (the race semantic itself is covered by the new safety test).

Co-Authored-By: Claude Opus 4.7 &lt;noreply@anthropic.com&gt;
diff --git a/packages/redis-worker/src/mollifier/drainer.test.ts b/packages/redis-worker/src/mollifier/drainer.test.ts
@@ -276,7 +276,13 @@ describe("MollifierDrainer.drainBatchSize", () => {
       handler: async (input) => {
         handled.push(input.runId);
       },
-      concurrency: 5,
+      // Concurrency=1 so the worker pool runs sequentially and pop calls
+      // can't race past the `skip.add(envId)` that fires after a pop
+      // failure. The semantic this test pins (one env's pop blowup
+      // aborts its batch and counts as one failure) is the deterministic
+      // case; multi-worker race semantics are exercised by the safety
+      // property test below.
+      concurrency: 1,
       maxAttempts: 3,
       isRetryable: () => false,
       drainBatchSize: 5,
@@ -484,6 +490,69 @@ describe("MollifierDrainer.drainBatchSize", () => {
     expect(r.failed).toBe(2);
   });
 
+  it("never has more than `concurrency` entries popped-but-not-acked at any moment", async () => {
+    // Regression guard for the DRAINING blast radius. Each pop+process
+    // happens inside a single pLimit slot, so at any instant the number
+    // of entries that have been popped (and therefore marked DRAINING in
+    // a real buffer) but not yet acked is bounded by `concurrency`. This
+    // matters because the stale sweep only catches DRAINING entries
+    // visibly after a threshold — a process crash with thousands of
+    // mid-flight entries would mean a long detection/recovery window.
+    const envCount = 10;
+    const perEnv = 20;
+    const queues = new Map<string, string[]>();
+    for (let i = 0; i < envCount; i++) {
+      queues.set(
+        `env_${i}`,
+        Array.from({ length: perEnv }, (_, j) => `env_${i}_run_${j}`),
+      );
+    }
+    let inflightPoppedNotAcked = 0;
+    let peak = 0;
+    const concurrency = 4;
+    const buffer = makeStubBuffer({
+      ...eachEnvAsOwnOrg([...queues.keys()]),
+      pop: async (envId: string) => {
+        const q = queues.get(envId);
+        if (!q || q.length === 0) return null;
+        const runId = q.shift()!;
+        inflightPoppedNotAcked += 1;
+        if (inflightPoppedNotAcked > peak) peak = inflightPoppedNotAcked;
+        return {
+          runId,
+          envId,
+          orgId: envId,
+          payload: "{}",
+          attempts: 0,
+          createdAt: new Date(),
+        } as any;
+      },
+      ack: async () => {
+        inflightPoppedNotAcked -= 1;
+      },
+    });
+
+    const drainer = new MollifierDrainer({
+      buffer,
+      handler: async () => {
+        // Force handler overlap if scheduling allowed it — without a
+        // tight per-slot bound the peak would visibly exceed `concurrency`.
+        await new Promise((r) => setTimeout(r, 15));
+      },
+      concurrency,
+      maxAttempts: 3,
+      isRetryable: () => false,
+      drainBatchSize: perEnv,
+      logger: new Logger("test-drainer", "log"),
+    });
+
+    const r = await drainer.runOnce();
+    expect(r.drained).toBe(envCount * perEnv);
+    expect(peak).toBeGreaterThan(1); // concurrency is real, not serialised
+    expect(peak).toBeLessThanOrEqual(concurrency); // and bounded by it
+    expect(inflightPoppedNotAcked).toBe(0); // everything settled
+  });
+
   it("stops popping early when the env's queue empties before reaching drainBatchSize", async () => {
     const queue = ["only_1", "only_2"];
     const handled: string[] = [];
@@ -511,7 +580,12 @@ describe("MollifierDrainer.drainBatchSize", () => {
       handler: async (input) => {
         handled.push(input.runId);
       },
-      concurrency: 5,
+      // Concurrency=1 isolates the "stop on empty" semantic from the
+      // worker pool's parallel-pick race: with multiple workers, several
+      // can pick env_a simultaneously and pop in parallel before any of
+      // them can `skip.add(envId)`, so the empty-pop count would be
+      // >1 nondeterministically.
+      concurrency: 1,
       maxAttempts: 3,
       isRetryable: () => false,
       drainBatchSize: 10,
diff --git a/packages/redis-worker/src/mollifier/drainer.ts b/packages/redis-worker/src/mollifier/drainer.ts
@@ -1,5 +1,4 @@
 import { Logger } from "@trigger.dev/core/logger";
-import pLimit from "p-limit";
 import { MollifierBuffer } from "./buffer.js";
 import { BufferEntry, deserialiseSnapshot } from "./schemas.js";
 
@@ -84,8 +83,8 @@ export class MollifierDrainer<TPayload = unknown> {
   private readonly pollIntervalMs: number;
   private readonly maxOrgsPerTick: number;
   private readonly drainBatchSize: number;
+  private readonly concurrency: number;
   private readonly logger: Logger;
-  private readonly limit: ReturnType<typeof pLimit>;
   // Rotation state. `orgCursor` advances through the active-orgs list.
   // Each org has its own internal cursor in `perOrgEnvCursors` for
   // cycling through that org's envs. Both reset on `start()`.
@@ -104,8 +103,8 @@ export class MollifierDrainer<TPayload = unknown> {
     this.pollIntervalMs = options.pollIntervalMs ?? 100;
     this.maxOrgsPerTick = options.maxOrgsPerTick ?? 500;
     this.drainBatchSize = Math.max(1, options.drainBatchSize ?? 1);
+    this.concurrency = Math.max(1, options.concurrency);
     this.logger = options.logger ?? new Logger("MollifierDrainer", "debug");
-    this.limit = pLimit(options.concurrency);
   }
 
   async runOnce(): Promise<DrainResult> {
@@ -133,64 +132,98 @@ export class MollifierDrainer<TPayload = unknown> {
       targets.push(envId);
     }
 
-    // Pop a batch from each target env in parallel. Within an env we pop
-    // sequentially (each Lua `pop` is atomic; back-to-back pops on the
-    // same env can't be concurrent without a `popBatch` Lua, and Redis
-    // RTT × drainBatchSize is cheap compared to the engine.trigger work
-    // that follows). A pop failure mid-batch aborts only that env's
-    // batch and counts as one failure — same semantics as the previous
-    // one-pop-per-env path, generalised.
-    const envBatches = await Promise.all(
-      targets.map(async (envId) => {
-        const entries: BufferEntry[] = [];
-        let popFailed = false;
-        for (let i = 0; i < this.drainBatchSize; i++) {
-          let entry: BufferEntry | null;
-          try {
-            entry = await this.buffer.pop(envId);
-          } catch (err) {
-            this.logger.error("MollifierDrainer.pop failed", { envId, err });
-            popFailed = true;
-            break;
-          }
-          if (!entry) break;
-          entries.push(entry);
+    if (targets.length === 0) return { drained: 0, failed: 0 };
+
+    // Worker-pool draining. We spawn up to `concurrency` workers; each
+    // worker repeatedly:
+    //   1. Picks the next env with budget remaining (round-robin),
+    //      atomically claiming one slot of that env's per-tick budget.
+    //   2. Pops one entry and processes it.
+    //   3. Repeats until pickNextEnv returns null.
+    //
+    // This pattern gives us both invariants the prior two designs traded
+    // off:
+    //   - Single-env bursts use the full `concurrency` budget. All
+    //     workers can pull from one env, processing `concurrency` entries
+    //     in parallel.
+    //   - The number of entries in "popped-but-not-acked" (DRAINING)
+    //     state at any moment is bounded by the worker count, i.e.
+    //     `concurrency` — same blast radius as the pre-batch
+    //     one-pop-per-env model. A process crash mid-tick strands at
+    //     most `concurrency` entries for stale-sweep to recover, not
+    //     `maxOrgsPerTick × drainBatchSize`.
+    //
+    // Fairness: pickNextEnv advances a cursor by 1 each successful pick,
+    // so workers round-robin across envs at the entry level. Combined
+    // with the per-env budget cap, an env contributes at most
+    // `drainBatchSize` entries per tick regardless of how many workers
+    // are free — a heavy env can't starve siblings within a tick.
+    const remaining = new Map<string, number>();
+    const skip = new Set<string>(); // envs with empty queue or pop failure this tick
+    for (const envId of targets) remaining.set(envId, this.drainBatchSize);
+
+    let cursor = 0;
+    const pickNextEnv = (): string | null => {
+      for (let i = 0; i < targets.length; i++) {
+        const idx = (cursor + i) % targets.length;
+        const envId = targets[idx]!;
+        if (skip.has(envId)) continue;
+        const r = remaining.get(envId) ?? 0;
+        if (r > 0) {
+          remaining.set(envId, r - 1);
+          cursor = (idx + 1) % targets.length;
+          return envId;
         }
-        return { entries, popFailed };
-      }),
-    );
+      }
+      return null;
+    };
 
-    const popFailures = envBatches.reduce((n, b) => n + (b.popFailed ? 1 : 0), 0);
-    const allEntries = envBatches.flatMap((b) => b.entries);
-    if (allEntries.length === 0) {
-      return { drained: 0, failed: popFailures };
-    }
+    let drained = 0;
+    let failed = 0;
 
-    // Dispatch every popped entry through the shared pLimit so the
-    // global in-flight cap is `concurrency` regardless of how many envs
-    // contributed entries this tick. Per-entry errors are caught inside
-    // the closure so a single bad entry can't poison the tick — same
-    // safety net the old `processOneFromEnv` provided.
-    const inflight = allEntries.map((entry) =>
-      this.limit(async () => {
+    const worker = async (): Promise<void> => {
+      while (true) {
+        const envId = pickNextEnv();
+        if (envId === null) return;
+        let entry: BufferEntry | null;
         try {
-          return await this.processEntry(entry);
+          entry = await this.buffer.pop(envId);
+        } catch (err) {
+          // A pop failure on one env aborts that env's batch for this
+          // tick (don't keep hammering a broken Redis) and counts as
+          // exactly one failure — same as the pre-batch path on a pop
+          // blowup. Other envs continue.
+          this.logger.error("MollifierDrainer.pop failed", { envId, err });
+          skip.add(envId);
+          failed += 1;
+          continue;
+        }
+        if (!entry) {
+          // Queue exhausted between scheduling and this pop. Mark the
+          // env skipped so siblings aren't held up by repeated empty pops.
+          skip.add(envId);
+          continue;
+        }
+        try {
+          const outcome = await this.processEntry(entry);
+          if (outcome === "drained") drained += 1;
+          else failed += 1;
         } catch (err) {
           this.logger.error("MollifierDrainer.processEntry failed", {
-            envId: entry.envId,
+            envId,
             runId: entry.runId,
             err,
           });
-          return "failed" as const;
+          failed += 1;
         }
-      }),
-    );
-
-    const results = await Promise.all(inflight);
-    return {
-      drained: results.filter((r) => r === "drained").length,
-      failed: results.filter((r) => r === "failed").length + popFailures,
+      }
     };
+
+    const totalBudget = targets.length * this.drainBatchSize;
+    const workerCount = Math.min(this.concurrency, totalBudget);
+    await Promise.all(Array.from({ length: workerCount }, () => worker()));
+
+    return { drained, failed };
   }
 
   start(): void {
