@shipfox/api-runners 2.0.0

package/src/db/reservations.test.ts

@@ -0,0 +1,399 @@
+import {and, eq, inArray, or, sql, sum} from 'drizzle-orm';
+import {db} from '#db/db.js';
+import {
+  deleteReservationsByIds,
+  pollDemandAndReserve,
+  releaseReservationUnits,
+} from '#db/reservations.js';
+import {reservations} from '#db/schema/reservations.js';
+import {pendingJobFactory, reservationFactory} from '#test/index.js';
+
+describe('pollDemandAndReserve', () => {
+  let workspaceId: string;
+  let provisionerId: string;
+
+  beforeEach(() => {
+    workspaceId = crypto.randomUUID();
+    provisionerId = crypto.randomUUID();
+  });
+
+  it('caps grants by unreserved demand, available slots, and max reservations', async () => {
+    await createPendingJobs(50, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 100,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 20)],
+    });
+
+    expect(result.reservations).toHaveLength(1);
+    expect(result.reservations[0]?.count).toBe(20);
+    expect(result.stats[0]).toMatchObject({labels: ['linux'], queued: 50, reserved: 20});
+  });
+
+  it('allocates overlapping label sets most-specific-first', async () => {
+    await createPendingJobs(2, ['linux']);
+    await createPendingJobs(1, ['linux', 'gpu']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 2,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1), template('linux-gpu', ['linux', 'gpu'], 1)],
+    });
+
+    expect(result.reservations).toEqual([
+      expect.objectContaining({labels: ['gpu', 'linux'], count: 1}),
+      expect.objectContaining({labels: ['linux'], count: 1}),
+    ]);
+  });
+
+  it('serializes multiple provisioners so total active reservations do not exceed queued demand', async () => {
+    await createPendingJobs(5, ['linux']);
+
+    await Promise.all([
+      pollDemandAndReserve({
+        workspaceId,
+        provisionerId: crypto.randomUUID(),
+        maxReservations: 5,
+        ttlSeconds: 60,
+        templates: [template('linux-a', ['linux'], 5)],
+      }),
+      pollDemandAndReserve({
+        workspaceId,
+        provisionerId: crypto.randomUUID(),
+        maxReservations: 5,
+        ttlSeconds: 60,
+        templates: [template('linux-b', ['linux'], 5)],
+      }),
+    ]);
+
+    const reserved = await activeReservedCount();
+    expect(reserved).toBeLessThanOrEqual(5);
+  });
+
+  it('does not count expired reservations against demand', async () => {
+    await createPendingJobs(1, ['linux']);
+    await reservationFactory.create({
+      workspaceId,
+      provisionerId,
+      requiredLabels: ['linux'],
+      count: 1,
+      expiresAt: new Date(Date.now() - 60_000),
+    });
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    expect(result.reservations).toHaveLength(1);
+    expect(result.reservations[0]?.count).toBe(1);
+    expect(result.stats[0]).toMatchObject({queued: 1, reserved: 1});
+  });
+
+  it('deducts this provisioner active reservations from advertised capacity', async () => {
+    await createPendingJobs(10, ['linux']);
+    await reservationFactory.create({
+      workspaceId,
+      provisionerId,
+      requiredLabels: ['linux'],
+      count: 5,
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 5,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 5)],
+    });
+
+    const reserved = await activeReservedCount();
+    expect(result.reservations).toEqual([]);
+    expect(result.stats[0]).toMatchObject({queued: 10, reserved: 5});
+    expect(reserved).toBe(5);
+  });
+
+  it('returns multiple reservation groups in one response', async () => {
+    await createPendingJobs(2, ['linux']);
+    await createPendingJobs(1, ['macos']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 3,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 2), template('macos', ['macos'], 1)],
+    });
+
+    expect(result.reservations).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({labels: ['linux'], count: 2}),
+        expect.objectContaining({labels: ['macos'], count: 1}),
+      ]),
+    );
+  });
+
+  it('splits max reservation budget across groups after the most-specific group', async () => {
+    await createPendingJobs(3, ['linux', 'gpu']);
+    await createPendingJobs(3, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 4,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 10), template('linux-gpu', ['linux', 'gpu'], 10)],
+    });
+
+    expect(result.reservations).toEqual([
+      expect.objectContaining({labels: ['gpu', 'linux'], count: 3}),
+      expect.objectContaining({labels: ['linux'], count: 1}),
+    ]);
+  });
+
+  it('excludes demand that no template can satisfy from stats', async () => {
+    await createPendingJobs(1, ['windows']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    expect(result).toEqual({stats: [], reservations: []});
+  });
+
+  it('returns stats without writing rows when max reservations is zero', async () => {
+    await createPendingJobs(1, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 0,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    const reserved = await activeReservedCount();
+    expect(result.stats).toEqual([
+      expect.objectContaining({labels: ['linux'], queued: 1, reserved: 0}),
+    ]);
+    expect(result.reservations).toEqual([]);
+    expect(reserved).toBe(0);
+  });
+
+  it('sets reservation expiry from database time', async () => {
+    await createPendingJobs(1, ['linux']);
+    const before = Date.now();
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    const expiresAt = result.reservations[0]?.expiresAt;
+    expect(expiresAt).toBeInstanceOf(Date);
+    expect(expiresAt?.getTime()).toBeGreaterThan(before + 50_000);
+  });
+
+  it('deletes reservations by id', async () => {
+    await reservationFactory.create({
+      workspaceId,
+      provisionerId,
+      requiredLabels: ['linux'],
+      count: 1,
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+    await reservationFactory.create({
+      workspaceId,
+      provisionerId,
+      requiredLabels: ['macos'],
+      count: 1,
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+    const beforeDelete = await reservationsForTest();
+    const deletedReservation = beforeDelete.find((reservation) =>
+      reservation.requiredLabels.includes('linux'),
+    );
+    if (!deletedReservation) throw new Error('Expected linux reservation');
+
+    const deleted = await deleteReservationsByIds([deletedReservation.id]);
+
+    const remaining = await reservationsForTest();
+    expect(deleted).toBe(1);
+    expect(remaining).toHaveLength(1);
+    expect(remaining[0]?.requiredLabels).toEqual(['macos']);
+  });
+
+  it('decrements reservation units inside a caller transaction', async () => {
+    await reservationFactory.create({
+      workspaceId,
+      provisionerId,
+      requiredLabels: ['linux'],
+      count: 3,
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+    const [reservation] = await reservationsForTest();
+    if (!reservation) throw new Error('Expected reservation');
+
+    const released = await db().transaction((tx) =>
+      releaseReservationUnits(tx, {
+        workspaceId,
+        provisionerId,
+        releases: [{reservationId: reservation.id, count: 2}],
+      }),
+    );
+
+    const rows = await reservationsForTest();
+    expect(released).toBe(2);
+    expect(rows[0]?.count).toBe(1);
+  });
+
+  it('deletes reservations when releasing all remaining units', async () => {
+    await reservationFactory.create({
+      workspaceId,
+      provisionerId,
+      requiredLabels: ['linux'],
+      count: 1,
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+    const [reservation] = await reservationsForTest();
+    if (!reservation) throw new Error('Expected reservation');
+
+    const released = await db().transaction((tx) =>
+      releaseReservationUnits(tx, {
+        workspaceId,
+        provisionerId,
+        releases: [{reservationId: reservation.id, count: 1}],
+      }),
+    );
+
+    const rows = await reservationsForTest();
+    expect(released).toBe(1);
+    expect(rows).toHaveLength(0);
+  });
+
+  it('does not release another workspace or provisioner reservation', async () => {
+    const otherWorkspace = await reservationFactory.create({
+      workspaceId: crypto.randomUUID(),
+      provisionerId,
+      requiredLabels: ['linux'],
+      count: 1,
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+    const otherProvisioner = await reservationFactory.create({
+      workspaceId,
+      provisionerId: crypto.randomUUID(),
+      requiredLabels: ['linux'],
+      count: 1,
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+    const rowsBefore = await db()
+      .select()
+      .from(reservations)
+      .where(
+        or(
+          and(
+            eq(reservations.workspaceId, otherWorkspace.workspaceId),
+            eq(reservations.provisionerId, otherWorkspace.provisionerId),
+          ),
+          and(
+            eq(reservations.workspaceId, otherProvisioner.workspaceId),
+            eq(reservations.provisionerId, otherProvisioner.provisionerId),
+          ),
+        ),
+      );
+
+    const released = await db().transaction((tx) =>
+      releaseReservationUnits(tx, {
+        workspaceId,
+        provisionerId,
+        releases: rowsBefore.map((reservation) => ({reservationId: reservation.id, count: 1})),
+      }),
+    );
+
+    const rowsAfter = await db()
+      .select()
+      .from(reservations)
+      .where(
+        inArray(
+          reservations.id,
+          rowsBefore.map((reservation) => reservation.id),
+        ),
+      );
+    expect(released).toBe(0);
+    expect(rowsAfter.map((reservation) => reservation.id).sort()).toEqual(
+      rowsBefore.map((reservation) => reservation.id).sort(),
+    );
+  });
+
+  it('credits only the deleted reservation count when release units exceed the row count', async () => {
+    await reservationFactory.create({
+      workspaceId,
+      provisionerId,
+      requiredLabels: ['linux'],
+      count: 1,
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+    const [reservation] = await reservationsForTest();
+    if (!reservation) throw new Error('Expected reservation');
+
+    const released = await db().transaction((tx) =>
+      releaseReservationUnits(tx, {
+        workspaceId,
+        provisionerId,
+        releases: [{reservationId: reservation.id, count: 3}],
+      }),
+    );
+
+    const rows = await reservationsForTest();
+    expect(released).toBe(1);
+    expect(rows).toHaveLength(0);
+  });
+
+  async function createPendingJobs(count: number, requiredLabels: string[]): Promise<void> {
+    for (let index = 0; index < count; index++) {
+      await pendingJobFactory.create({workspaceId, requiredLabels});
+    }
+  }
+
+  async function activeReservedCount(): Promise<number> {
+    const [row] = await db()
+      .select({value: sum(reservations.count)})
+      .from(reservations)
+      .where(
+        sql`${reservations.workspaceId} = ${workspaceId} and ${reservations.expiresAt} > now()`,
+      );
+    return Number(row?.value ?? 0);
+  }
+
+  async function reservationsForTest() {
+    return await db()
+      .select()
+      .from(reservations)
+      .where(
+        and(
+          eq(reservations.workspaceId, workspaceId),
+          eq(reservations.provisionerId, provisionerId),
+        ),
+      );
+  }
+
+  function template(templateKey: string, labels: string[], availableSlots: number) {
+    return {templateKey, labels, availableSlots, starting: 0, running: 0};
+  }
+});

package/src/db/reservations.ts

@@ -0,0 +1,296 @@
+import {canonicalizeLabels} from '@shipfox/runner-labels';
+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 {reservations} from './schema/reservations.js';
+
+export interface ReservationTemplate {
+  templateKey: string;
+  labels: string[];
+  availableSlots: number;
+  starting: number;
+  running: number;
+}
+
+export interface DemandStat {
+  labels: string[];
+  queued: number;
+  reserved: number;
+  oldestQueuedAt: Date;
+}
+
+export interface ReservationGrant {
+  reservationId: string;
+  labels: string[];
+  count: number;
+  expiresAt: Date;
+}
+
+export interface PollDemandAndReserveParams {
+  workspaceId: string;
+  provisionerId: string;
+  maxReservations: number;
+  ttlSeconds: number;
+  templates: ReservationTemplate[];
+}
+
+interface NormalizedTemplate {
+  templateKey: string;
+  labels: string[];
+  remainingSlots: number;
+}
+
+interface DemandRow {
+  requiredLabels: string[];
+  queued: number;
+  oldestQueuedAt: Date;
+}
+
+export async function pollDemandAndReserve(
+  params: PollDemandAndReserveParams,
+): Promise<{stats: DemandStat[]; reservations: ReservationGrant[]}> {
+  return await db().transaction(async (tx) => {
+    return await pollDemandAndReserveTx(tx, params);
+  });
+}
+
+export async function pollDemandAndReserveTx(
+  tx: Tx,
+  params: PollDemandAndReserveParams,
+): Promise<{stats: DemandStat[]; reservations: ReservationGrant[]}> {
+  await tx.execute(sql`select pg_advisory_xact_lock(hashtext(${params.workspaceId}))`);
+
+  const demandRows = (
+    await tx
+      .select({
+        requiredLabels: pendingJobExecutions.requiredLabels,
+        queued: sql<number>`count(*)::int`,
+        oldestQueuedAt: sql<Date | string>`min(${pendingJobExecutions.createdAt})`,
+      })
+      .from(pendingJobExecutions)
+      .where(eq(pendingJobExecutions.workspaceId, params.workspaceId))
+      .groupBy(pendingJobExecutions.requiredLabels)
+  ).map((row) => ({
+    ...row,
+    oldestQueuedAt: new Date(row.oldestQueuedAt),
+  }));
+
+  const activeReservationRows = await tx
+    .select({
+      requiredLabels: reservations.requiredLabels,
+      reserved: sql<number>`coalesce(sum(${reservations.count}), 0)::int`,
+    })
+    .from(reservations)
+    .where(
+      and(eq(reservations.workspaceId, params.workspaceId), gt(reservations.expiresAt, sql`now()`)),
+    )
+    .groupBy(reservations.requiredLabels);
+
+  const reservedByLabels = new Map(
+    activeReservationRows.map((row) => [labelKey(row.requiredLabels), row.reserved]),
+  );
+  const activeProvisionerReservationRows = await tx
+    .select({
+      requiredLabels: reservations.requiredLabels,
+      reserved: sql<number>`coalesce(sum(${reservations.count}), 0)::int`,
+    })
+    .from(reservations)
+    .where(
+      and(
+        eq(reservations.workspaceId, params.workspaceId),
+        eq(reservations.provisionerId, params.provisionerId),
+        gt(reservations.expiresAt, sql`now()`),
+      ),
+    )
+    .groupBy(reservations.requiredLabels);
+
+  const templates = params.templates.map((template) => ({
+    templateKey: template.templateKey,
+    labels: [...canonicalizeLabels(template.labels)],
+    remainingSlots: template.availableSlots,
+  }));
+  deductProvisionerReservations(templates, activeProvisionerReservationRows);
+  const stats: DemandStat[] = [];
+  const grants: ReservationGrant[] = [];
+  let remainingMaxReservations = params.maxReservations;
+
+  for (const demand of sortDemandRows(demandRows)) {
+    const satisfyingTemplates = templates
+      .filter((template) => isSubset(demand.requiredLabels, template.labels))
+      .sort(
+        (a, b) => a.labels.length - b.labels.length || a.templateKey.localeCompare(b.templateKey),
+      );
+
+    if (satisfyingTemplates.length === 0) continue;
+
+    const reserved = reservedByLabels.get(labelKey(demand.requiredLabels)) ?? 0;
+    const unreserved = Math.max(0, demand.queued - reserved);
+    const capacity = satisfyingTemplates.reduce(
+      (total, template) => total + template.remainingSlots,
+      0,
+    );
+    const grant = Math.min(unreserved, capacity, remainingMaxReservations);
+    let reservedAfterGrant = reserved;
+
+    if (grant > 0 && params.maxReservations > 0) {
+      const [inserted] = await tx
+        .insert(reservations)
+        .values({
+          workspaceId: params.workspaceId,
+          provisionerId: params.provisionerId,
+          requiredLabels: demand.requiredLabels,
+          count: grant,
+          expiresAt: sql`now() + (${params.ttlSeconds} || ' seconds')::interval`,
+        })
+        .returning({id: reservations.id, expiresAt: reservations.expiresAt});
+
+      if (!inserted) throw new Error('Insert returned no rows');
+
+      remainingMaxReservations -= grant;
+      drawSlots(satisfyingTemplates, grant);
+      reservedAfterGrant += grant;
+      grants.push({
+        reservationId: inserted.id,
+        labels: demand.requiredLabels,
+        count: grant,
+        expiresAt: inserted.expiresAt,
+      });
+    }
+
+    stats.push({
+      labels: demand.requiredLabels,
+      queued: demand.queued,
+      reserved: reservedAfterGrant,
+      oldestQueuedAt: demand.oldestQueuedAt,
+    });
+  }
+
+  return {stats, reservations: grants};
+}
+
+export async function deleteExpiredReservations(params?: {limit?: number}): Promise<number> {
+  const expiredIds = db()
+    .select({id: reservations.id})
+    .from(reservations)
+    .where(lt(reservations.expiresAt, sql`now()`))
+    .orderBy(asc(reservations.expiresAt))
+    .limit(params?.limit ?? 1000);
+
+  const deleted = await db()
+    .delete(reservations)
+    .where(inArray(reservations.id, expiredIds))
+    .returning({id: reservations.id});
+
+  return deleted.length;
+}
+
+export async function deleteReservationsByIds(ids: string[]): Promise<number> {
+  if (ids.length === 0) return 0;
+
+  const deleted = await db()
+    .delete(reservations)
+    .where(inArray(reservations.id, ids))
+    .returning({id: reservations.id});
+
+  return deleted.length;
+}
+
+export async function releaseReservationUnits(
+  tx: Tx,
+  params: {
+    workspaceId: string;
+    provisionerId: string;
+    releases: Array<{reservationId: string; count: number}>;
+  },
+): Promise<number> {
+  let released = 0;
+
+  for (const release of params.releases) {
+    if (release.count <= 0) continue;
+
+    const decremented = await tx
+      .update(reservations)
+      .set({count: sql`${reservations.count} - ${release.count}`})
+      .where(
+        and(
+          eq(reservations.id, release.reservationId),
+          eq(reservations.workspaceId, params.workspaceId),
+          eq(reservations.provisionerId, params.provisionerId),
+          gt(reservations.count, release.count),
+          gt(reservations.expiresAt, sql`now()`),
+        ),
+      )
+      .returning({id: reservations.id});
+
+    if (decremented.length > 0) {
+      released += release.count;
+      continue;
+    }
+
+    const deleted = await tx
+      .delete(reservations)
+      .where(
+        and(
+          eq(reservations.id, release.reservationId),
+          eq(reservations.workspaceId, params.workspaceId),
+          eq(reservations.provisionerId, params.provisionerId),
+          sql`${reservations.count} <= ${release.count}`,
+          gt(reservations.expiresAt, sql`now()`),
+        ),
+      )
+      .returning({count: reservations.count});
+
+    released += deleted.reduce((total, row) => total + row.count, 0);
+  }
+
+  return released;
+}
+
+function deductProvisionerReservations(
+  templates: NormalizedTemplate[],
+  activeReservations: {requiredLabels: string[]; reserved: number}[],
+): void {
+  for (const reservation of sortReservationRows(activeReservations)) {
+    const satisfyingTemplates = templates
+      .filter((template) => isSubset(reservation.requiredLabels, template.labels))
+      .sort(
+        (a, b) => a.labels.length - b.labels.length || a.templateKey.localeCompare(b.templateKey),
+      );
+    drawSlots(satisfyingTemplates, reservation.reserved);
+  }
+}
+
+function sortDemandRows(rows: DemandRow[]): DemandRow[] {
+  return [...rows].sort((a, b) => {
+    const specificity = b.requiredLabels.length - a.requiredLabels.length;
+    if (specificity !== 0) return specificity;
+    return a.oldestQueuedAt.getTime() - b.oldestQueuedAt.getTime();
+  });
+}
+
+function sortReservationRows<T extends {requiredLabels: string[]}>(rows: T[]): T[] {
+  return [...rows].sort((a, b) => {
+    const specificity = b.requiredLabels.length - a.requiredLabels.length;
+    if (specificity !== 0) return specificity;
+    return labelKey(a.requiredLabels).localeCompare(labelKey(b.requiredLabels));
+  });
+}
+
+function isSubset(requiredLabels: string[], availableLabels: string[]): boolean {
+  return requiredLabels.every((label) => availableLabels.includes(label));
+}
+
+function drawSlots(templates: NormalizedTemplate[], count: number): void {
+  let remaining = count;
+  for (const template of templates) {
+    if (remaining === 0) return;
+    const used = Math.min(template.remainingSlots, remaining);
+    template.remainingSlots -= used;
+    remaining -= used;
+  }
+}
+
+function labelKey(labels: string[]): string {
+  return JSON.stringify(labels);
+}