@shipfox/api-runners 13.0.0 → 13.1.0

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@shipfox/api-runners",
   "license": "MIT",
-  "version": "13.0.0",
+  "version": "13.1.0",
   "repository": {
     "type": "git",
     "url": "git+https://github.com/ShipfoxHQ/shipfox.git",
@@ -25,7 +25,7 @@
     "@shipfox/api-auth-context": "12.2.0",
     "@shipfox/api-auth-dto": "12.0.0",
     "@shipfox/api-runners-dto": "13.0.0",
-    "@shipfox/api-workflows-dto": "13.0.0",
+    "@shipfox/api-workflows-dto": "13.1.0",
     "@shipfox/config": "1.2.4",
     "@shipfox/inter-module": "0.2.3",
     "@shipfox/node-drizzle": "0.3.5",

package/src/db/job-executions.test.ts

@@ -482,6 +482,79 @@ describe('claimPendingJobExecution', () => {
     expect(runner?.firstClaimedAt).toBeInstanceOf(Date);
   });
 
+  it('releases a provisioned runner reservation on its first claim only', async () => {
+    const provisionerId = crypto.randomUUID();
+    const providerRunnerId = `provisioned-runner-${crypto.randomUUID()}`;
+    const [reservation] = await db()
+      .insert(reservations)
+      .values({
+        workspaceId,
+        provisionerId,
+        requiredLabels: sessionLabels,
+        count: 2,
+        expiresAt: new Date(Date.now() + 60_000),
+      })
+      .returning({id: reservations.id});
+    if (!reservation) throw new Error('Expected reservation');
+
+    const runner = await providerRunnerFactory.create({
+      workspaceId,
+      provisionerId,
+      providerRunnerId,
+      reservationId: reservation.id,
+      runnerSessionId,
+      state: 'running',
+      labels: sessionLabels,
+    });
+    await db()
+      .update(runnerSessions)
+      .set({
+        registrationTokenKind: 'ephemeral',
+        maxClaims: 2,
+        provisionerId,
+        providerRunnerId,
+      })
+      .where(eq(runnerSessions.id, runnerSessionId));
+
+    const firstPending = await pendingJobFactory.create({workspaceId});
+    const firstClaim = await claimPendingJobExecution({
+      workspaceId,
+      runnerSessionId,
+      maxClaims: 2,
+    });
+    const [afterFirstClaim] = await db()
+      .select({count: reservations.count})
+      .from(reservations)
+      .where(eq(reservations.id, reservation.id));
+    const [runnerAfterFirstClaim] = await db()
+      .select({reservationReleasedAt: providerRunners.reservationReleasedAt})
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    const firstReleaseAt = runnerAfterFirstClaim?.reservationReleasedAt;
+
+    const secondPending = await pendingJobFactory.create({workspaceId});
+    const secondClaim = await claimPendingJobExecution({
+      workspaceId,
+      runnerSessionId,
+      maxClaims: 2,
+    });
+    const [afterSecondClaim] = await db()
+      .select({count: reservations.count})
+      .from(reservations)
+      .where(eq(reservations.id, reservation.id));
+    const [runnerAfterSecondClaim] = await db()
+      .select({reservationReleasedAt: providerRunners.reservationReleasedAt})
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+
+    expect(firstClaim?.jobExecutionId).toBe(firstPending.jobExecutionId);
+    expect(afterFirstClaim?.count).toBe(1);
+    expect(firstReleaseAt).toBeInstanceOf(Date);
+    expect(secondClaim?.jobExecutionId).toBe(secondPending.jobExecutionId);
+    expect(afterSecondClaim?.count).toBe(1);
+    expect(runnerAfterSecondClaim?.reservationReleasedAt).toEqual(firstReleaseAt);
+  });
+
   it('allows a manual session to claim repeatedly', async () => {
     const first = await pendingJobFactory.create({workspaceId});
     const second = await pendingJobFactory.create({workspaceId});
@@ -1037,7 +1110,7 @@ describe('reconcileTerminalJobExecution', () => {
     expect(rows[0]?.cancellationRequestedAt).not.toBeNull();
   });
 
-  it('releases an already-terminal runner reservation after cancelling its last lease', async () => {
+  it('does not double-release a reservation when a claimed runner becomes terminal', async () => {
     const provisionerId = crypto.randomUUID();
     const providerRunnerId = crypto.randomUUID();
     const [reservation] = await db()
@@ -1046,7 +1119,7 @@ describe('reconcileTerminalJobExecution', () => {
         workspaceId,
         provisionerId,
         requiredLabels: sessionLabels,
-        count: 1,
+        count: 2,
         expiresAt: new Date(Date.now() + 60_000),
       })
       .returning({id: reservations.id});
@@ -1057,12 +1130,29 @@ describe('reconcileTerminalJobExecution', () => {
       providerRunnerId,
       reservationId: reservation.id,
       runnerSessionId,
-      state: 'terminated',
-      terminatedAt: new Date(),
+      state: 'running',
     });
+    await db()
+      .update(runnerSessions)
+      .set({registrationTokenKind: 'ephemeral', maxClaims: 1, provisionerId, providerRunnerId})
+      .where(eq(runnerSessions.id, runnerSessionId));
     const pending = await pendingJobFactory.create({workspaceId});
-    const claimed = await claimPendingJobExecution({workspaceId, runnerSessionId, maxClaims: null});
+    const claimed = await claimPendingJobExecution({workspaceId, runnerSessionId, maxClaims: 1});
     expect(claimed?.jobExecutionId).toBe(pending.jobExecutionId);
+    const [afterClaim] = await db()
+      .select({count: reservations.count})
+      .from(reservations)
+      .where(eq(reservations.id, reservation.id));
+    const [runnerAfterClaim] = await db()
+      .select({reservationReleasedAt: providerRunners.reservationReleasedAt})
+      .from(providerRunners)
+      .where(eq(providerRunners.providerRunnerId, providerRunnerId));
+    expect(afterClaim?.count).toBe(1);
+    expect(runnerAfterClaim?.reservationReleasedAt).toBeInstanceOf(Date);
+    await db()
+      .update(providerRunners)
+      .set({state: 'terminated', terminatedAt: new Date()})
+      .where(eq(providerRunners.providerRunnerId, providerRunnerId));
     await db()
       .update(runningJobExecutions)
       .set({provisionerId, providerRunnerId})
@@ -1070,9 +1160,11 @@ describe('reconcileTerminalJobExecution', () => {
 
     await reconcileTerminalJobExecution({jobExecutionId: pending.jobExecutionId});
 
-    expect(
-      await db().select().from(reservations).where(eq(reservations.id, reservation.id)),
-    ).toHaveLength(0);
+    const [afterTerminal] = await db()
+      .select({count: reservations.count})
+      .from(reservations)
+      .where(eq(reservations.id, reservation.id));
+    expect(afterTerminal?.count).toBe(1);
     const [runner] = await db()
       .select({reservationReleasedAt: providerRunners.reservationReleasedAt})
       .from(providerRunners)

package/src/db/job-executions.ts

@@ -38,9 +38,14 @@ import {
 } from '#metrics/instance.js';
 import type {Tx} from './db.js';
 import {db} from './db.js';
-import {releaseTerminalRunnerInstanceReservationsByIds} from './reservations.js';
+import {lockRunnerReservationAdvisoryKeysTx} from './reservation-locks.js';
+import {
+  releaseReservationUnits,
+  releaseTerminalRunnerInstanceReservationsByIds,
+} from './reservations.js';
 import {runnersOutbox} from './schema/outbox.js';
 import {pendingJobExecutions} from './schema/pending-job-executions.js';
+import {reservations} from './schema/reservations.js';
 import {providerRunners} from './schema/runner-instances.js';
 import {runnerSessions} from './schema/runner-sessions.js';
 import {runningJobExecutions} from './schema/running-job-executions.js';
@@ -262,6 +267,32 @@ export async function claimPendingJobExecution(params: {
       providerRunnerId = session.providerRunnerId;
     }
 
+    const runnerInstanceCondition = runnerInstanceId
+      ? eq(providerRunners.id, runnerInstanceId)
+      : provisionerId && providerRunnerId
+        ? and(
+            eq(providerRunners.provisionerId, provisionerId),
+            eq(providerRunners.providerRunnerId, providerRunnerId),
+          )
+        : null;
+    if (runnerInstanceCondition && provisionerId) {
+      const [runner] = await tx
+        .select({
+          reservationId: providerRunners.reservationId,
+          intendedReservationId: providerRunners.intendedReservationId,
+        })
+        .from(providerRunners)
+        .where(runnerInstanceCondition)
+        .limit(1);
+      await lockRunnerReservationAdvisoryKeysTx(tx, {
+        provisionerId,
+        reservationIds: [runner?.reservationId, runner?.intendedReservationId].filter(
+          (reservationId): reservationId is string =>
+            reservationId !== null && reservationId !== undefined,
+        ),
+      });
+    }
+
     // `id` is a uuidv7 (time-ordered), so it is a deterministic FIFO tiebreaker
     // for rows sharing a created_at within a batch. Lock only the FIFO candidate before
     // attempting its execution advisory lock; putting pg_try_advisory_xact_lock in this
@@ -328,14 +359,6 @@ export async function claimPendingJobExecution(params: {
       launchKind: params.maxClaims === null ? 'manual' : 'unknown',
     };
 
-    const runnerInstanceCondition = runnerInstanceId
-      ? eq(providerRunners.id, runnerInstanceId)
-      : provisionerId && providerRunnerId
-        ? and(
-            eq(providerRunners.provisionerId, provisionerId),
-            eq(providerRunners.providerRunnerId, providerRunnerId),
-          )
-        : null;
     if (runnerInstanceCondition) {
       const [claimedRunner] = await tx
         .update(providerRunners)
@@ -350,6 +373,8 @@ export async function claimPendingJobExecution(params: {
           provider: providerRunners.providerKind,
           launchKind: providerRunners.launchKind,
           runnerInstanceId: providerRunners.id,
+          reservationId: providerRunners.reservationId,
+          intendedReservationId: providerRunners.intendedReservationId,
           sessionCreatedAtEpochMs: sql<number | null>`(
             select extract(epoch from ${runnerSessions.createdAt})::double precision * 1000
             from ${runnerSessions}
@@ -360,6 +385,47 @@ export async function claimPendingJobExecution(params: {
         queueTimeObservation.provider = claimedRunner.provider;
         queueTimeObservation.launchKind = claimedRunner.launchKind;
       }
+      if (claimedRunner?.isFirstClaim) {
+        const [releasedRunner] = await tx
+          .update(providerRunners)
+          .set({
+            intendedReservationId: null,
+            reservationReleasedAt: sql`now()`,
+            updatedAt: sql`now()`,
+          })
+          .where(
+            and(
+              eq(providerRunners.id, claimedRunner.runnerInstanceId),
+              or(
+                isNotNull(providerRunners.reservationId),
+                isNotNull(providerRunners.intendedReservationId),
+              ),
+              isNull(providerRunners.reservationReleasedAt),
+            ),
+          )
+          .returning({
+            provisionerId: providerRunners.provisionerId,
+          });
+        const reservationId = claimedRunner.intendedReservationId ?? claimedRunner.reservationId;
+        if (releasedRunner?.provisionerId && reservationId) {
+          const [reservation] = await tx
+            .select({workspaceId: reservations.workspaceId})
+            .from(reservations)
+            .where(
+              and(
+                eq(reservations.id, reservationId),
+                eq(reservations.provisionerId, releasedRunner.provisionerId),
+              ),
+            )
+            .limit(1);
+          if (reservation)
+            await releaseReservationUnits(tx, {
+              workspaceId: reservation.workspaceId,
+              provisionerId: releasedRunner.provisionerId,
+              releases: [{reservationId, count: 1}],
+            });
+        }
+      }
       if (
         claimedRunner?.isFirstClaim &&
         claimedRunner.firstClaimedAt !== null &&