@shipfox/api-runners 12.2.0 → 12.3.0

package/src/db/reservations.test.ts

@@ -86,6 +86,350 @@ describe('pollDemandAndReserve', () => {
     });
   });
 
+  it('binds a runner whose intended reservation has passed its activation grace period', async () => {
+    const intendedReservation = await createIntendedReservation({
+      workspaceId,
+      expiresAt: new Date(Date.now() - 60_000),
+    });
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      intendedReservationId: intendedReservation.id,
+    });
+    await createPendingJobs(1, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(storedRunner).toMatchObject({
+      workspaceId,
+      reservationId: result.reservations[0]?.reservationId,
+      intendedReservationId: null,
+      assignedAt: expect.any(Date),
+    });
+  });
+
+  it('binds a runner whose intended reservation was deleted after its grace period', async () => {
+    const intendedReservation = await createIntendedReservation({
+      workspaceId,
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+    await db().delete(reservations).where(eq(reservations.id, intendedReservation.id));
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      intendedReservationId: intendedReservation.id,
+    });
+    await createPendingJobs(1, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(storedRunner).toMatchObject({
+      workspaceId,
+      reservationId: result.reservations[0]?.reservationId,
+      intendedReservationId: null,
+      assignedAt: expect.any(Date),
+    });
+  });
+
+  it('binds an assigned runner after its reservation grace period expires', async () => {
+    const staleReservation = await createIntendedReservation({
+      workspaceId,
+      expiresAt: new Date(Date.now() - 60_000),
+    });
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      workspaceId,
+      reservationId: staleReservation.id,
+      intendedReservationId: staleReservation.id,
+    });
+    const [activationToken] = await db()
+      .insert(runnerActivationTokens)
+      .values({
+        runnerInstanceId: runner.id,
+        hashedToken: crypto.randomUUID(),
+        prefix: 'test',
+        expiresAt: new Date(Date.now() + 60_000),
+      })
+      .returning();
+    if (!activationToken) throw new Error('Expected activation token');
+    await createPendingJobs(1, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    const [storedToken] = await db()
+      .select()
+      .from(runnerActivationTokens)
+      .where(eq(runnerActivationTokens.id, activationToken.id));
+    expect(storedRunner).toMatchObject({
+      workspaceId,
+      reservationId: result.reservations[0]?.reservationId,
+      intendedReservationId: null,
+      assignedAt: expect.any(Date),
+    });
+    expect(storedToken?.revokedAt).toBeInstanceOf(Date);
+  });
+
+  it('binds an assigned runner whose reservation was deleted', async () => {
+    const staleReservation = await createIntendedReservation({
+      workspaceId,
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+    await db().delete(reservations).where(eq(reservations.id, staleReservation.id));
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      workspaceId,
+      reservationId: staleReservation.id,
+      intendedReservationId: staleReservation.id,
+    });
+    await createPendingJobs(1, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(storedRunner).toMatchObject({
+      workspaceId,
+      reservationId: result.reservations[0]?.reservationId,
+      intendedReservationId: null,
+      assignedAt: expect.any(Date),
+    });
+  });
+
+  it('does not bind an assigned runner within its activation grace period', async () => {
+    const liveReservation = await createIntendedReservation({
+      workspaceId,
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      workspaceId,
+      reservationId: liveReservation.id,
+    });
+    await createPendingJobs(2, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 2,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 2)],
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(result.reservations[0]?.count).toBe(1);
+    expect(storedRunner).toMatchObject({
+      workspaceId,
+      reservationId: liveReservation.id,
+      intendedReservationId: null,
+      assignedAt: null,
+    });
+  });
+
+  it('does not bind a runner whose intended reservation is within its activation grace period', async () => {
+    const intendedReservation = await createIntendedReservation({
+      workspaceId: crypto.randomUUID(),
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      intendedReservationId: intendedReservation.id,
+    });
+    await createPendingJobs(1, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(result.reservations[0]?.count).toBe(1);
+    expect(storedRunner).toMatchObject({
+      workspaceId: null,
+      reservationId: null,
+      intendedReservationId: intendedReservation.id,
+      assignedAt: null,
+    });
+  });
+
+  it('does not rebind a stale runner from another workspace in a workspace-scoped poll', async () => {
+    const previousWorkspaceId = crypto.randomUUID();
+    const staleReservation = await createIntendedReservation({
+      workspaceId: previousWorkspaceId,
+      expiresAt: new Date(Date.now() - 60_000),
+    });
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      workspaceId: previousWorkspaceId,
+      reservationId: staleReservation.id,
+      intendedReservationId: staleReservation.id,
+    });
+    await createPendingJobs(1, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(result.reservations[0]?.count).toBe(1);
+    expect(storedRunner).toMatchObject({
+      workspaceId: previousWorkspaceId,
+      reservationId: staleReservation.id,
+      intendedReservationId: staleReservation.id,
+    });
+  });
+
+  it('rebinds an unowned runner whose reservation was deleted', async () => {
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      reservationId: crypto.randomUUID(),
+    });
+    await createPendingJobs(1, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(storedRunner).toMatchObject({
+      workspaceId,
+      reservationId: result.reservations[0]?.reservationId,
+      intendedReservationId: null,
+      assignedAt: expect.any(Date),
+    });
+  });
+
+  it('rebinds a stale runner across workspaces for an installation-scoped poll', async () => {
+    const previousWorkspaceId = crypto.randomUUID();
+    const targetWorkspaceId = crypto.randomUUID();
+    const staleReservation = await createIntendedReservation({
+      workspaceId: previousWorkspaceId,
+      expiresAt: new Date(Date.now() - 60_000),
+    });
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      workspaceId: previousWorkspaceId,
+      reservationId: staleReservation.id,
+      intendedReservationId: staleReservation.id,
+    });
+    await pendingJobFactory.create({workspaceId: targetWorkspaceId, requiredLabels: ['linux']});
+
+    const result = await pollInstallationDemandAndReserve({
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+      capabilityWindowSeconds: 60,
+      eligibleWorkspaceIds: new Set([targetWorkspaceId]),
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(result.reservations[0]).toMatchObject({workspaceId: targetWorkspaceId, count: 1});
+    expect(storedRunner).toMatchObject({
+      workspaceId: targetWorkspaceId,
+      reservationId: result.reservations[0]?.reservationId,
+      intendedReservationId: null,
+      assignedAt: expect.any(Date),
+    });
+  });
+
+  it('does not rebind a released runner after its reservation expires', async () => {
+    const staleReservation = await createIntendedReservation({
+      workspaceId,
+      expiresAt: new Date(Date.now() - 60_000),
+    });
+    const releasedAt = new Date(Date.now() - 30_000);
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      workspaceId,
+      reservationId: staleReservation.id,
+      intendedReservationId: staleReservation.id,
+      reservationReleasedAt: releasedAt,
+    });
+    await createPendingJobs(1, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(result.reservations[0]?.count).toBe(1);
+    expect(storedRunner).toMatchObject({
+      workspaceId,
+      reservationId: staleReservation.id,
+      intendedReservationId: staleReservation.id,
+      reservationReleasedAt: releasedAt,
+    });
+  });
+
   it('does not bind a running runner without an active control session', async () => {
     const [runner] = await db()
       .insert(providerRunners)
@@ -808,19 +1152,43 @@ describe('pollDemandAndReserve', () => {
     }
   }
 
+  async function createIntendedReservation(params: {
+    workspaceId: string;
+    expiresAt: Date;
+  }): Promise<{id: string}> {
+    const [reservation] = await db()
+      .insert(reservations)
+      .values({
+        workspaceId: params.workspaceId,
+        provisionerId,
+        requiredLabels: ['linux'],
+        count: 1,
+        expiresAt: params.expiresAt,
+      })
+      .returning({id: reservations.id});
+    if (!reservation) throw new Error('Expected reservation');
+    return reservation;
+  }
+
   async function createIdleRunner(params: {
     labels: string[];
     createdAt?: Date;
     controlSessionExpiresAt?: Date;
-    intendedReservationId?: string;
+    workspaceId?: string | null;
+    reservationId?: string | null;
+    intendedReservationId?: string | null;
+    reservationReleasedAt?: Date | null;
   }) {
     const createdAt = params.createdAt ?? new Date();
     const [runner] = await db()
       .insert(providerRunners)
       .values({
         provisionerId,
+        workspaceId: params.workspaceId,
+        reservationId: params.reservationId,
         providerRunnerId: crypto.randomUUID(),
         intendedReservationId: params.intendedReservationId,
+        reservationReleasedAt: params.reservationReleasedAt,
         labels: params.labels,
         state: 'running',
         reportedAt: createdAt,

package/src/db/reservations.ts

@@ -10,6 +10,7 @@ import {
   isNotNull,
   isNull,
   lt,
+  notExists,
   or,
   sql,
 } from 'drizzle-orm';
@@ -67,6 +68,12 @@ export interface InstallationPollDemandAndReserveParams {
   onReservations?: (reservations: ReservationGrant[]) => void;
 }
 
+type DemandScope = 'installation' | 'workspace';
+
+type PollDemandAndReserveLockedParams = PollDemandAndReserveParams & {
+  scope: DemandScope;
+};
+
 interface NormalizedTemplate {
   templateKey: string;
   labels: string[];
@@ -92,7 +99,7 @@ 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);
+  return await pollDemandAndReserveLockedTx(tx, {...params, scope: 'workspace'});
 }
 
 export async function pollInstallationDemandAndReserve(
@@ -126,6 +133,7 @@ export async function pollInstallationDemandAndReserve(
           availableSlots: template.remainingSlots,
         })),
         capabilityWindowSeconds: params.capabilityWindowSeconds,
+        scope: 'installation',
       });
     });
     results.push(result);
@@ -158,7 +166,7 @@ function consumeInstallationTemplateSlots(
 
 async function pollDemandAndReserveLockedTx(
   tx: Tx,
-  params: PollDemandAndReserveParams,
+  params: PollDemandAndReserveLockedParams,
 ): Promise<{stats: DemandStat[]; reservations: ReservationGrant[]}> {
   let demandRows = (
     await tx
@@ -265,6 +273,7 @@ async function pollDemandAndReserveLockedTx(
         requiredLabels: demand.requiredLabels,
         count: grant,
         workspaceId: params.workspaceId,
+        scope: params.scope,
       });
       reservedAfterGrant += grant;
       grants.push({
@@ -296,18 +305,61 @@ async function bindIdleRunnerInstancesTx(
     workspaceId: string;
     requiredLabels: string[];
     count: number;
+    scope: DemandScope;
   },
 ): Promise<void> {
+  // Reservation expiry is the activation grace deadline. A live reservation
+  // protects a runner that is still booting; an expired or missing row is stale.
+  const canBindRunner = () =>
+    and(
+      or(
+        and(isNull(providerRunners.workspaceId), isNull(providerRunners.reservationId)),
+        and(
+          isNotNull(providerRunners.reservationId),
+          params.scope === 'installation'
+            ? undefined
+            : or(
+                isNull(providerRunners.workspaceId),
+                eq(providerRunners.workspaceId, params.workspaceId),
+              ),
+          notExists(
+            tx
+              .select({id: reservations.id})
+              .from(reservations)
+              .where(
+                and(
+                  eq(reservations.id, providerRunners.reservationId),
+                  gt(reservations.expiresAt, sql`now()`),
+                ),
+              ),
+          ),
+        ),
+      ),
+      or(
+        isNull(providerRunners.intendedReservationId),
+        notExists(
+          tx
+            .select({id: reservations.id})
+            .from(reservations)
+            .where(
+              and(
+                eq(reservations.id, providerRunners.intendedReservationId),
+                gt(reservations.expiresAt, sql`now()`),
+              ),
+            ),
+        ),
+      ),
+      isNull(providerRunners.reservationReleasedAt),
+      isNull(providerRunners.runnerSessionId),
+    );
+
   const idleRunners = await tx
     .select({id: providerRunners.id})
     .from(providerRunners)
     .where(
       and(
         eq(providerRunners.provisionerId, params.provisionerId),
-        isNull(providerRunners.workspaceId),
-        isNull(providerRunners.reservationId),
-        isNull(providerRunners.intendedReservationId),
-        isNull(providerRunners.runnerSessionId),
+        canBindRunner(),
         isNotNull(providerRunners.providerRunnerId),
         eq(providerRunners.state, 'running'),
         arrayContains(providerRunners.labels, params.requiredLabels),
@@ -332,11 +384,26 @@ async function bindIdleRunnerInstancesTx(
 
   if (idleRunners.length === 0) return;
 
+  await tx
+    .update(runnerActivationTokens)
+    .set({revokedAt: sql`now()`})
+    .where(
+      and(
+        inArray(
+          runnerActivationTokens.runnerInstanceId,
+          idleRunners.map((runner) => runner.id),
+        ),
+        isNull(runnerActivationTokens.consumedAt),
+        isNull(runnerActivationTokens.revokedAt),
+      ),
+    );
+
   await tx
     .update(providerRunners)
     .set({
       workspaceId: params.workspaceId,
       reservationId: params.reservationId,
+      intendedReservationId: null,
       assignedAt: sql`now()`,
       updatedAt: sql`now()`,
     })
@@ -346,9 +413,7 @@ async function bindIdleRunnerInstancesTx(
           providerRunners.id,
           idleRunners.map((runner) => runner.id),
         ),
-        isNull(providerRunners.workspaceId),
-        isNull(providerRunners.reservationId),
-        isNull(providerRunners.runnerSessionId),
+        canBindRunner(),
       ),
     );
 }