Stop a spent claim budget starving the same job types

apps/worker/src/__tests__/health.test.ts

@@ -148,10 +148,12 @@ describe('buildHealthResponse', () => {
     // ─── busy is not stalled ───────────────────────────────────────────────
     //
     // THE REGRESSION THIS MODEL EXISTS FOR, and the one that got the previous
-    // attempt reverted. Two facts together: poll() stamps lastPollTime and then
-    // awaits its jobs, and claimJobsByType awaits each TYPE's batch sequentially
-    // (worker.ts) — so a single poll may legitimately run the SUM of every
-    // registered type's timeout. Against the config in index.ts that is 930s.
+    // attempt reverted. poll() stamps lastPollTime and then awaits every job it
+    // claimed, so a single poll legitimately runs as long as its slowest job.
+    // When claimJobsByType awaited each type's batch in turn, that was the sum
+    // of every registered type's timeout, 930s against the config in index.ts.
+    // Claims now all happen before dispatch, but the status writes after each
+    // handler are untimed, so the poll still has no fixed upper bound.
     //
     // Both earlier bounds were therefore wrong: `now - lastPollTime > 60s`, and
     // its replacement `pollDuration > max(jobTimeouts) + 60s` = 360s. Either one
@@ -168,7 +170,8 @@ describe('buildHealthResponse', () => {
             ...over,
         });
 
-        // 930_000 is the real sequential worst case for the ten registered types.
+        // 930_000 was the worst case when each type's batch ran in turn. It stays
+        // as a long-poll case: the probe must not care how the poll batches.
         it.each([61_000, 360_001, 600_000, 930_000, 1_200_000])(
             'reports 200 after %ims of polling while jobs are in flight and none is overdue',
             (pollAgeMs) => {
@@ -740,20 +743,20 @@ describe('a poll that keeps failing is not healthy', () => {
 });
 
 // The claim clock runs from the later of "this poll started" and "a job last
-// settled". claimJobsByType awaits each type's batch sequentially, so after a
-// long batch finishes, the remaining claim queries each run with
-// activeJobCount === 0 while pollDuration already carries that batch. Measuring
-// the whole poll reported a healthy worker stalled — the reverted defect, moved
-// rather than removed.
+// settled". Claims all happen before dispatch, but after the last job of a long
+// batch settles there is a short gap before the poll returns. In that gap
+// activeJobCount === 0 while pollDuration already carries the batch. Measuring
+// the whole poll would report a healthy worker stalled, the reverted defect
+// moved rather than removed.
 describe('a settled job resets the claim clock', () => {
     it('stays 200 when a batch settled recently, however long the poll has run', () => {
         const afterLongBatch: WorkerHealthStatus = {
             ...WORKER_HEALTH,
             activeJobCount: 0,
             pollStartedAt: new Date(NOW - 118_000),
             lastPollTime: new Date(NOW - 118_000),
-            // The AI_RESPONSE batch finished a moment ago; the poll is now
-            // issuing the next type's claim query.
+            // The AI_RESPONSE batch finished a moment ago and the poll is
+            // about to return.
             lastJobSettledAt: new Date(NOW - 200),
         };
 

apps/worker/src/health.ts

@@ -78,14 +78,14 @@ export const STALE_POLL_MS = 60_000;
  * query, so a poll that has not returned in a minute with no job to blame is
  * blocked on the database rather than working.
  *
- * Deliberately NOT a bound on poll duration in general. claimJobsByType awaits
- * each type's batch sequentially, so one poll can legitimately run the sum of
- * every registered type's timeout — 930s against the real worker config. An
- * earlier version of this module bounded the whole poll at
- * `max(jobTimeouts) + 60s` = 360s and reported a healthy worker stalled, which
- * is what got the previous attempt at this endpoint reverted. Whether the jobs
- * themselves are overdue is the worker's own question to answer, and it answers
- * it per job in `overdueJobCount`.
+ * Deliberately NOT a bound on poll duration in general. A poll claims every
+ * type up front and then waits for the whole batch, so it legitimately runs as
+ * long as its slowest job, and longer if a status write is slow. An earlier
+ * version of this module bounded the whole poll at `max(jobTimeouts) + 60s`,
+ * back when claiming awaited each type's batch in turn, and reported a healthy
+ * worker stalled. That is what got the previous attempt at this endpoint
+ * reverted. Whether the jobs themselves are overdue is the worker's own
+ * question to answer, and it answers it per job in `overdueJobCount`.
  */
 export const CLAIM_STALL_MS = 60_000;
 
@@ -439,16 +439,15 @@ export function buildHealthResponse(
 
     if (workerHealth.pollStartedAt) {
         // Nothing claimed and still not back: the claim query itself is blocked.
-        // With jobs in flight the duration says nothing, since a poll waits out
-        // every type's batch in turn.
+        // With jobs in flight the duration says nothing, since a poll waits for
+        // its whole claimed batch to settle.
         // Measured from the later of "this poll started" and "a job last
-        // settled", NOT from the start of the poll. claimJobsByType awaits each
-        // type's batch sequentially, so after a 118s batch of AI_RESPONSE jobs
-        // finishes, the nine remaining claim queries each run with
-        // activeJobCount === 0 and a pollDuration already carrying those 118s.
-        // Bounding the whole poll reported a healthy worker as stalled — the
-        // same mistake that got the previous two attempts reverted, moved rather
-        // than removed.
+        // settled", NOT from the start of the poll. Claims all happen before
+        // dispatch, but once the last job of a 118s batch settles there is a
+        // short gap before the poll returns. In that gap activeJobCount is 0 and
+        // the poll already carries those 118s. Bounding the whole poll would read
+        // that as a blocked claim and report a healthy worker as stalled, the
+        // same mistake that got the previous two attempts reverted.
         const claimingSince = Math.max(
             workerHealth.pollStartedAt.getTime(),
             workerHealth.lastJobSettledAt?.getTime() ?? 0,

packages/outpost/queue/src/__tests__/worker-concurrency.test.ts

@@ -816,3 +816,207 @@ describe('a recovering worker stops reporting failure at the claim, not at the e
         await vi.advanceTimersByTimeAsync(0);
     });
 });
+
+// ─── Claim-before-process (#232) ────────────────────────────────────────────
+
+/**
+ * The global budget used to be charged after `processClaimedJobs` had already
+ * awaited the batch, so work that had finished still held slots. With a broad
+ * backlog the budget ran out partway down `this.handlers` insertion order and
+ * the same tail types were never reached — deterministically, because the next
+ * poll started from the same place against the same state.
+ */
+describe('Worker claim budget and ordering', () => {
+    let worker: InstanceType<typeof Worker>;
+
+    /** Types in registration order; the last few are the ones that starved. */
+    const TYPES = [
+        JobType.AI_RESPONSE,
+        JobType.ESCALATION,
+        JobType.SLA_CHECK,
+        JobType.TRACKER_SYNC,
+        JobType.JOB_CLEANUP,
+    ];
+
+    /** `claimJobsForType` interpolates type third: jsonb map, default timeout, type, limit. */
+    const claimedType = (call: unknown[]) => call[3] as string;
+
+    beforeEach(() => {
+        vi.clearAllMocks();
+        vi.useFakeTimers();
+        mockPrisma.$executeRaw.mockResolvedValue(0);
+        mockPrismaJob.updateMany.mockResolvedValue({ count: 1 });
+        mockPrismaJob.update.mockResolvedValue({});
+    });
+
+    afterEach(async () => {
+        if (worker) await worker.stop();
+        vi.useRealTimers();
+    });
+
+    it('reaches every registered type across polls when every type has a backlog', async () => {
+        worker = new Worker({
+            pollIntervalMs: 10,
+            maxConcurrency: 4, // smaller than the backlog, so the budget must run out
+            batchSize: 2,
+            concurrencyByType: Object.fromEntries(TYPES.map((t) => [t, 2])),
+        });
+
+        for (const t of TYPES) worker.on(t, async () => ({ success: true }));
+
+        // Every type always has work waiting.
+        let seq = 0;
+        mockPrisma.$queryRaw.mockImplementation((...call: unknown[]) => {
+            const type = claimedType(call);
+            return Promise.resolve([
+                makeJobRow({ id: `${type}-${seq++}`, type }),
+                makeJobRow({ id: `${type}-${seq++}`, type }),
+            ]);
+        });
+
+        worker.start();
+        // Several polls: any fair scheme reaches all five well inside this.
+        await vi.advanceTimersByTimeAsync(200);
+
+        const reached = new Set(mockPrisma.$queryRaw.mock.calls.map(claimedType));
+        for (const t of TYPES) {
+            expect(reached, `type ${t} was never claimed`).toContain(t);
+        }
+    });
+
+    it('issues every claim before any handler runs, so a slow type cannot delay the rest', async () => {
+        worker = new Worker({
+            pollIntervalMs: 10,
+            maxConcurrency: 10,
+            batchSize: 1,
+            concurrencyByType: { [JobType.AI_RESPONSE]: 2, [JobType.ESCALATION]: 2 },
+        });
+
+        const events: string[] = [];
+        let releaseSlow: () => void = () => {};
+        const slowDone = new Promise<void>((r) => (releaseSlow = r));
+
+        worker.on(JobType.AI_RESPONSE, async () => {
+            events.push('ai-handler-start');
+            await slowDone;
+            return { success: true };
+        });
+        worker.on(JobType.ESCALATION, async () => {
+            events.push('esc-handler-start');
+            return { success: true };
+        });
+
+        mockPrisma.$queryRaw.mockImplementation((...call: unknown[]) => {
+            const type = claimedType(call);
+            events.push(`claim:${type}`);
+            if (type === JobType.AI_RESPONSE)
+                return Promise.resolve([makeJobRow({ id: 'ai-1', type })]);
+            if (type === JobType.ESCALATION)
+                return Promise.resolve([makeJobRow({ id: 'esc-1', type })]);
+            return Promise.resolve([]);
+        });
+
+        worker.start();
+        await vi.advanceTimersByTimeAsync(0);
+
+        // The ESCALATION claim must be issued even though the AI handler is parked.
+        // Both events must be present, or the index comparison below proves nothing.
+        expect(events).toContain(`claim:${JobType.ESCALATION}`);
+        expect(events).toContain('ai-handler-start');
+        expect(events.indexOf(`claim:${JobType.ESCALATION}`)).toBeLessThan(
+            events.indexOf('ai-handler-start'),
+        );
+
+        releaseSlow();
+        await vi.advanceTimersByTimeAsync(0);
+    });
+
+    it('never claims more than the global ceiling across all types in one poll', async () => {
+        // Budget 4 with a per-type limit of 2 means the ceiling can only be
+        // honoured by summing across types: no single type can reach it alone.
+        worker = new Worker({
+            pollIntervalMs: 10,
+            maxConcurrency: 4,
+            batchSize: 5,
+            concurrencyByType: Object.fromEntries(TYPES.map((t) => [t, 2])),
+        });
+
+        // Every handler parks, so the first poll cannot finish and cannot
+        // schedule a second one. Whatever is claimed is one poll's worth.
+        // Released before the assertions so `stop()` can drain in afterEach.
+        let release: () => void = () => {};
+        const parked = new Promise<void>((r) => (release = r));
+        for (const t of TYPES)
+            worker.on(t, async () => {
+                await parked;
+                return { success: true };
+            });
+
+        let claimed = 0;
+        const typesClaimed = new Set<string>();
+        mockPrisma.$queryRaw.mockImplementation((...call: unknown[]) => {
+            const type = claimedType(call);
+            const limit = call[4] as number;
+            const rows = Array.from({ length: limit }, (_, i) =>
+                makeJobRow({ id: `${type}-${i}`, type }),
+            );
+            claimed += rows.length;
+            typesClaimed.add(type);
+            return Promise.resolve(rows);
+        });
+
+        worker.start();
+        await vi.advanceTimersByTimeAsync(0);
+
+        const claimedInOnePoll = claimed;
+        const spannedTypes = typesClaimed.size;
+        release();
+        await vi.advanceTimersByTimeAsync(0);
+
+        // Lower bound first: a claim path that did nothing would satisfy the
+        // ceiling trivially.
+        expect(claimedInOnePoll).toBeGreaterThan(0);
+        expect(spannedTypes).toBeGreaterThan(1);
+        expect(claimedInOnePoll).toBeLessThanOrEqual(4);
+    });
+
+    it('dispatches jobs already claimed when a later claim fails', async () => {
+        // Claiming every type before processing any of them means a rejection
+        // partway through must not discard the rows already claimed: they are
+        // PROCESSING with a live token, and dropping them leaves them for the
+        // reclaim sweep, which consumes an attempt.
+        worker = new Worker({
+            pollIntervalMs: 10,
+            maxConcurrency: 10,
+            batchSize: 1,
+            concurrencyByType: {
+                [JobType.AI_RESPONSE]: 2,
+                [JobType.ESCALATION]: 2,
+                [JobType.SLA_CHECK]: 2,
+            },
+        });
+
+        const ran: string[] = [];
+        for (const t of [JobType.AI_RESPONSE, JobType.ESCALATION, JobType.SLA_CHECK])
+            worker.on(t, async (payload) => {
+                ran.push(String((payload as { id?: string }).id ?? 'x'));
+                return { success: true };
+            });
+
+        let call = 0;
+        mockPrisma.$queryRaw.mockImplementation((...args: unknown[]) => {
+            const type = claimedType(args);
+            call += 1;
+            if (call === 3) return Promise.reject(new Error('connection pool timeout'));
+            return Promise.resolve([
+                makeJobRow({ id: `${type}-1`, type, payload: { id: `${type}-1` } }),
+            ]);
+        });
+
+        worker.start();
+        await vi.advanceTimersByTimeAsync(0);
+
+        // The two claims that succeeded before the failure must still have run.
+        expect(ran).toHaveLength(2);
+    });
+});

packages/outpost/queue/src/types.ts

@@ -262,10 +262,10 @@ export interface WorkerHealthStatus {
     /**
      * In-flight jobs that have outlived their OWN timeout plus a grace.
      *
-     * The honest liveness signal. Poll duration is not: claimJobsByType awaits
-     * each type's batch sequentially, so a single poll may legitimately run the
-     * sum of every registered type's timeout. Non-zero here means the timeout
-     * machinery failed, not that the worker is slow.
+     * The honest liveness signal. Poll duration is not: a poll waits for its
+     * whole claimed batch, so it legitimately runs as long as its slowest job
+     * plus an untimed status write. Non-zero here means the timeout machinery
+     * failed, not that the worker is slow.
      */
     overdueJobCount: number;
     /** When a job last settled, whatever the outcome. Null before the first one. */