@shipfox/api-runners 12.2.0 → 12.4.0

package/src/db/runner-assignments.test.ts

@@ -1,6 +1,7 @@
 import {eq} from 'drizzle-orm';
 import {
   ReservationExpiredError,
+  ReservationNotFoundError,
   RunnerInstanceAlreadyAssignedError,
   RunnerInstanceNotAssignableError,
 } from '#core/errors.js';
@@ -41,6 +42,50 @@ describe('assignRunnerInstances', () => {
     });
   });
 
+  it('repairs a provider-reported reservation before assignment commit', async () => {
+    const reservation = await createReservation();
+    const [runner] = await db()
+      .insert(providerRunners)
+      .values({
+        provisionerId,
+        intendedReservationId: reservation.id,
+        reservationId: reservation.id,
+        providerRunnerId: crypto.randomUUID(),
+        labels: ['linux'],
+        state: 'running',
+        reportedAt: new Date(),
+      })
+      .returning();
+    if (!runner) throw new Error('Runner instance insert returned no row');
+    await db()
+      .insert(runnerControlSessions)
+      .values({
+        runnerInstanceId: runner.id,
+        provisionerId,
+        hashedToken: crypto.randomUUID(),
+        prefix: 'test',
+        expiresAt: new Date(Date.now() + 60_000),
+      });
+
+    const assigned = await assignRunnerInstances({
+      provisionerId,
+      reservationId: reservation.id,
+      runnerInstanceIds: [runner.id],
+    });
+
+    expect(assigned).toEqual([runner.id]);
+    const [stored] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(stored).toMatchObject({
+      workspaceId,
+      reservationId: reservation.id,
+      intendedReservationId: null,
+      assignedAt: expect.any(Date),
+    });
+  });
+
   it('is idempotent for concurrent retries of the same assignment', async () => {
     const reservation = await createReservation();
     const runner = await createEnrolledRunner();
@@ -80,6 +125,32 @@ describe('assignRunnerInstances', () => {
     expect(assigned).toEqual([runner.id]);
   });
 
+  it('is idempotent for a committed bound reservation assignment', async () => {
+    const reservation = await createReservation({kind: 'bound'});
+    const runner = await createEnrolledRunner();
+    await db()
+      .update(providerRunners)
+      .set({
+        workspaceId,
+        reservationId: reservation.id,
+        assignedAt: new Date(),
+      })
+      .where(eq(providerRunners.id, runner.id));
+
+    const assigned = await assignRunnerInstances({
+      provisionerId,
+      reservationId: reservation.id,
+      runnerInstanceIds: [runner.id],
+    });
+
+    expect(assigned).toEqual([runner.id]);
+    const [stored] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(stored?.intendedReservationId).toBeNull();
+  });
+
   it('rejects expired reservations', async () => {
     const reservation = await createReservation({expiresAt: new Date(Date.now() - 1_000)});
     const runner = await createEnrolledRunner();
@@ -93,6 +164,19 @@ describe('assignRunnerInstances', () => {
     await expect(assignment).rejects.toThrow(ReservationExpiredError);
   });
 
+  it('rejects bound reservations', async () => {
+    const reservation = await createReservation({kind: 'bound'});
+    const runner = await createEnrolledRunner();
+
+    const assignment = assignRunnerInstances({
+      provisionerId,
+      reservationId: reservation.id,
+      runnerInstanceIds: [runner.id],
+    });
+
+    await expect(assignment).rejects.toThrow(ReservationNotFoundError);
+  });
+
   it('rejects unenrolled or incompatible runners', async () => {
     const reservation = await createReservation({requiredLabels: ['linux', 'gpu']});
     const runner = await createEnrolledRunner({labels: ['linux']});
@@ -139,6 +223,25 @@ describe('assignRunnerInstances', () => {
     await expect(assignment).rejects.toThrow(RunnerInstanceNotAssignableError);
   });
 
+  it('charges capacity for a committed runner replayed alongside a new one', async () => {
+    const reservation = await createReservation();
+    const firstRunner = await createEnrolledRunner();
+    const secondRunner = await createEnrolledRunner();
+    await assignRunnerInstances({
+      provisionerId,
+      reservationId: reservation.id,
+      runnerInstanceIds: [firstRunner.id],
+    });
+
+    const assignment = assignRunnerInstances({
+      provisionerId,
+      reservationId: reservation.id,
+      runnerInstanceIds: [firstRunner.id, secondRunner.id],
+    });
+
+    await expect(assignment).rejects.toThrow(RunnerInstanceNotAssignableError);
+  });
+
   it('does not charge capacity for an already released reservation intent', async () => {
     const reservation = await createReservation();
     await db().insert(providerRunners).values({
@@ -202,7 +305,7 @@ describe('assignRunnerInstances', () => {
   });
 
   async function createReservation(
-    overrides: Partial<{expiresAt: Date; requiredLabels: string[]}> = {},
+    overrides: Partial<{expiresAt: Date; requiredLabels: string[]; kind: 'bound' | 'launch'}> = {},
   ) {
     const [reservation] = await db()
       .insert(reservations)
@@ -211,6 +314,7 @@ describe('assignRunnerInstances', () => {
         provisionerId,
         requiredLabels: overrides.requiredLabels ?? ['linux'],
         count: 1,
+        kind: overrides.kind ?? 'launch',
         expiresAt: overrides.expiresAt ?? new Date(Date.now() + 60_000),
       })
       .returning();

package/src/db/runner-assignments.ts

@@ -1,4 +1,4 @@
-import {and, eq, inArray, isNull, notInArray, or, sql} from 'drizzle-orm';
+import {and, eq, inArray, isNull, ne, notInArray, or, sql} from 'drizzle-orm';
 import {
   ReservationExpiredError,
   ReservationNotFoundError,
@@ -7,7 +7,8 @@ import {
 } from '#core/errors.js';
 import type {Tx} from './db.js';
 import {db} from './db.js';
-import {terminalStates} from './runner-instances.js';
+import {lockRunnerReservationAdvisoryKeysTx} from './reservation-locks.js';
+import {terminalStates} from './runner-states.js';
 import {reservations} from './schema/reservations.js';
 import {runnerControlSessions} from './schema/runner-control-sessions.js';
 import {providerRunners} from './schema/runner-instances.js';
@@ -30,15 +31,17 @@ export async function assignRunnerInstancesTx(
   },
 ): Promise<string[]> {
   const runnerInstanceIds = [...params.runnerInstanceIds].sort();
-  await tx.execute(
-    sql`select pg_advisory_xact_lock(hashtext(${`runners_assignment:${params.provisionerId}:${params.reservationId}`}))`,
-  );
+  await lockRunnerReservationAdvisoryKeysTx(tx, {
+    provisionerId: params.provisionerId,
+    reservationIds: [params.reservationId],
+  });
   const runnerRows = await tx
     .select({
       id: providerRunners.id,
       reservationId: providerRunners.reservationId,
       intendedReservationId: providerRunners.intendedReservationId,
       workspaceId: providerRunners.workspaceId,
+      assignedAt: providerRunners.assignedAt,
       providerRunnerId: providerRunners.providerRunnerId,
       labels: providerRunners.labels,
       state: providerRunners.state,
@@ -53,11 +56,16 @@ export async function assignRunnerInstancesTx(
     .for('update');
 
   // The assignment outlives its short reservation row. A retry after maintenance
-  // cleanup is complete when every owned runner already carries the requested assignment.
+  // cleanup is complete when every owned runner already carries the committed assignment.
   if (
     runnerInstanceIds.length > 0 &&
     runnerRows.length === runnerInstanceIds.length &&
-    runnerRows.every((runner) => runner.reservationId === params.reservationId)
+    runnerRows.every(
+      (runner) =>
+        runner.reservationId === params.reservationId &&
+        runner.assignedAt !== null &&
+        runner.workspaceId !== null,
+    )
   ) {
     await tx
       .update(providerRunners)
@@ -78,6 +86,7 @@ export async function assignRunnerInstancesTx(
       and(
         eq(reservations.id, params.reservationId),
         eq(reservations.provisionerId, params.provisionerId),
+        eq(reservations.kind, 'launch'),
       ),
     )
     .limit(1)
@@ -110,13 +119,20 @@ export async function assignRunnerInstancesTx(
   const alreadyAssigned = runners.filter((runner) => runner.reservationId !== null);
   if (alreadyAssigned.some((runner) => runner.reservationId !== reservation.id))
     throw new RunnerInstanceAlreadyAssignedError(alreadyAssigned[0]?.id ?? '');
-  const newRunners = runners.filter((runner) => runner.reservationId === null);
+  const newRunners = runners.filter(
+    (runner) =>
+      runner.reservationId === null ||
+      (runner.reservationId === reservation.id &&
+        (runner.assignedAt === null || runner.workspaceId === null)),
+  );
+  const newRunnerIds = newRunners.map((runner) => runner.id);
   const assignedCount = await tx
     .select({count: sql<number>`count(*)::int`})
     .from(providerRunners)
     .where(
       and(
         eq(providerRunners.provisionerId, params.provisionerId),
+        notInArray(providerRunners.id, newRunnerIds),
         or(
           and(
             eq(providerRunners.reservationId, reservation.id),
@@ -124,7 +140,6 @@ export async function assignRunnerInstancesTx(
           ),
           and(
             eq(providerRunners.intendedReservationId, reservation.id),
-            notInArray(providerRunners.id, runnerInstanceIds),
             isNull(providerRunners.reservationReleasedAt),
             notInArray(providerRunners.state, [...terminalStates]),
           ),
@@ -161,3 +176,157 @@ export async function assignRunnerInstancesTx(
   }
   return params.runnerInstanceIds;
 }
+
+export type RunnerReservationCapacityFailureReason =
+  | 'reservation-not-found'
+  | 'reservation-kind-mismatch'
+  | 'reservation-expired'
+  | 'capacity-exhausted';
+
+export interface RunnerReservationCapacityValidation {
+  acceptedByReservation: Map<string, number>;
+  unavailableByReservation: Map<
+    string,
+    {reason: RunnerReservationCapacityFailureReason; count: number}
+  >;
+}
+
+export async function validateRunnerReservationCapacityTx(
+  tx: Tx,
+  params: {
+    provisionerId: string;
+    requests: readonly {reservationId: string; count: number}[];
+  },
+  options: {advisoryLocksHeld?: boolean} = {},
+): Promise<RunnerReservationCapacityValidation> {
+  const requestedByReservation = new Map<string, number>();
+  for (const request of params.requests) {
+    if (request.count <= 0) continue;
+    requestedByReservation.set(
+      request.reservationId,
+      (requestedByReservation.get(request.reservationId) ?? 0) + request.count,
+    );
+  }
+  const reservationIds = [...requestedByReservation.keys()].sort();
+  if (reservationIds.length === 0)
+    return {acceptedByReservation: new Map(), unavailableByReservation: new Map()};
+
+  if (!options.advisoryLocksHeld)
+    await lockRunnerReservationAdvisoryKeysTx(tx, {
+      provisionerId: params.provisionerId,
+      reservationIds,
+    });
+
+  const reservationRows = await tx
+    .select({
+      id: reservations.id,
+      count: reservations.count,
+      expiresAt: reservations.expiresAt,
+      isExpired: sql<boolean>`${reservations.expiresAt} <= now()`,
+    })
+    .from(reservations)
+    .where(
+      and(
+        eq(reservations.provisionerId, params.provisionerId),
+        inArray(reservations.id, reservationIds),
+        eq(reservations.kind, 'launch'),
+      ),
+    )
+    .for('update');
+  const nonLaunchReservationRows = await tx
+    .select({id: reservations.id})
+    .from(reservations)
+    .where(
+      and(
+        eq(reservations.provisionerId, params.provisionerId),
+        inArray(reservations.id, reservationIds),
+        ne(reservations.kind, 'launch'),
+      ),
+    )
+    .for('update');
+  const reservationsById = new Map(
+    reservationRows.map((reservation) => [reservation.id, reservation]),
+  );
+  const nonLaunchReservationIds = new Set(
+    nonLaunchReservationRows.map((reservation) => reservation.id),
+  );
+  const requestedReservationIds = new Set(reservationIds);
+  const usedRunnerIdsByReservation = new Map<string, Set<string>>();
+  const usedRunnerRows = await tx
+    .select({
+      id: providerRunners.id,
+      reservationId: providerRunners.reservationId,
+      intendedReservationId: providerRunners.intendedReservationId,
+      state: providerRunners.state,
+    })
+    .from(providerRunners)
+    .where(
+      and(
+        eq(providerRunners.provisionerId, params.provisionerId),
+        isNull(providerRunners.reservationReleasedAt),
+        or(
+          inArray(providerRunners.reservationId, reservationIds),
+          and(
+            inArray(providerRunners.intendedReservationId, reservationIds),
+            notInArray(providerRunners.state, [...terminalStates]),
+          ),
+        ),
+      ),
+    );
+  for (const runner of usedRunnerRows) {
+    if (runner.reservationId && requestedReservationIds.has(runner.reservationId)) {
+      const usedRunnerIds =
+        usedRunnerIdsByReservation.get(runner.reservationId) ?? new Set<string>();
+      usedRunnerIds.add(runner.id);
+      usedRunnerIdsByReservation.set(runner.reservationId, usedRunnerIds);
+    }
+    if (
+      runner.intendedReservationId &&
+      requestedReservationIds.has(runner.intendedReservationId) &&
+      !terminalStates.includes(runner.state as (typeof terminalStates)[number])
+    ) {
+      const usedRunnerIds =
+        usedRunnerIdsByReservation.get(runner.intendedReservationId) ?? new Set<string>();
+      usedRunnerIds.add(runner.id);
+      usedRunnerIdsByReservation.set(runner.intendedReservationId, usedRunnerIds);
+    }
+  }
+
+  const acceptedByReservation = new Map<string, number>();
+  const unavailableByReservation = new Map<
+    string,
+    {reason: RunnerReservationCapacityFailureReason; count: number}
+  >();
+  for (const reservationId of reservationIds) {
+    const requested = requestedByReservation.get(reservationId) ?? 0;
+    const reservation = reservationsById.get(reservationId);
+    if (!reservation) {
+      unavailableByReservation.set(reservationId, {
+        reason: nonLaunchReservationIds.has(reservationId)
+          ? 'reservation-kind-mismatch'
+          : 'reservation-not-found',
+        count: requested,
+      });
+      continue;
+    }
+    if (reservation.isExpired) {
+      unavailableByReservation.set(reservationId, {
+        reason: 'reservation-expired',
+        count: requested,
+      });
+      continue;
+    }
+
+    const used = usedRunnerIdsByReservation.get(reservationId)?.size ?? 0;
+    const accepted = Math.min(requested, Math.max(0, reservation.count - used));
+    acceptedByReservation.set(reservationId, accepted);
+    if (accepted < requested) {
+      unavailableByReservation.set(reservationId, {
+        reason: 'capacity-exhausted',
+        count: requested - accepted,
+      });
+    }
+  }
+
+  return {acceptedByReservation, unavailableByReservation};
+}