@shipfox/api-runners 5.0.0 → 7.0.1

package/src/db/reservations.test.ts

@@ -3,6 +3,7 @@ import {db} from '#db/db.js';
 import {
   deleteReservationsByIds,
   pollDemandAndReserve,
+  pollInstallationDemandAndReserve,
   releaseReservationUnits,
 } from '#db/reservations.js';
 import {reservations} from '#db/schema/reservations.js';
@@ -75,6 +76,54 @@ describe('pollDemandAndReserve', () => {
     expect(reserved).toBeLessThanOrEqual(5);
   });
 
+  it('allocates eligible workspace heads in oldest-demand order without overselling templates', async () => {
+    const olderWorkspaceId = crypto.randomUUID();
+    const newerWorkspaceId = crypto.randomUUID();
+    await pendingJobFactory.create({workspaceId: olderWorkspaceId, requiredLabels: ['linux']});
+    await pendingJobFactory.create({workspaceId: newerWorkspaceId, requiredLabels: ['linux']});
+
+    const result = await pollInstallationDemandAndReserve({
+      provisionerId,
+      maxReservations: 5,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+      capabilityWindowSeconds: 60,
+      eligibleWorkspaceIds: new Set([olderWorkspaceId, newerWorkspaceId]),
+    });
+
+    expect(result.reservations).toEqual([
+      expect.objectContaining({workspaceId: olderWorkspaceId, labels: ['linux'], count: 1}),
+    ]);
+  });
+
+  it('reports committed installation grants before an aborted poll stops allocating', async () => {
+    const firstWorkspaceId = crypto.randomUUID();
+    const secondWorkspaceId = crypto.randomUUID();
+    const abortController = new AbortController();
+    const reportedReservations: string[] = [];
+    await pendingJobFactory.create({workspaceId: firstWorkspaceId, requiredLabels: ['linux']});
+    await pendingJobFactory.create({workspaceId: secondWorkspaceId, requiredLabels: ['linux']});
+
+    const result = await pollInstallationDemandAndReserve({
+      provisionerId,
+      maxReservations: 2,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 2)],
+      capabilityWindowSeconds: 60,
+      eligibleWorkspaceIds: new Set([firstWorkspaceId, secondWorkspaceId]),
+      signal: abortController.signal,
+      onReservations: (reservations) => {
+        reportedReservations.push(...reservations.map((reservation) => reservation.reservationId));
+        abortController.abort();
+      },
+    });
+
+    expect(result.reservations).toHaveLength(1);
+    expect(reportedReservations).toEqual(
+      result.reservations.map((reservation) => reservation.reservationId),
+    );
+  });
+
   it('does not count expired reservations against demand', async () => {
     await createPendingJobs(1, ['linux']);
     await reservationFactory.create({

package/src/db/reservations.ts

@@ -3,6 +3,8 @@ import {and, asc, eq, gt, inArray, lt, sql} from 'drizzle-orm';
 import type {Tx} from './db.js';
 import {db} from './db.js';
 import {pendingJobExecutions} from './schema/pending-job-executions.js';
+import {provisionerCapabilitySnapshots} from './schema/provisioner-capability-snapshots.js';
+import {provisionerTokens} from './schema/provisioner-tokens.js';
 import {reservations} from './schema/reservations.js';
 
 export interface ReservationTemplate {
@@ -14,6 +16,7 @@ export interface ReservationTemplate {
 }
 
 export interface DemandStat {
+  workspaceId?: string;
   labels: string[];
   queued: number;
   reserved: number;
@@ -22,6 +25,7 @@ export interface DemandStat {
 
 export interface ReservationGrant {
   reservationId: string;
+  workspaceId?: string;
   labels: string[];
   count: number;
   expiresAt: Date;
@@ -33,6 +37,18 @@ export interface PollDemandAndReserveParams {
   maxReservations: number;
   ttlSeconds: number;
   templates: ReservationTemplate[];
+  capabilityWindowSeconds?: number;
+}
+
+export interface InstallationPollDemandAndReserveParams {
+  provisionerId: string;
+  maxReservations: number;
+  ttlSeconds: number;
+  templates: ReservationTemplate[];
+  capabilityWindowSeconds: number;
+  eligibleWorkspaceIds: ReadonlySet<string>;
+  signal?: AbortSignal;
+  onReservations?: (reservations: ReservationGrant[]) => void;
 }
 
 interface NormalizedTemplate {
@@ -60,8 +76,75 @@ 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);
+}
+
+export async function pollInstallationDemandAndReserve(
+  params: InstallationPollDemandAndReserveParams,
+): Promise<{stats: DemandStat[]; reservations: ReservationGrant[]}> {
+  const candidateWorkspaceIds = await listInstallationDemandWorkspaceIds(
+    params.eligibleWorkspaceIds,
+  );
+  const results: Array<{stats: DemandStat[]; reservations: ReservationGrant[]}> = [];
+  let remainingMaxReservations = params.maxReservations;
+  const remainingTemplates = params.templates.map((template) => ({
+    ...template,
+    labels: [...canonicalizeLabels(template.labels)],
+    remainingSlots: template.availableSlots,
+  }));
+  for (const workspaceId of candidateWorkspaceIds) {
+    if (params.signal?.aborted || remainingMaxReservations === 0) break;
+    const result = await db().transaction(async (tx) => {
+      const lockResult = await tx.execute<{locked: boolean}>(
+        sql`select pg_try_advisory_xact_lock(hashtext(${workspaceId})) as locked`,
+      );
+      const locked = lockResult.rows[0];
+      if (!locked?.locked) return {stats: [], reservations: []};
+      return await pollDemandAndReserveLockedTx(tx, {
+        workspaceId,
+        provisionerId: params.provisionerId,
+        maxReservations: remainingMaxReservations,
+        ttlSeconds: params.ttlSeconds,
+        templates: remainingTemplates.map((template) => ({
+          ...template,
+          availableSlots: template.remainingSlots,
+        })),
+        capabilityWindowSeconds: params.capabilityWindowSeconds,
+      });
+    });
+    results.push(result);
+    consumeInstallationTemplateSlots(remainingTemplates, result.reservations);
+    params.onReservations?.(result.reservations);
+    remainingMaxReservations -= result.reservations.reduce(
+      (total, reservation) => total + reservation.count,
+      0,
+    );
+  }
+  return {
+    stats: results.flatMap((result) => result.stats),
+    reservations: results.flatMap((result) => result.reservations),
+  };
+}
 
-  const demandRows = (
+function consumeInstallationTemplateSlots(
+  templates: NormalizedTemplate[],
+  reservations: ReservationGrant[],
+): void {
+  for (const reservation of reservations) {
+    const satisfyingTemplates = templates
+      .filter((template) => isSubset(reservation.labels, template.labels))
+      .sort(
+        (a, b) => a.labels.length - b.labels.length || a.templateKey.localeCompare(b.templateKey),
+      );
+    drawSlots(satisfyingTemplates, reservation.count);
+  }
+}
+
+async function pollDemandAndReserveLockedTx(
+  tx: Tx,
+  params: PollDemandAndReserveParams,
+): Promise<{stats: DemandStat[]; reservations: ReservationGrant[]}> {
+  let demandRows = (
     await tx
       .select({
         requiredLabels: pendingJobExecutions.requiredLabels,
@@ -75,6 +158,15 @@ export async function pollDemandAndReserveTx(
     ...row,
     oldestQueuedAt: new Date(row.oldestQueuedAt),
   }));
+  if (params.capabilityWindowSeconds !== undefined) {
+    const capabilityLabels = await listActiveWorkspaceCapabilityLabelsTx(tx, {
+      workspaceId: params.workspaceId,
+      windowSeconds: params.capabilityWindowSeconds,
+    });
+    demandRows = demandRows.filter(
+      (demand) => !capabilityLabels.some((labels) => isSubset(demand.requiredLabels, labels)),
+    );
+  }
 
   const activeReservationRows = await tx
     .select({
@@ -114,6 +206,8 @@ export async function pollDemandAndReserveTx(
   const stats: DemandStat[] = [];
   const grants: ReservationGrant[] = [];
   let remainingMaxReservations = params.maxReservations;
+  const workspaceIdField =
+    params.capabilityWindowSeconds === undefined ? {} : {workspaceId: params.workspaceId};
 
   for (const demand of sortDemandRows(demandRows)) {
     const satisfyingTemplates = templates
@@ -152,6 +246,7 @@ export async function pollDemandAndReserveTx(
       reservedAfterGrant += grant;
       grants.push({
         reservationId: inserted.id,
+        ...workspaceIdField,
         labels: demand.requiredLabels,
         count: grant,
         expiresAt: inserted.expiresAt,
@@ -159,6 +254,7 @@ export async function pollDemandAndReserveTx(
     }
 
     stats.push({
+      ...workspaceIdField,
       labels: demand.requiredLabels,
       queued: demand.queued,
       reserved: reservedAfterGrant,
@@ -169,6 +265,55 @@ export async function pollDemandAndReserveTx(
   return {stats, reservations: grants};
 }
 
+async function listInstallationDemandWorkspaceIds(eligibleWorkspaceIds: ReadonlySet<string>) {
+  if (eligibleWorkspaceIds.size === 0) return [];
+  const rows = await db()
+    .select({
+      workspaceId: pendingJobExecutions.workspaceId,
+      oldestQueuedAt: sql<Date>`min(${pendingJobExecutions.createdAt})`,
+    })
+    .from(pendingJobExecutions)
+    .where(inArray(pendingJobExecutions.workspaceId, [...eligibleWorkspaceIds]))
+    .groupBy(pendingJobExecutions.workspaceId)
+    .orderBy(
+      asc(sql`min(${pendingJobExecutions.createdAt})`),
+      asc(pendingJobExecutions.workspaceId),
+    );
+  return rows.map((row) => row.workspaceId);
+}
+
+export async function listQueuedDemandWorkspaceIds(): Promise<string[]> {
+  const rows = await db()
+    .select({workspaceId: pendingJobExecutions.workspaceId})
+    .from(pendingJobExecutions)
+    .groupBy(pendingJobExecutions.workspaceId);
+  return rows.map((row) => row.workspaceId);
+}
+
+async function listActiveWorkspaceCapabilityLabelsTx(
+  tx: Tx,
+  params: {workspaceId: string; windowSeconds: number},
+): Promise<string[][]> {
+  const rows = await tx
+    .select({labels: provisionerCapabilitySnapshots.labels})
+    .from(provisionerCapabilitySnapshots)
+    .innerJoin(
+      provisionerTokens,
+      eq(provisionerTokens.id, provisionerCapabilitySnapshots.provisionerId),
+    )
+    .where(
+      and(
+        eq(provisionerCapabilitySnapshots.workspaceId, params.workspaceId),
+        eq(provisionerTokens.workspaceId, params.workspaceId),
+        eq(provisionerTokens.scope, 'workspace'),
+        sql`${provisionerCapabilitySnapshots.advertisedAt} > now() - (${params.windowSeconds} || ' seconds')::interval`,
+        sql`${provisionerTokens.revokedAt} is null`,
+        sql`(${provisionerTokens.expiresAt} is null or ${provisionerTokens.expiresAt} > now())`,
+      ),
+    );
+  return rows.map((row) => row.labels);
+}
+
 export async function deleteExpiredReservations(params?: {limit?: number}): Promise<number> {
   const expiredIds = db()
     .select({id: reservations.id})

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

@@ -0,0 +1,163 @@
+import {eq} from 'drizzle-orm';
+import {
+  ReservationExpiredError,
+  RunnerInstanceAlreadyAssignedError,
+  RunnerInstanceNotAssignableError,
+} from '#core/errors.js';
+import {db} from '#db/db.js';
+import {assignRunnerInstances} from '#db/runner-assignments.js';
+import {reservations} from '#db/schema/reservations.js';
+import {runnerControlSessions} from '#db/schema/runner-control-sessions.js';
+import {providerRunners} from '#db/schema/runner-instances.js';
+
+describe('assignRunnerInstances', () => {
+  let workspaceId: string;
+  let provisionerId: string;
+
+  beforeEach(() => {
+    workspaceId = crypto.randomUUID();
+    provisionerId = crypto.randomUUID();
+  });
+
+  it('writes the reservation and workspace after enrollment', async () => {
+    const reservation = await createReservation();
+    const runner = await createEnrolledRunner();
+
+    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,
+      assignedAt: expect.any(Date),
+    });
+  });
+
+  it('is idempotent for concurrent retries of the same assignment', async () => {
+    const reservation = await createReservation();
+    const runner = await createEnrolledRunner();
+
+    const results = await Promise.all([
+      assignRunnerInstances({
+        provisionerId,
+        reservationId: reservation.id,
+        runnerInstanceIds: [runner.id],
+      }),
+      assignRunnerInstances({
+        provisionerId,
+        reservationId: reservation.id,
+        runnerInstanceIds: [runner.id],
+      }),
+    ]);
+
+    expect(results).toEqual([[runner.id], [runner.id]]);
+  });
+
+  it('rejects expired reservations', async () => {
+    const reservation = await createReservation({expiresAt: new Date(Date.now() - 1_000)});
+    const runner = await createEnrolledRunner();
+
+    const assignment = assignRunnerInstances({
+      provisionerId,
+      reservationId: reservation.id,
+      runnerInstanceIds: [runner.id],
+    });
+
+    await expect(assignment).rejects.toThrow(ReservationExpiredError);
+  });
+
+  it('rejects unenrolled or incompatible runners', async () => {
+    const reservation = await createReservation({requiredLabels: ['linux', 'gpu']});
+    const runner = await createEnrolledRunner({labels: ['linux']});
+
+    const assignment = assignRunnerInstances({
+      provisionerId,
+      reservationId: reservation.id,
+      runnerInstanceIds: [runner.id],
+    });
+
+    await expect(assignment).rejects.toThrow(RunnerInstanceNotAssignableError);
+  });
+
+  it('rejects a runner assigned to a different reservation', async () => {
+    const firstReservation = await createReservation();
+    const secondReservation = await createReservation();
+    const runner = await createEnrolledRunner();
+    await assignRunnerInstances({
+      provisionerId,
+      reservationId: firstReservation.id,
+      runnerInstanceIds: [runner.id],
+    });
+
+    const assignment = assignRunnerInstances({
+      provisionerId,
+      reservationId: secondReservation.id,
+      runnerInstanceIds: [runner.id],
+    });
+
+    await expect(assignment).rejects.toThrow(RunnerInstanceAlreadyAssignedError);
+  });
+
+  it('rejects assignments that exceed reservation capacity', async () => {
+    const reservation = await createReservation();
+    const firstRunner = await createEnrolledRunner();
+    const secondRunner = await createEnrolledRunner();
+
+    const assignment = assignRunnerInstances({
+      provisionerId,
+      reservationId: reservation.id,
+      runnerInstanceIds: [firstRunner.id, secondRunner.id],
+    });
+
+    await expect(assignment).rejects.toThrow(RunnerInstanceNotAssignableError);
+  });
+
+  async function createReservation(
+    overrides: Partial<{expiresAt: Date; requiredLabels: string[]}> = {},
+  ) {
+    const [reservation] = await db()
+      .insert(reservations)
+      .values({
+        workspaceId,
+        provisionerId,
+        requiredLabels: overrides.requiredLabels ?? ['linux'],
+        count: 1,
+        expiresAt: overrides.expiresAt ?? new Date(Date.now() + 60_000),
+      })
+      .returning();
+    if (!reservation) throw new Error('Reservation insert returned no row');
+    return reservation;
+  }
+
+  async function createEnrolledRunner(overrides: Partial<{labels: string[]}> = {}) {
+    const [runner] = await db()
+      .insert(providerRunners)
+      .values({
+        provisionerId,
+        providerRunnerId: crypto.randomUUID(),
+        labels: overrides.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),
+      });
+    return runner;
+  }
+});

package/src/db/runner-assignments.ts

@@ -0,0 +1,115 @@
+import {and, eq, inArray, isNull, sql} from 'drizzle-orm';
+import {
+  ReservationExpiredError,
+  ReservationNotFoundError,
+  RunnerInstanceAlreadyAssignedError,
+  RunnerInstanceNotAssignableError,
+} from '#core/errors.js';
+import {db} from './db.js';
+import {reservations} from './schema/reservations.js';
+import {runnerControlSessions} from './schema/runner-control-sessions.js';
+import {providerRunners} from './schema/runner-instances.js';
+
+/** Atomically consumes reservation units by writing the immutable assignment on each runner. */
+export async function assignRunnerInstances(params: {
+  provisionerId: string;
+  reservationId: string;
+  runnerInstanceIds: string[];
+}): Promise<string[]> {
+  const runnerInstanceIds = [...params.runnerInstanceIds].sort();
+  return await db().transaction(async (tx) => {
+    await tx.execute(
+      sql`select pg_advisory_xact_lock(hashtext(${`runners_assignment:${params.provisionerId}:${params.reservationId}`}))`,
+    );
+    const [reservation] = await tx
+      .select()
+      .from(reservations)
+      .where(
+        and(
+          eq(reservations.id, params.reservationId),
+          eq(reservations.provisionerId, params.provisionerId),
+        ),
+      )
+      .limit(1)
+      .for('update');
+    if (!reservation) throw new ReservationNotFoundError(params.reservationId);
+    if (reservation.expiresAt <= new Date())
+      throw new ReservationExpiredError(params.reservationId);
+
+    const runnerRows = await tx
+      .select({
+        id: providerRunners.id,
+        reservationId: providerRunners.reservationId,
+        workspaceId: providerRunners.workspaceId,
+        providerRunnerId: providerRunners.providerRunnerId,
+        labels: providerRunners.labels,
+        state: providerRunners.state,
+      })
+      .from(providerRunners)
+      .where(
+        and(
+          eq(providerRunners.provisionerId, params.provisionerId),
+          inArray(providerRunners.id, runnerInstanceIds),
+        ),
+      )
+      .for('update');
+    const activeControlSessions = await tx
+      .select({runnerInstanceId: runnerControlSessions.runnerInstanceId})
+      .from(runnerControlSessions)
+      .where(
+        and(
+          inArray(
+            runnerControlSessions.runnerInstanceId,
+            runnerRows.map((runner) => runner.id),
+          ),
+          isNull(runnerControlSessions.closedAt),
+        ),
+      );
+    const runnerInstanceIdsWithControlSession = new Set(
+      activeControlSessions.map((session) => session.runnerInstanceId),
+    );
+    const runners = runnerRows.map((runner) => ({
+      ...runner,
+      controlSessionId: runnerInstanceIdsWithControlSession.has(runner.id) ? runner.id : null,
+    }));
+    if (runners.length !== runnerInstanceIds.length)
+      throw new RunnerInstanceNotAssignableError(runnerInstanceIds[0] ?? '');
+
+    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 assignedCount = await tx
+      .select({count: sql<number>`count(*)::int`})
+      .from(providerRunners)
+      .where(eq(providerRunners.reservationId, reservation.id));
+    if ((assignedCount[0]?.count ?? 0) + newRunners.length > reservation.count)
+      throw new RunnerInstanceNotAssignableError(newRunners[0]?.id ?? '');
+    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 (newRunners.length > 0) {
+      await tx
+        .update(providerRunners)
+        .set({
+          workspaceId: reservation.workspaceId,
+          reservationId: reservation.id,
+          assignedAt: sql`now()`,
+          updatedAt: sql`now()`,
+        })
+        .where(
+          inArray(
+            providerRunners.id,
+            newRunners.map((runner) => runner.id),
+          ),
+        );
+    }
+    return params.runnerInstanceIds;
+  });
+}