fix(selfhost): stop the queue drain from claiming new work during shutdown; add shutdown headroom to compose (#9007)

docker-compose.yml

@@ -87,6 +87,16 @@ services:
         INSTALL_AI_CLIS: "${INSTALL_AI_CLIS:-true}"
         INSTALL_VISUAL_REVIEW: "${INSTALL_VISUAL_REVIEW:-false}"
     restart: unless-stopped
+    # #9007: a review pass routinely runs tens of seconds to minutes (AI calls + GitHub round-trips), but
+    # Docker's default stop_grace_period is only 10s -- far shorter than the queue's own graceful drain
+    # (pg-queue.ts/sqlite-queue.ts stop(), which now stops CLAIMING new work immediately but still waits for
+    # whatever pass is already in flight to finish). Without headroom here, every `docker compose
+    # restart`/`stop`/redeploy during active review traffic sends SIGKILL before that drain completes,
+    # severing the in-flight pass mid-way and leaving its Redis locks/DB state exactly where it stood --
+    # the root cause behind a whole class of "PR looks stuck" incidents. 300s covers the slowest realistic
+    # pass (see AI_REVIEW_LOCK_TTL_SECONDS's own 30-minute crash-safety backstop for the orphan case this
+    # doesn't fully eliminate) while still bounding a genuinely wedged shutdown.
+    stop_grace_period: 300s
     <<: *default-logging
     ports:
       # Remove this when using the caddy profile — Caddy becomes the public listener.

src/selfhost/pg-queue.ts

@@ -263,6 +263,13 @@ export function createPgQueue(
   const processingTimeoutMs = queueProcessingTimeoutMs();
 
   let running = false;
+  // #9007: distinct from `running`, which also governs pre-start()/direct pump()-drain() usage (binding.send()'s
+  // fire-and-forget kickOne(), and drain()'s own direct pump() call, both legitimately run a pump loop before
+  // start() has ever been called or after stop() -- gating the loop below on `running` itself broke exactly
+  // that path). `shuttingDown` instead means ONLY "stop() has been called and hasn't been superseded by a
+  // later start()" -- it starts false, flips true at the very top of stop() (before its own drain-wait), and
+  // resets false in start() so a stop()-then-restart cycle claims fresh again.
+  let shuttingDown = false;
   let active = 0;
   let activeBackground = 0;
   const activeJobIds = new Set<string>();
@@ -1472,8 +1479,15 @@ export function createPgQueue(
     if (active >= concurrency) return;
     active++;
     try {
-      while (await processOne()) {
-        /* drain due jobs */
+      // #9007: `shuttingDown` is re-checked on every iteration (not just at pump() entry) so that stop() --
+      // which sets it synchronously, before its own `while (active > 0)` wait -- actually bounds this loop to
+      // the job already in flight. Without this guard, a pump already draining due work keeps calling
+      // processOne() for as long as more due jobs exist, entirely blind to a concurrent stop(): shutdown then
+      // waits for the WHOLE backlog to drain rather than for in-flight work to finish, so under sustained load
+      // stop() never returns before the container's SIGKILL grace period expires -- the root cause of every
+      // review pass severed mid-flight by a deploy.
+      while (!shuttingDown && (await processOne())) {
+        /* drain due jobs, but stop as soon as shutdown begins */
       }
     } catch (error) {
       // claimNext()/reclaimExpiredProcessingJobs() run OUTSIDE processOne's own try/finally, so a raw pool
@@ -1542,6 +1556,7 @@ export function createPgQueue(
     start() {
       if (running) return;
       running = true;
+      shuttingDown = false; // #9007: a stop()-then-restart cycle claims fresh again
       const tick = (): void => {
         /* v8 ignore next */ // stop() clears the timer before the next tick can fire with running=false
         if (!running) return;
@@ -1561,6 +1576,10 @@ export function createPgQueue(
       );
     },
     async stop() {
+      // #9007: set BEFORE the drain-wait below (and before `running = false`) so any pump loop still actively
+      // draining due work sees it on its very next iteration and stops claiming new jobs immediately, rather
+      // than continuing to drain the whole backlog while this wait loop sits idle.
+      shuttingDown = true;
       running = false;
       if (timer) clearTimeout(timer);
       if (deadLetterReviveTimer) clearInterval(deadLetterReviveTimer);

src/selfhost/sqlite-queue.ts

@@ -272,6 +272,13 @@ export function createSqliteQueue(
       }),
     );
   let running = false;
+  // #9007: distinct from `running`, which also governs pre-start()/direct pump()-drain() usage (binding.send()'s
+  // fire-and-forget kickOne(), and drain()'s own direct pump() call, both legitimately run a pump loop before
+  // start() has ever been called or after stop() -- gating the loop below on `running` itself broke exactly
+  // that path). `shuttingDown` instead means ONLY "stop() has been called and hasn't been superseded by a
+  // later start()" -- it starts false, flips true at the very top of stop() (before its own drain-wait), and
+  // resets false in start() so a stop()-then-restart cycle claims fresh again.
+  let shuttingDown = false;
   let active = 0; // number of concurrent pump() loops currently draining jobs
   let activeBackground = 0;
   const activeJobIds = new Set<number>();
@@ -1153,8 +1160,15 @@ export function createSqliteQueue(
     if (active >= concurrency) return;
     active++;
     try {
-      while (await processOne()) {
-        /* keep draining due jobs */
+      // #9007: `shuttingDown` is re-checked on every iteration (not just at pump() entry) so that stop() --
+      // which sets it synchronously, before its own `while (active > 0)` wait -- actually bounds this loop to
+      // the job already in flight. Without this guard, a pump already draining due work keeps calling
+      // processOne() for as long as more due jobs exist, entirely blind to a concurrent stop(): shutdown then
+      // waits for the WHOLE backlog to drain rather than for in-flight work to finish, so under sustained load
+      // stop() never returns before the container's SIGKILL grace period expires -- the root cause of every
+      // review pass severed mid-flight by a deploy.
+      while (!shuttingDown && (await processOne())) {
+        /* keep draining due jobs, but stop as soon as shutdown begins */
       }
     } catch (error) {
       // claimNext()/reclaimExpiredProcessingJobs() run OUTSIDE processOne's own try/finally, so a raw driver
@@ -1230,6 +1244,7 @@ export function createSqliteQueue(
     start() {
       if (running) return;
       running = true;
+      shuttingDown = false; // #9007: a stop()-then-restart cycle claims fresh again
       const tick = (): void => {
         /* v8 ignore next */ // stop() clears the timer, so a tick never fires with running=false
         if (!running) return;
@@ -1246,6 +1261,10 @@ export function createSqliteQueue(
       foregroundLivenessTimer = setInterval(() => void releaseStaleForegroundDeferralsSafely(), foregroundLivenessConfig.checkIntervalMs);
     },
     async stop() {
+      // #9007: set BEFORE the drain-wait below (and before `running = false`) so any pump loop still actively
+      // draining due work sees it on its very next iteration and stops claiming new jobs immediately, rather
+      // than continuing to drain the whole backlog while this wait loop sits idle.
+      shuttingDown = true;
       running = false;
       if (timer) clearTimeout(timer);
       if (deadLetterReviveTimer) clearInterval(deadLetterReviveTimer);

test/unit/selfhost-pg-queue.test.ts

@@ -3108,6 +3108,51 @@ describe("createPgQueue (durable #977)", () => {
     }
   });
 
+  it("REGRESSION (#9007): stop() lets only the in-flight job finish -- it does not drain the whole preloaded backlog", async () => {
+    // Reproduces the shutdown bug directly: pump()'s `while (await processOne())` loop used to be blind to
+    // shutdown state, so a single already-looping pump kept claiming every subsequent due job regardless of a
+    // concurrent stop() -- shutdown was bounded by "drain everything still due", not by "finish what's
+    // already in flight". With concurrency:1 and a huge pollIntervalMs, start() fires exactly one kickAll(),
+    // which starts exactly one pump() that (pre-fix) would have claimed all 5 preloaded jobs back-to-back.
+    const m = makePool();
+    const consumed: string[] = [];
+    const q = createPgQueue(
+      m.pool,
+      async (j) => {
+        consumed.push(typeOf(j));
+        await new Promise((r) => setTimeout(r, 25));
+      },
+      { concurrency: 1, pollIntervalMs: 100_000 },
+    );
+    await q.init();
+    m.enqueueJob("1", { type: "a" });
+    m.enqueueJob("2", { type: "b" });
+    m.enqueueJob("3", { type: "c" });
+    m.enqueueJob("4", { type: "d" });
+    m.enqueueJob("5", { type: "e" });
+    q.start(); // the one tick's kickAll() starts the lone pump, which claims "a" and enters its 25ms consume
+    await new Promise((r) => setTimeout(r, 5));
+    await q.stop(); // must return once "a" finishes -- NOT after b/c/d/e also drain
+    expect(consumed).toEqual(["a"]);
+  });
+
+  it("REGRESSION (#9007): a pump kicked via send()/drain() BEFORE start() is ever called is unaffected by shuttingDown", async () => {
+    // Guards against a real mistake made while fixing the bug above: gating the pump loop on `running` itself
+    // (instead of a separate `shuttingDown` flag) would have made every pre-start() drain()/send() no-op,
+    // since `running` starts false and isn't set until start() is first called -- breaking the extremely
+    // common test/production pattern of `send()` + `drain()` with no start()/stop() lifecycle involved at all.
+    const m = makePool();
+    m.enqueueJob("1", { type: "never-started" });
+    let calls = 0;
+    const q = createPgQueue(m.pool, async () => {
+      calls += 1;
+    });
+    await q.init();
+    await q.binding.send(msg("never-started"));
+    await q.drain();
+    expect(calls).toBe(1);
+  });
+
   it("start() is idempotent", async () => {
     const { pool } = makePool();
     const q = createPgQueue(pool, async () => undefined, { pollIntervalMs: 100_000 });

test/unit/selfhost-sqlite-queue.test.ts

@@ -3352,6 +3352,42 @@ describe("createSqliteQueue (durable #980)", () => {
     expect(done).toBe(true);
   });
 
+  it("REGRESSION (#9007): stop() does not drain the whole backlog -- it lets only the in-flight job finish", async () => {
+    // Reproduces the bug directly: a single pump() invocation's `while (await processOne())` loop used to be
+    // blind to shutdown state, so once one pump claimed the first due job it kept claiming every subsequent
+    // due job regardless of a concurrent stop() -- shutdown wasn't bounded by "finish what's in flight", it
+    // was bounded by "drain everything still pending". With concurrency:1, sendBatch's 5 kickOne() calls all
+    // fire synchronously; only the FIRST actually starts a pump (the other 4 see active>=concurrency and
+    // no-op) -- so whether jobs 2-5 get processed during this one stop() call depends entirely on whether the
+    // loop condition checks shutdown state.
+    const consumed: string[] = [];
+    const q = createSqliteQueue(makeDriver(), async (message) => {
+      consumed.push((message as unknown as { type: string }).type);
+      await new Promise((r) => setTimeout(r, 25));
+    }, { concurrency: 1, pollIntervalMs: 100_000 }); // huge poll interval: no periodic kickAll to confound this
+    await q.binding.sendBatch([msg("a"), msg("b"), msg("c"), msg("d"), msg("e")].map((body) => ({ body })));
+    q.start();
+    await new Promise((r) => setTimeout(r, 5)); // let the lone pump claim "a" and enter its 25ms consume
+    await q.stop(); // must return once "a" finishes -- NOT after b/c/d/e also drain
+    expect(consumed).toEqual(["a"]);
+    expect(await q.size()).toBe(4); // b, c, d, e remain pending, untouched by this stop()
+  });
+
+  it("REGRESSION (#9007): a pump kicked via send()/drain() BEFORE start() is ever called is unaffected by shuttingDown", async () => {
+    // Guards against a real mistake made while fixing the bug above: gating the pump loop on `running` itself
+    // (instead of a separate `shuttingDown` flag) would have made every pre-start() drain()/send() no-op,
+    // since `running` starts false and isn't set until start() is first called -- breaking the extremely
+    // common test/production pattern of `send()` + `drain()` with no start()/stop() lifecycle involved at all.
+    let calls = 0;
+    const q = createSqliteQueue(makeDriver(), async () => {
+      calls += 1;
+    }, {});
+    await q.binding.send(msg("never-started"));
+    await q.drain();
+    expect(calls).toBe(1);
+    expect(await q.size()).toBe(0);
+  });
+
   describe("maintenance-admission pressure gating (#selfhost-runtime-pressure)", () => {
     const envKeys = [
       "MAINTENANCE_ADMISSION_ENABLED",