@shipfox/api-runners 12.3.0 → 12.4.1

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,13 +1,19 @@
-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,
   RunnerInstanceAlreadyAssignedError,
   RunnerInstanceNotAssignableError,
 } from '#core/errors.js';
+import {
+  type ProviderRunnerAssignmentObservation,
+  type RunnerAssignmentSurface,
+  recordProviderRunnerControlSessionToAssignment,
+} from '#metrics/instance.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';
@@ -18,7 +24,17 @@ export async function assignRunnerInstances(params: {
   reservationId: string;
   runnerInstanceIds: string[];
 }): Promise<string[]> {
-  return await db().transaction(async (tx) => assignRunnerInstancesTx(tx, params));
+  const result = await db().transaction(async (tx) =>
+    assignRunnerInstancesTx(tx, {...params, surface: 'provisioner'}),
+  );
+  for (const observation of result.controlSessionToAssignment)
+    recordProviderRunnerControlSessionToAssignment(observation);
+  return result.runnerInstanceIds;
+}
+
+interface AssignRunnerInstancesTxResult {
+  runnerInstanceIds: string[];
+  controlSessionToAssignment: ProviderRunnerAssignmentObservation[];
 }
 
 export async function assignRunnerInstancesTx(
@@ -27,19 +43,24 @@ export async function assignRunnerInstancesTx(
     provisionerId: string;
     reservationId: string;
     runnerInstanceIds: string[];
+    surface: RunnerAssignmentSurface;
   },
-): Promise<string[]> {
+): Promise<AssignRunnerInstancesTxResult> {
   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,
+      provider: providerRunners.providerKind,
+      launchKind: providerRunners.launchKind,
       labels: providerRunners.labels,
       state: providerRunners.state,
     })
@@ -53,11 +74,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)
@@ -68,7 +94,10 @@ export async function assignRunnerInstancesTx(
           inArray(providerRunners.id, runnerInstanceIds),
         ),
       );
-    return params.runnerInstanceIds;
+    return {
+      runnerInstanceIds: params.runnerInstanceIds,
+      controlSessionToAssignment: [],
+    };
   }
 
   const [reservation] = await tx
@@ -78,6 +107,7 @@ export async function assignRunnerInstancesTx(
       and(
         eq(reservations.id, params.reservationId),
         eq(reservations.provisionerId, params.provisionerId),
+        eq(reservations.kind, 'launch'),
       ),
     )
     .limit(1)
@@ -85,10 +115,13 @@ export async function assignRunnerInstancesTx(
   if (!reservation) throw new ReservationNotFoundError(params.reservationId);
   if (reservation.expiresAt <= new Date()) throw new ReservationExpiredError(params.reservationId);
   if (runnerRows.length !== runnerInstanceIds.length)
-    throw new RunnerInstanceNotAssignableError(runnerInstanceIds[0] ?? '');
+    throw new RunnerInstanceNotAssignableError(runnerInstanceIds[0] ?? '', 'runner-not-found');
 
   const activeControlSessions = await tx
-    .select({runnerInstanceId: runnerControlSessions.runnerInstanceId})
+    .select({
+      runnerInstanceId: runnerControlSessions.runnerInstanceId,
+      createdAt: runnerControlSessions.createdAt,
+    })
     .from(runnerControlSessions)
     .where(
       and(
@@ -102,21 +135,32 @@ export async function assignRunnerInstancesTx(
   const runnerInstanceIdsWithControlSession = new Set(
     activeControlSessions.map((session) => session.runnerInstanceId),
   );
+  const controlSessionCreatedAtByRunner = new Map(
+    activeControlSessions.map((session) => [session.runnerInstanceId, session.createdAt]),
+  );
   const runners = runnerRows.map((runner) => ({
     ...runner,
     controlSessionId: runnerInstanceIdsWithControlSession.has(runner.id) ? runner.id : null,
+    controlSessionCreatedAt: controlSessionCreatedAtByRunner.get(runner.id) ?? null,
   }));
 
   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 +168,6 @@ export async function assignRunnerInstancesTx(
           ),
           and(
             eq(providerRunners.intendedReservationId, reservation.id),
-            notInArray(providerRunners.id, runnerInstanceIds),
             isNull(providerRunners.reservationReleasedAt),
             notInArray(providerRunners.state, [...terminalStates]),
           ),
@@ -132,18 +175,19 @@ export async function assignRunnerInstancesTx(
       ),
     );
   if ((assignedCount[0]?.count ?? 0) + newRunners.length > reservation.count)
-    throw new RunnerInstanceNotAssignableError(newRunners[0]?.id ?? '');
+    throw new RunnerInstanceNotAssignableError(newRunners[0]?.id ?? '', 'capacity-exhausted');
   for (const runner of newRunners) {
-    if (
-      runner.state !== 'running' ||
-      !runner.providerRunnerId ||
-      !runner.controlSessionId ||
-      !reservation.requiredLabels.every((label) => runner.labels.includes(label))
-    )
-      throw new RunnerInstanceNotAssignableError(runner.id);
+    if (runner.state !== 'running')
+      throw new RunnerInstanceNotAssignableError(runner.id, 'runner-not-running');
+    if (!runner.providerRunnerId)
+      throw new RunnerInstanceNotAssignableError(runner.id, 'provider-identity-missing');
+    if (!runner.controlSessionId)
+      throw new RunnerInstanceNotAssignableError(runner.id, 'control-session-not-active');
+    if (!reservation.requiredLabels.every((label) => runner.labels.includes(label)))
+      throw new RunnerInstanceNotAssignableError(runner.id, 'labels-mismatch');
   }
   if (newRunners.length > 0) {
-    await tx
+    const assignedRows = await tx
       .update(providerRunners)
       .set({
         workspaceId: reservation.workspaceId,
@@ -157,7 +201,183 @@ export async function assignRunnerInstancesTx(
           providerRunners.id,
           newRunners.map((runner) => runner.id),
         ),
-      );
+      )
+      .returning({id: providerRunners.id, assignedAt: providerRunners.assignedAt});
+    const runnersById = new Map(newRunners.map((runner) => [runner.id, runner]));
+    const controlSessionToAssignment: ProviderRunnerAssignmentObservation[] = [];
+    for (const assigned of assignedRows) {
+      const runner = runnersById.get(assigned.id);
+      const controlSessionCreatedAt = runner?.controlSessionCreatedAt;
+      if (!runner || !assigned.assignedAt || !controlSessionCreatedAt) continue;
+      controlSessionToAssignment.push({
+        durationMilliseconds: assigned.assignedAt.getTime() - controlSessionCreatedAt.getTime(),
+        provider: runner.provider,
+        launchKind: runner.launchKind,
+        surface: params.surface,
+        runnerInstanceId: runner.id,
+      });
+    }
+    return {
+      runnerInstanceIds: params.runnerInstanceIds,
+      controlSessionToAssignment,
+    };
   }
-  return params.runnerInstanceIds;
+  return {
+    runnerInstanceIds: params.runnerInstanceIds,
+    controlSessionToAssignment: [],
+  };
+}
+
+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};
 }