@shipfox/api-runners 12.3.0 → 12.4.0

package/src/core/runner-control-sessions.ts

@@ -12,13 +12,17 @@ import {
 } from '#core/errors.js';
 import {sanitizeRunnerLabels} from '#core/runner-labels.js';
 import {db, type schema, type Tx} from '#db/db.js';
-import {assignRunnerInstancesTx} from '#db/runner-assignments.js';
-import {terminalStates} from '#db/runner-instances.js';
+import {
+  assignRunnerInstancesTx,
+  validateRunnerReservationCapacityTx,
+} from '#db/runner-assignments.js';
+import {terminalStates} from '#db/runner-states.js';
 import {provisionerTokens} from '#db/schema/provisioner-tokens.js';
 import {runnerBootstrapTokens, runnerControlSessions} from '#db/schema/runner-control-sessions.js';
 import {providerRunners} from '#db/schema/runner-instances.js';
 import {
   type RunnerReservationPromotionFailureReason,
+  recordRunnerReservationCapacityFailure,
   recordRunnerReservationPromotionFailure,
 } from '#metrics/index.js';
 import {issueRunnerActivationTokenTx} from './runner-activation.js';
@@ -45,27 +49,72 @@ export async function createRunnerInstancesWithBootstrapTokens(params: {
     reservationId?: string | null;
   }>;
   ttlSeconds: number;
-}): Promise<Array<{runnerInstanceId: string; bootstrapToken: string}>> {
+}): Promise<{
+  runnerInstances: Array<{
+    runnerInstanceId: string;
+    bootstrapToken: string;
+    requestIndex: number;
+  }>;
+  reservationUnavailable: boolean;
+}> {
   const expiresAt = new Date(Date.now() + params.ttlSeconds * 1000);
   return await db().transaction(async (tx) => {
+    const reservationValidation = await validateRunnerReservationCapacityTx(tx, {
+      provisionerId: params.provisionerId,
+      requests: params.runnerInstances.reduce(
+        (requests, runner) => {
+          if (!runner.reservationId) return requests;
+          const request = requests.find(
+            (candidate) => candidate.reservationId === runner.reservationId,
+          );
+          if (request) request.count += 1;
+          else requests.push({reservationId: runner.reservationId, count: 1});
+          return requests;
+        },
+        [] as Array<{reservationId: string; count: number}>,
+      ),
+    });
+    for (const {reason, count} of reservationValidation.unavailableByReservation.values())
+      recordRunnerReservationCapacityFailure(reason, count);
+    const remainingAcceptedByReservation = new Map(reservationValidation.acceptedByReservation);
+    const runnerInstances = params.runnerInstances.flatMap((runner, requestIndex) => {
+      if (!runner.reservationId) return [{...runner, requestIndex}];
+      const remaining = remainingAcceptedByReservation.get(runner.reservationId) ?? 0;
+      if (remaining === 0) return [];
+      remainingAcceptedByReservation.set(runner.reservationId, remaining - 1);
+      return [{...runner, requestIndex}];
+    });
+
+    if (runnerInstances.length === 0)
+      return {
+        runnerInstances: [],
+        reservationUnavailable: reservationValidation.unavailableByReservation.size > 0,
+      };
+
     const instances = await tx
       .insert(providerRunners)
       .values(
-        params.runnerInstances.map((runner) => ({
+        runnerInstances.map((runner) => ({
           provisionerId: params.provisionerId,
           providerKind: params.providerKind ?? null,
           templateKey: runner.templateKey ?? null,
           intendedReservationId: runner.reservationId ?? null,
+          launchKind: runner.reservationId ? ('demand' as const) : ('warm' as const),
           state: 'starting' as const,
           labels: [],
           reportedAt: new Date(),
         })),
       )
       .returning({id: providerRunners.id});
-    const results = instances.map((instance) => ({
-      runnerInstanceId: instance.id,
-      bootstrapToken: generateOpaqueToken('runnerBootstrapToken'),
-    }));
+    const results = instances.map((instance, index) => {
+      const runner = runnerInstances[index];
+      if (!runner) throw new Error('Runner instance insert returned an unexpected row count');
+      return {
+        runnerInstanceId: instance.id,
+        bootstrapToken: generateOpaqueToken('runnerBootstrapToken'),
+        requestIndex: runner.requestIndex,
+      };
+    });
     await tx.insert(runnerBootstrapTokens).values(
       results.map((result) => ({
         runnerInstanceId: result.runnerInstanceId,
@@ -75,7 +124,10 @@ export async function createRunnerInstancesWithBootstrapTokens(params: {
         expiresAt,
       })),
     );
-    return results;
+    return {
+      runnerInstances: results,
+      reservationUnavailable: reservationValidation.unavailableByReservation.size > 0,
+    };
   });
 }
 
@@ -227,6 +279,7 @@ export async function enrollRunnerControlSession(params: {
           runnerInstanceId: params.runnerInstanceId,
           provisionerId: params.provisionerId,
           ttlSeconds: config.RUNNER_ACTIVATION_TOKEN_TTL_SECONDS,
+          surface: 'enrollment',
         });
       });
     } catch (error) {

package/src/core/runner-instances.test.ts

@@ -54,18 +54,37 @@ describe('reconcileRunnerInstancesFromDbResult', () => {
 
     expect(result[0]?.desiredIntent).toBe('terminate');
   });
+
+  it('returns the intended reservation when a runner is not assigned yet', () => {
+    const intendedReservationId = crypto.randomUUID();
+    const result = reconcileRunnerInstancesFromDbResult({
+      observedRunnerInstanceIds: ['provisioned-runner-1'],
+      observedRows: [
+        providerRunner({providerRunnerId: 'provisioned-runner-1', intendedReservationId}),
+      ],
+      boundJobExecutionsByRunnerInstanceId: new Map(),
+    });
+
+    expect(result[0]).toMatchObject({
+      reservationId: null,
+      intendedReservationId,
+    });
+  });
 });
 
 function providerRunner(params: {
   providerRunnerId: string;
   state?: 'starting' | 'running' | 'stopping' | 'stopped' | 'failed' | 'terminated';
+  intendedReservationId?: string | null;
 }) {
   return {
     id: crypto.randomUUID(),
     workspaceId: crypto.randomUUID(),
     provisionerId: crypto.randomUUID(),
     providerRunnerId: params.providerRunnerId,
+    intendedReservationId: params.intendedReservationId ?? null,
     reservationId: null,
+    launchKind: 'manual' as const,
     templateKey: 'linux',
     labels: ['linux'],
     state: params.state ?? 'running',

package/src/core/runner-instances.ts

@@ -62,6 +62,7 @@ export interface ReconciledBoundJobExecution {
 export interface ReconciledRunnerInstance {
   providerRunnerId: string;
   state: RunnerInstanceState | null;
+  intendedReservationId: string | null;
   reservationId: string | null;
   runnerSessionId: string | null;
   boundJobExecution: ReconciledBoundJobExecution | null;
@@ -157,6 +158,7 @@ export function reconcileRunnerInstancesFromDbResult(params: {
     return {
       providerRunnerId,
       state: row?.state ?? null,
+      intendedReservationId: row?.intendedReservationId ?? null,
       reservationId: row?.reservationId ?? null,
       runnerSessionId: row?.runnerSessionId ?? null,
       boundJobExecution: boundJobExecution

package/src/core/runner-sessions.test.ts

@@ -1,9 +1,15 @@
+import {vi} from '@shipfox/vitest/vi';
 import {and, eq} from 'drizzle-orm';
 import {db} from '#db/db.js';
 import {ephemeralRegistrationTokens} from '#db/schema/ephemeral-registration-tokens.js';
 import {runnerActivationTokens} from '#db/schema/runner-activation-tokens.js';
+import type {RunnerInstanceInsertDb} from '#db/schema/runner-instances.js';
 import {providerRunners} from '#db/schema/runner-instances.js';
 import {runnerSessions} from '#db/schema/runner-sessions.js';
+import {
+  type RunnerActivationTokenNotIssuedReason,
+  runnerActivationTokenNotIssuedCount,
+} from '#metrics/instance.js';
 import {
   ephemeralRegistrationTokenFactory,
   getRunnerSessionTokenClaims,
@@ -18,9 +24,34 @@ import {
   RegistrationTokenWorkspaceMismatchError,
   RunnerLabelsReservedError,
 } from './errors.js';
-import {issueRunnerActivationToken} from './runner-activation.js';
+import {getRunnerAssignment, issueRunnerActivationToken} from './runner-activation.js';
 import {registerRunnerSession} from './runner-sessions.js';
 
+const activationTokenNotIssuedCases: Array<{
+  reason: RunnerActivationTokenNotIssuedReason;
+  update: Partial<RunnerInstanceInsertDb>;
+  provisionerId?: string;
+}> = [
+  {reason: 'runner-not-found', update: {}, provisionerId: crypto.randomUUID()},
+  {reason: 'missing-workspace', update: {workspaceId: null}},
+  {reason: 'existing-session', update: {runnerSessionId: crypto.randomUUID()}},
+  {reason: 'not-running', update: {state: 'starting'}},
+];
+
+function activationTokenMetricCalls(spy: {mock: {calls: unknown[][]}}): unknown[][] {
+  // The test setup uses a shared NoopMeter counter, so filter calls to this metric before counting.
+  return spy.mock.calls.filter(([, attributes]) => {
+    if (typeof attributes !== 'object' || attributes === null) return false;
+    const reason = (attributes as {reason?: unknown}).reason;
+    const surface = (attributes as {surface?: unknown}).surface;
+    return (
+      ['runner-not-found', 'missing-workspace', 'existing-session', 'not-running'].includes(
+        String(reason),
+      ) && ['enrollment', 'poll'].includes(String(surface))
+    );
+  });
+}
+
 describe('registerRunnerSession', () => {
   let workspaceId: string;
   let registrationTokenId: string;
@@ -245,11 +276,13 @@ describe('activation runner sessions', () => {
       runnerInstanceId,
       provisionerId,
       ttlSeconds: 60,
+      surface: 'poll',
     });
     const replacementToken = await issueRunnerActivationToken({
       runnerInstanceId,
       provisionerId,
       ttlSeconds: 60,
+      surface: 'poll',
     });
 
     const tokens = await db()
@@ -265,11 +298,97 @@ describe('activation runner sessions', () => {
     expect(tokens.filter((token) => token.revokedAt !== null)).toHaveLength(1);
   });
 
+  it.each(
+    activationTokenNotIssuedCases,
+  )('records $reason when direct activation-token issuance is skipped', async ({
+    reason,
+    update,
+    provisionerId: caseProvisionerId,
+  }) => {
+    const addSpy = vi.spyOn(runnerActivationTokenNotIssuedCount, 'add');
+
+    try {
+      if (Object.keys(update).length > 0) {
+        await db()
+          .update(providerRunners)
+          .set(update)
+          .where(eq(providerRunners.id, runnerInstanceId));
+      }
+
+      const activationToken = await issueRunnerActivationToken({
+        runnerInstanceId,
+        provisionerId: caseProvisionerId ?? provisionerId,
+        ttlSeconds: 60,
+        surface: 'poll',
+      });
+
+      expect(activationToken).toBeNull();
+      expect(activationTokenMetricCalls(addSpy)).toEqual([[1, {reason, surface: 'poll'}]]);
+    } finally {
+      addSpy.mockRestore();
+    }
+  });
+
+  it('does not record a metric when assignment polling finds no assignment', async () => {
+    const addSpy = vi.spyOn(runnerActivationTokenNotIssuedCount, 'add');
+
+    try {
+      await db()
+        .update(providerRunners)
+        .set({workspaceId: null})
+        .where(eq(providerRunners.id, runnerInstanceId));
+
+      const assignment = await getRunnerAssignment({runnerInstanceId, provisionerId});
+
+      expect(assignment).toBeNull();
+      expect(activationTokenMetricCalls(addSpy)).toHaveLength(0);
+    } finally {
+      addSpy.mockRestore();
+    }
+  });
+
+  it('does not poll for an assignment while the runner is not running', async () => {
+    const addSpy = vi.spyOn(runnerActivationTokenNotIssuedCount, 'add');
+
+    try {
+      await db()
+        .update(providerRunners)
+        .set({state: 'starting'})
+        .where(eq(providerRunners.id, runnerInstanceId));
+
+      const assignment = await getRunnerAssignment({runnerInstanceId, provisionerId});
+
+      expect(assignment).toBeNull();
+      expect(activationTokenMetricCalls(addSpy)).toHaveLength(0);
+    } finally {
+      addSpy.mockRestore();
+    }
+  });
+
+  it('does not record a metric when activation-token issuance succeeds', async () => {
+    const addSpy = vi.spyOn(runnerActivationTokenNotIssuedCount, 'add');
+
+    try {
+      const activationToken = await issueRunnerActivationToken({
+        runnerInstanceId,
+        provisionerId,
+        ttlSeconds: 60,
+        surface: 'poll',
+      });
+
+      expect(activationToken).toEqual(expect.any(String));
+      expect(activationTokenMetricCalls(addSpy)).toHaveLength(0);
+    } finally {
+      addSpy.mockRestore();
+    }
+  });
+
   it('allows only one concurrent registration to consume an activation token', async () => {
     const rawToken = await issueRunnerActivationToken({
       runnerInstanceId,
       provisionerId,
       ttlSeconds: 60,
+      surface: 'poll',
     });
     const [activationToken] = await db()
       .select()
@@ -314,6 +433,7 @@ describe('activation runner sessions', () => {
       runnerInstanceId,
       provisionerId,
       ttlSeconds: 60,
+      surface: 'poll',
     });
     const [activationToken] = await db()
       .select()

package/src/db/ephemeral-registration-tokens.test.ts

@@ -3,6 +3,7 @@ import {and, eq, inArray} from 'drizzle-orm';
 import {
   ActiveEphemeralRegistrationTokensExistError,
   type RegistrationTokenBatchExceedsReservationError,
+  ReservationNotFoundError,
 } from '#core/errors.js';
 import {db} from '#db/db.js';
 import {
@@ -56,6 +57,20 @@ describe('createEphemeralRegistrationTokensBatch', () => {
     ).toEqual(expect.objectContaining({providerRunnerId: 'provisioned-runner-a'}));
   });
 
+  it('rejects bound reservations', async () => {
+    const boundReservationId = await createReservation({count: 1, kind: 'bound'});
+
+    await expect(
+      createEphemeralRegistrationTokensBatch({
+        workspaceId,
+        provisionerId,
+        reservationId: boundReservationId,
+        expiresAt: new Date(Date.now() + 300_000),
+        rows: [row('bound-runner', generateOpaqueToken('ephemeralRegistrationToken'))],
+      }),
+    ).rejects.toBeInstanceOf(ReservationNotFoundError);
+  });
+
   it('rejects the batch and inserts no rows when any provisioned runner already has an active token', async () => {
     await ephemeralRegistrationTokenFactory.create({
       workspaceId,
@@ -186,7 +201,10 @@ describe('createEphemeralRegistrationTokensBatch', () => {
     } satisfies Partial<RegistrationTokenBatchExceedsReservationError>);
   });
 
-  async function createReservation(params: {count: number}): Promise<string> {
+  async function createReservation(params: {
+    count: number;
+    kind?: 'bound' | 'launch';
+  }): Promise<string> {
     const [reservation] = await db()
       .insert(reservations)
       .values({
@@ -194,6 +212,7 @@ describe('createEphemeralRegistrationTokensBatch', () => {
         provisionerId,
         requiredLabels: ['linux'],
         count: params.count,
+        kind: params.kind ?? 'launch',
         expiresAt: new Date(Date.now() + 60_000),
       })
       .returning({id: reservations.id});

package/src/db/ephemeral-registration-tokens.ts

@@ -133,6 +133,7 @@ export async function createEphemeralRegistrationTokensBatch(
           eq(reservations.id, params.reservationId),
           eq(reservations.workspaceId, params.workspaceId),
           eq(reservations.provisionerId, params.provisionerId),
+          eq(reservations.kind, 'launch'),
         ),
       )
       .limit(1);

package/src/db/reservation-locks.ts

@@ -0,0 +1,22 @@
+import {sql} from 'drizzle-orm';
+import type {Tx} from './db.js';
+
+export async function lockRunnerReservationAdvisoryKeysTx(
+  tx: Tx,
+  params: {provisionerId: string; reservationIds: readonly string[]},
+): Promise<void> {
+  const reservationIds = [...new Set(params.reservationIds)].sort();
+  if (reservationIds.length === 0) return;
+
+  const lockKeys = sql.join(
+    reservationIds.map(
+      (reservationId) => sql`(${`runners_assignment:${params.provisionerId}:${reservationId}`})`,
+    ),
+    sql`, `,
+  );
+  await tx.execute(sql`
+    select pg_advisory_xact_lock(hashtext(lock_key))
+    from (values ${lockKeys}) as reservation_locks(lock_key)
+    order by lock_key
+  `);
+}