@shipfox/api-runners 12.1.0 → 12.3.0

package/src/db/reservations.ts

@@ -10,6 +10,8 @@ import {
   isNotNull,
   isNull,
   lt,
+  notExists,
+  or,
   sql,
 } from 'drizzle-orm';
 import type {Tx} from './db.js';
@@ -66,6 +68,12 @@ export interface InstallationPollDemandAndReserveParams {
   onReservations?: (reservations: ReservationGrant[]) => void;
 }
 
+type DemandScope = 'installation' | 'workspace';
+
+type PollDemandAndReserveLockedParams = PollDemandAndReserveParams & {
+  scope: DemandScope;
+};
+
 interface NormalizedTemplate {
   templateKey: string;
   labels: string[];
@@ -91,7 +99,7 @@ export async function pollDemandAndReserveTx(
   params: PollDemandAndReserveParams,
 ): Promise<{stats: DemandStat[]; reservations: ReservationGrant[]}> {
   await tx.execute(sql`select pg_advisory_xact_lock(hashtext(${params.workspaceId}))`);
-  return await pollDemandAndReserveLockedTx(tx, params);
+  return await pollDemandAndReserveLockedTx(tx, {...params, scope: 'workspace'});
 }
 
 export async function pollInstallationDemandAndReserve(
@@ -125,6 +133,7 @@ export async function pollInstallationDemandAndReserve(
           availableSlots: template.remainingSlots,
         })),
         capabilityWindowSeconds: params.capabilityWindowSeconds,
+        scope: 'installation',
       });
     });
     results.push(result);
@@ -157,7 +166,7 @@ function consumeInstallationTemplateSlots(
 
 async function pollDemandAndReserveLockedTx(
   tx: Tx,
-  params: PollDemandAndReserveParams,
+  params: PollDemandAndReserveLockedParams,
 ): Promise<{stats: DemandStat[]; reservations: ReservationGrant[]}> {
   let demandRows = (
     await tx
@@ -264,6 +273,7 @@ async function pollDemandAndReserveLockedTx(
         requiredLabels: demand.requiredLabels,
         count: grant,
         workspaceId: params.workspaceId,
+        scope: params.scope,
       });
       reservedAfterGrant += grant;
       grants.push({
@@ -295,18 +305,61 @@ async function bindIdleRunnerInstancesTx(
     workspaceId: string;
     requiredLabels: string[];
     count: number;
+    scope: DemandScope;
   },
 ): Promise<void> {
+  // Reservation expiry is the activation grace deadline. A live reservation
+  // protects a runner that is still booting; an expired or missing row is stale.
+  const canBindRunner = () =>
+    and(
+      or(
+        and(isNull(providerRunners.workspaceId), isNull(providerRunners.reservationId)),
+        and(
+          isNotNull(providerRunners.reservationId),
+          params.scope === 'installation'
+            ? undefined
+            : or(
+                isNull(providerRunners.workspaceId),
+                eq(providerRunners.workspaceId, params.workspaceId),
+              ),
+          notExists(
+            tx
+              .select({id: reservations.id})
+              .from(reservations)
+              .where(
+                and(
+                  eq(reservations.id, providerRunners.reservationId),
+                  gt(reservations.expiresAt, sql`now()`),
+                ),
+              ),
+          ),
+        ),
+      ),
+      or(
+        isNull(providerRunners.intendedReservationId),
+        notExists(
+          tx
+            .select({id: reservations.id})
+            .from(reservations)
+            .where(
+              and(
+                eq(reservations.id, providerRunners.intendedReservationId),
+                gt(reservations.expiresAt, sql`now()`),
+              ),
+            ),
+        ),
+      ),
+      isNull(providerRunners.reservationReleasedAt),
+      isNull(providerRunners.runnerSessionId),
+    );
+
   const idleRunners = await tx
     .select({id: providerRunners.id})
     .from(providerRunners)
     .where(
       and(
         eq(providerRunners.provisionerId, params.provisionerId),
-        isNull(providerRunners.workspaceId),
-        isNull(providerRunners.reservationId),
-        isNull(providerRunners.intendedReservationId),
-        isNull(providerRunners.runnerSessionId),
+        canBindRunner(),
         isNotNull(providerRunners.providerRunnerId),
         eq(providerRunners.state, 'running'),
         arrayContains(providerRunners.labels, params.requiredLabels),
@@ -331,11 +384,26 @@ async function bindIdleRunnerInstancesTx(
 
   if (idleRunners.length === 0) return;
 
+  await tx
+    .update(runnerActivationTokens)
+    .set({revokedAt: sql`now()`})
+    .where(
+      and(
+        inArray(
+          runnerActivationTokens.runnerInstanceId,
+          idleRunners.map((runner) => runner.id),
+        ),
+        isNull(runnerActivationTokens.consumedAt),
+        isNull(runnerActivationTokens.revokedAt),
+      ),
+    );
+
   await tx
     .update(providerRunners)
     .set({
       workspaceId: params.workspaceId,
       reservationId: params.reservationId,
+      intendedReservationId: null,
       assignedAt: sql`now()`,
       updatedAt: sql`now()`,
     })
@@ -345,9 +413,7 @@ async function bindIdleRunnerInstancesTx(
           providerRunners.id,
           idleRunners.map((runner) => runner.id),
         ),
-        isNull(providerRunners.workspaceId),
-        isNull(providerRunners.reservationId),
-        isNull(providerRunners.runnerSessionId),
+        canBindRunner(),
       ),
     );
 }
@@ -432,16 +498,26 @@ async function deleteReservationsWithCleanupTx(
 ): Promise<number> {
   if (ids.length === 0) return 0;
 
-  const assignedRunners = await tx
-    .select({id: providerRunners.id, reservationId: providerRunners.reservationId})
+  const affectedRunners = await tx
+    .select({
+      id: providerRunners.id,
+      reservationId: providerRunners.reservationId,
+      intendedReservationId: providerRunners.intendedReservationId,
+      runnerSessionId: providerRunners.runnerSessionId,
+      reservationReleasedAt: providerRunners.reservationReleasedAt,
+    })
     .from(providerRunners)
     .where(
-      and(
-        inArray(providerRunners.reservationId, ids),
-        isNull(providerRunners.runnerSessionId),
-        isNull(providerRunners.reservationReleasedAt),
+      or(
+        and(
+          inArray(providerRunners.reservationId, ids),
+          isNull(providerRunners.runnerSessionId),
+          isNull(providerRunners.reservationReleasedAt),
+        ),
+        inArray(providerRunners.intendedReservationId, ids),
       ),
     )
+    .orderBy(asc(providerRunners.id))
     .for('update');
   const reservationRows = await tx
     .select({id: reservations.id})
@@ -457,8 +533,20 @@ async function deleteReservationsWithCleanupTx(
 
   if (reservationIds.length === 0) return 0;
 
-  const assignedRunnerIds = assignedRunners
-    .filter((runner) => runner.reservationId && reservationIds.includes(runner.reservationId))
+  const assignedRunnerIds = affectedRunners
+    .filter(
+      (runner) =>
+        runner.reservationId &&
+        reservationIds.includes(runner.reservationId) &&
+        !runner.runnerSessionId &&
+        !runner.reservationReleasedAt,
+    )
+    .map((runner) => runner.id);
+  const intendedRunnerIds = affectedRunners
+    .filter(
+      (runner) =>
+        runner.intendedReservationId && reservationIds.includes(runner.intendedReservationId),
+    )
     .map((runner) => runner.id);
   if (assignedRunnerIds.length > 0) {
     await tx
@@ -487,6 +575,17 @@ async function deleteReservationsWithCleanupTx(
         ),
       );
   }
+  if (intendedRunnerIds.length > 0) {
+    await tx
+      .update(providerRunners)
+      .set({intendedReservationId: null, updatedAt: sql`now()`})
+      .where(
+        and(
+          inArray(providerRunners.id, intendedRunnerIds),
+          inArray(providerRunners.intendedReservationId, reservationIds),
+        ),
+      );
+  }
 
   const deleted = await tx
     .delete(reservations)