@shipfox/api-runners 12.6.0 → 13.0.0

package/src/db/job-executions.ts

@@ -1,7 +1,6 @@
 import {
   RUNNER_JOB_CLAIMED,
   RUNNER_JOB_LEASE_EXPIRED,
-  RUNNER_JOB_QUEUED,
   type RunnersEventMap,
   type RunnerToolCapabilitiesDto,
 } from '@shipfox/api-runners-dto';
@@ -39,6 +38,7 @@ import {
 } from '#metrics/instance.js';
 import type {Tx} from './db.js';
 import {db} from './db.js';
+import {releaseTerminalRunnerInstanceReservationsByIds} from './reservations.js';
 import {runnersOutbox} from './schema/outbox.js';
 import {pendingJobExecutions} from './schema/pending-job-executions.js';
 import {providerRunners} from './schema/runner-instances.js';
@@ -53,6 +53,39 @@ async function lockJobExecution(tx: Tx, jobExecutionId: string): Promise<void> {
   );
 }
 
+async function releaseReservationsForTerminalRunningRows(
+  tx: Tx,
+  rows: ReadonlyArray<{
+    provisionerId: string | null;
+    providerRunnerId: string | null;
+  }>,
+): Promise<void> {
+  // Converge after the lease fact changes. The reservation helper locks each provider runner
+  // and releases only when the runner is terminal and no uncancelled lease remains.
+  const providerRunnerIdsByProvisionerId = new Map<string, Set<string>>();
+  for (const row of rows) {
+    if (row.provisionerId === null || row.providerRunnerId === null) continue;
+    const providerRunnerIds =
+      providerRunnerIdsByProvisionerId.get(row.provisionerId) ?? new Set<string>();
+    providerRunnerIds.add(row.providerRunnerId);
+    providerRunnerIdsByProvisionerId.set(row.provisionerId, providerRunnerIds);
+  }
+
+  for (const provisionerId of [...providerRunnerIdsByProvisionerId.keys()].sort()) {
+    const providerRunnerIds = providerRunnerIdsByProvisionerId.get(provisionerId);
+    if (!providerRunnerIds) continue;
+    await releaseTerminalRunnerInstanceReservationsByIds(tx, {
+      workspaceId: null,
+      provisionerId,
+      providerRunnerIds: [...providerRunnerIds].sort(),
+      requireUnlinkedSession: false,
+      // Lock and re-check the runner row locally so lease finalization remains retryable
+      // without a workflow/runner scope lock.
+      requireTerminalState: false,
+    });
+  }
+}
+
 export interface EnqueueJobExecutionParams {
   workspaceId: string;
   workflowRunId: string;
@@ -61,6 +94,7 @@ export interface EnqueueJobExecutionParams {
   jobExecutionId: string;
   projectId: string;
   requiredLabels: string[];
+  queuedAt: Date;
 }
 
 export async function getWorkspaceJobCounts(params: {
@@ -91,11 +125,9 @@ export async function getWorkspaceJobCounts(params: {
   }));
 }
 
-// Idempotent while the job execution is still pending: a duplicate jobExecutionId already in
-// `runners_pending_jobs` is a no-op. Temporal retries the enqueue activity at-least-once, so a
-// unique-violation throw on a retry-after-lost-result would permanently fail a healthy execution.
-// The per-execution advisory lock serializes retries with lease expiry/reconciliation, and the
-// durable lease-expired event prevents a retry from re-queueing an execution already reaped.
+// The workflows outbox delivers queue facts at least once. The per-execution advisory lock and
+// pending/running checks make replay a no-op across the claim transition. A durable lease-expired
+// event also prevents a delayed replay from resurrecting an execution already reaped locally.
 export async function enqueueJobExecution(params: EnqueueJobExecutionParams): Promise<void> {
   const requiredLabels = [...canonicalizeLabels(params.requiredLabels)];
   if (requiredLabels.length === 0) throw new EmptyRequiredLabelsError();
@@ -103,6 +135,13 @@ export async function enqueueJobExecution(params: EnqueueJobExecutionParams): Pr
   const enqueued = await db().transaction(async (tx) => {
     await lockJobExecution(tx, params.jobExecutionId);
 
+    const [running] = await tx
+      .select({jobExecutionId: runningJobExecutions.jobExecutionId})
+      .from(runningJobExecutions)
+      .where(eq(runningJobExecutions.jobExecutionId, params.jobExecutionId))
+      .limit(1);
+    if (running) return false;
+
     const [leaseExpired] = await tx
       .select({id: runnersOutbox.id})
       .from(runnersOutbox)
@@ -125,26 +164,12 @@ export async function enqueueJobExecution(params: EnqueueJobExecutionParams): Pr
         jobExecutionId: params.jobExecutionId,
         projectId: params.projectId,
         requiredLabels,
+        createdAt: params.queuedAt,
       })
       .onConflictDoNothing({target: pendingJobExecutions.jobExecutionId})
       .returning({createdAt: pendingJobExecutions.createdAt});
 
-    // A retry that hits the conflict inserts nothing: the first enqueue already
-    // emitted the queued event (durably, in the outbox), so re-emitting would
-    // only add a redundant row the subscriber coalesces away. Skip it.
-    if (!inserted) return false;
-
-    await writeOutboxEvent<RunnersEventMap>(tx, runnersOutbox, {
-      type: RUNNER_JOB_QUEUED,
-      payload: {
-        workflowRunId: params.workflowRunId,
-        workflowRunAttemptId: params.workflowRunAttemptId,
-        jobId: params.jobId,
-        jobExecutionId: params.jobExecutionId,
-        queuedAt: inserted.createdAt.toISOString(),
-      },
-    });
-    return true;
+    return inserted !== undefined;
   });
 
   if (enqueued) jobExecutionEnqueuedCount.add(1);
@@ -402,43 +427,18 @@ async function touchRunnerSessionLiveness(params: {
     );
 }
 
-/**
- * Releases a job execution's lease when the orchestration workflow finalizes it: deletes the
- * running-job-execution row AND any lingering pending row for the same execution, in one tx.
- * Idempotent (0-row no-op), no token scope (the workflow is authoritative), and
- * emits no event: the workflow already owns the outcome. Sweeping the pending row
- * too closes the at-least-once window where an enqueue retry left an orphan that a
- * later claim would otherwise pick up for an already-finished job execution.
- */
-export async function releaseJobExecution(params: {jobExecutionId: string}): Promise<void> {
-  await db().transaction(async (tx) => {
-    await lockJobExecution(tx, params.jobExecutionId);
-
-    // Delete pending before running to match `claimPendingJobExecution`'s lock-acquisition
-    // order (it locks the pending row first, then the running row). A concurrent
-    // claim picking up an orphan pending row for this same job execution would otherwise
-    // deadlock against the reverse order here.
-    await tx
-      .delete(pendingJobExecutions)
-      .where(eq(pendingJobExecutions.jobExecutionId, params.jobExecutionId));
-    await tx
-      .delete(runningJobExecutions)
-      .where(eq(runningJobExecutions.jobExecutionId, params.jobExecutionId));
-  });
-}
-
 /**
  * Reaps stale leases (bounded by `limit`), emitting one
  * `runners.job.lease_expired` event per reaped job execution.
  *
  * The cutoff is re-checked in the DELETE, not just the locking subquery, so a
  * heartbeat landing mid-call spares the live row. Each reaped execution also sweeps its
- * pending row: a failed best-effort `releaseJobExecution` would otherwise leave an orphan
- * that a later claim re-runs as an already-finished job execution.
+ * pending row so an at-least-once queue replay cannot leave an orphan.
  *
- * Locks pending-then-running to match `claimPendingJobExecution`, `releaseJobExecution`, and
- * `reconcileTerminalJobExecution`. The stale candidate scan intentionally does not lock running
- * rows first; the running-row DELETE re-checks the stale predicate after the pending-row sweep.
+ * Locks pending-then-running to match `claimPendingJobExecution` and
+ * `reconcileTerminalJobExecution`. The stale candidate scan intentionally does not lock
+ * running rows first; the running-row DELETE re-checks the stale predicate after the
+ * pending-row sweep.
  */
 export async function expireStuckJobExecutions(params: {
   thresholdSeconds: number;
@@ -515,10 +515,14 @@ export async function expireStuckJobExecutions(params: {
         workflowRunAttemptId: runningJobExecutions.workflowRunAttemptId,
         jobId: runningJobExecutions.jobId,
         jobExecutionId: runningJobExecutions.jobExecutionId,
+        provisionerId: runningJobExecutions.provisionerId,
+        providerRunnerId: runningJobExecutions.providerRunnerId,
       });
 
     if (deleted.length === 0) return [];
 
+    await releaseReservationsForTerminalRunningRows(tx, deleted);
+
     await writeOutboxEvents<RunnersEventMap>(
       tx,
       runnersOutbox,
@@ -533,7 +537,12 @@ export async function expireStuckJobExecutions(params: {
       })),
     );
 
-    return deleted;
+    return deleted.map(({workflowRunId, workflowRunAttemptId, jobId, jobExecutionId}) => ({
+      workflowRunId,
+      workflowRunAttemptId,
+      jobId,
+      jobExecutionId,
+    }));
   });
 
   if (reaped.length > 0) jobExecutionLeaseExpiredCount.add(reaped.length);
@@ -781,8 +790,8 @@ export async function recordHeartbeat(params: {
 
 /**
  * Reconciles a terminal job execution with runner state in one transaction: removes its pending
- * queue row and requests cancellation on its running lease, if either exists. The operation is
- * idempotent and preserves the first cancellation-request timestamp.
+ * queue row, requests cancellation on its running lease, and converges any linked reservation.
+ * The operation is idempotent and preserves the first cancellation-request timestamp.
  */
 export async function reconcileTerminalJobExecution(params: {
   jobExecutionId: string;
@@ -796,25 +805,17 @@ export async function reconcileTerminalJobExecution(params: {
     await tx
       .delete(pendingJobExecutions)
       .where(eq(pendingJobExecutions.jobExecutionId, params.jobExecutionId));
-    await tx
+    const cancelledRunningRows = await tx
       .update(runningJobExecutions)
       .set({
         cancellationRequestedAt: sql`COALESCE(${runningJobExecutions.cancellationRequestedAt}, now())`,
       })
-      .where(eq(runningJobExecutions.jobExecutionId, params.jobExecutionId));
-  });
-}
-
-export async function cancelRunnerJobs(params: {jobIds: string[]}): Promise<void> {
-  if (params.jobIds.length === 0) return;
+      .where(eq(runningJobExecutions.jobExecutionId, params.jobExecutionId))
+      .returning({
+        provisionerId: runningJobExecutions.provisionerId,
+        providerRunnerId: runningJobExecutions.providerRunnerId,
+      });
 
-  await db().transaction(async (tx) => {
-    await tx.delete(pendingJobExecutions).where(inArray(pendingJobExecutions.jobId, params.jobIds));
-    await tx
-      .update(runningJobExecutions)
-      .set({
-        cancellationRequestedAt: sql`COALESCE(${runningJobExecutions.cancellationRequestedAt}, now())`,
-      })
-      .where(inArray(runningJobExecutions.jobId, params.jobIds));
+    await releaseReservationsForTerminalRunningRows(tx, cancelledRunningRows);
   });
 }

package/src/db/provisioner-tokens.test.ts

@@ -10,6 +10,7 @@ import {
   touchProvisionerLastSeen,
 } from '#db/provisioner-tokens.js';
 import {provisionerTokens} from '#db/schema/provisioner-tokens.js';
+import {reservations} from '#db/schema/reservations.js';
 import {runnerActivationTokens} from '#db/schema/runner-activation-tokens.js';
 import {runnerBootstrapTokens, runnerControlSessions} from '#db/schema/runner-control-sessions.js';
 import {providerRunners} from '#db/schema/runner-instances.js';
@@ -105,10 +106,22 @@ describe('provisioner token db', () => {
   it("cascades revocation to the provisioner's unclaimed runner credentials and sessions", async () => {
     const workspaceId = crypto.randomUUID();
     const provisioner = await provisionerTokenFactory.create({workspaceId});
+    const [reservation] = await db()
+      .insert(reservations)
+      .values({
+        workspaceId,
+        provisionerId: provisioner.id,
+        requiredLabels: ['linux'],
+        count: 1,
+        expiresAt: new Date(Date.now() + 60_000),
+      })
+      .returning();
+    if (!reservation) throw new Error('Expected reservation');
     const runner = await providerRunnerFactory.create({
       workspaceId,
       provisionerId: provisioner.id,
       runnerSessionId: null,
+      reservationId: reservation.id,
     });
     const future = new Date(Date.now() + 60_000);
     const tokenValue = (name: string) => hashOpaqueToken(`${name}-${crypto.randomUUID()}`);
@@ -192,6 +205,10 @@ describe('provisioner token db', () => {
       .select()
       .from(providerRunners)
       .where(eq(providerRunners.id, runner.id));
+    const [releasedReservation] = await db()
+      .select()
+      .from(reservations)
+      .where(eq(reservations.id, reservation.id));
     const [revokedSession] = await db()
       .select()
       .from(runnerSessions)
@@ -207,6 +224,8 @@ describe('provisioner token db', () => {
     expect(activation?.revokedAt).toBeInstanceOf(Date);
     expect(terminatedRunner).toMatchObject({state: 'terminated'});
     expect(terminatedRunner?.terminatedAt).toBeInstanceOf(Date);
+    expect(terminatedRunner?.reservationReleasedAt).toBeInstanceOf(Date);
+    expect(releasedReservation).toBeUndefined();
     expect(revokedSession?.revokedAt).toBeInstanceOf(Date);
     expect(preservedSession?.revokedAt).toBeNull();
   });

package/src/db/provisioner-tokens.ts

@@ -12,6 +12,7 @@ import type {
 } from '#core/entities/provisioner-token.js';
 import {ProvisionerAdminIdempotencyKeyReuseError} from '#core/errors.js';
 import {db} from './db.js';
+import {releaseTerminalRunnerInstanceReservationsByIds} from './reservations.js';
 import {runnersAdminCommandResults} from './schema/admin-command-results.js';
 import {runnersOutbox} from './schema/outbox.js';
 import {provisionerTokens, toProvisionerToken} from './schema/provisioner-tokens.js';
@@ -62,7 +63,7 @@ async function cascadeProvisionerRevocation(tx: Tx, provisionerId: string) {
         ),
       ),
     );
-  await tx
+  const terminatedRows = await tx
     .update(providerRunners)
     .set({state: 'terminated', terminatedAt: sql`now()`, updatedAt: sql`now()`})
     .where(
@@ -70,7 +71,17 @@ async function cascadeProvisionerRevocation(tx: Tx, provisionerId: string) {
         eq(providerRunners.provisionerId, provisionerId),
         isNull(providerRunners.runnerSessionId),
       ),
-    );
+    )
+    .returning({id: providerRunners.id});
+  const runnerInstanceIds = terminatedRows.map((row) => row.id);
+  if (runnerInstanceIds.length > 0) {
+    await releaseTerminalRunnerInstanceReservationsByIds(tx, {
+      workspaceId: null,
+      provisionerId,
+      runnerInstanceIds,
+      requireUnlinkedSession: false,
+    });
+  }
   await tx
     .update(runnerSessions)
     .set({revokedAt: sql`now()`})

package/src/db/reservations.test.ts

@@ -535,6 +535,66 @@ describe('pollDemandAndReserve', () => {
     expect(result.reservations).toEqual([]);
   });
 
+  it('deducts a terminal assigned runner reservation unit as pending', async () => {
+    const reservation = await createIntendedReservation({
+      workspaceId,
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+    await createIdleRunner({
+      labels: ['linux'],
+      reservationId: reservation.id,
+      state: 'terminated',
+    });
+    await createPendingJobs(1, ['linux', 'gpu']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux-gpu', ['linux', 'gpu'], 1)],
+    });
+
+    expect(result.reservations).toEqual([]);
+  });
+
+  it('counts only active assigned runners when deducting provisioner reservations', async () => {
+    const [reservation] = await db()
+      .insert(reservations)
+      .values({
+        workspaceId,
+        provisionerId,
+        requiredLabels: ['linux'],
+        count: 2,
+        expiresAt: new Date(Date.now() + 60_000),
+      })
+      .returning({id: reservations.id});
+    if (!reservation) throw new Error('Expected reservation');
+    await createIdleRunner({
+      labels: ['linux'],
+      reservationId: reservation.id,
+      state: 'running',
+    });
+    await createIdleRunner({
+      labels: ['linux'],
+      reservationId: reservation.id,
+      state: 'failed',
+    });
+    await createPendingJobs(2, ['linux', 'gpu']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 2,
+      ttlSeconds: 60,
+      templates: [template('linux-gpu', ['linux', 'gpu'], 2)],
+    });
+
+    expect(result.reservations).toHaveLength(1);
+    expect(result.reservations[0]?.count).toBe(1);
+    expect(result.stats[0]).toMatchObject({queued: 2, reserved: 1});
+  });
+
   it('binds a runner whose intended reservation has passed its activation grace period', async () => {
     const intendedReservation = await createIntendedReservation({
       workspaceId,
@@ -1716,6 +1776,84 @@ describe('pollDemandAndReserve', () => {
     expect(rows[0]?.count).toBe(1);
   });
 
+  it('releases units from multiple reservations in one transaction', async () => {
+    const [first] = await db()
+      .insert(reservations)
+      .values({
+        workspaceId,
+        provisionerId,
+        requiredLabels: ['linux'],
+        count: 3,
+        expiresAt: new Date(Date.now() + 60_000),
+      })
+      .returning({id: reservations.id});
+    const [second] = await db()
+      .insert(reservations)
+      .values({
+        workspaceId,
+        provisionerId,
+        requiredLabels: ['linux', 'gpu'],
+        count: 2,
+        expiresAt: new Date(Date.now() + 60_000),
+      })
+      .returning({id: reservations.id});
+    if (!first || !second) throw new Error('Expected reservations');
+
+    const released = await db().transaction((tx) =>
+      releaseReservationUnits(tx, {
+        workspaceId,
+        provisionerId,
+        releases: [
+          {reservationId: first.id, count: 2},
+          {reservationId: second.id, count: 1},
+        ],
+      }),
+    );
+
+    const rows = await reservationsForTest();
+    expect(released).toBe(3);
+    expect(rows).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({id: first.id, count: 1}),
+        expect.objectContaining({id: second.id, count: 1}),
+      ]),
+    );
+  });
+
+  it('does not strand units when concurrent releases drain one reservation', async () => {
+    const [reservation] = await db()
+      .insert(reservations)
+      .values({
+        workspaceId,
+        provisionerId,
+        requiredLabels: ['linux'],
+        count: 2,
+        expiresAt: new Date(Date.now() + 60_000),
+      })
+      .returning({id: reservations.id});
+    if (!reservation) throw new Error('Expected reservation');
+
+    const [firstReleased, secondReleased] = await Promise.all([
+      db().transaction((tx) =>
+        releaseReservationUnits(tx, {
+          workspaceId,
+          provisionerId,
+          releases: [{reservationId: reservation.id, count: 1}],
+        }),
+      ),
+      db().transaction((tx) =>
+        releaseReservationUnits(tx, {
+          workspaceId,
+          provisionerId,
+          releases: [{reservationId: reservation.id, count: 1}],
+        }),
+      ),
+    ]);
+
+    expect(firstReleased + secondReleased).toBe(2);
+    expect(await reservationsForTest()).toHaveLength(0);
+  });
+
   it('deletes reservations when releasing all remaining units', async () => {
     await reservationFactory.create({
       workspaceId,