@shipfox/api-runners 9.3.0 → 10.1.0

package/src/db/runner-assignments.ts

@@ -1,11 +1,13 @@
-import {and, eq, inArray, isNull, sql} from 'drizzle-orm';
+import {and, eq, inArray, isNull, notInArray, or, sql} from 'drizzle-orm';
 import {
   ReservationExpiredError,
   ReservationNotFoundError,
   RunnerInstanceAlreadyAssignedError,
   RunnerInstanceNotAssignableError,
 } from '#core/errors.js';
+import type {Tx} from './db.js';
 import {db} from './db.js';
+import {terminalStates} from './runner-instances.js';
 import {reservations} from './schema/reservations.js';
 import {runnerControlSessions} from './schema/runner-control-sessions.js';
 import {providerRunners} from './schema/runner-instances.js';
@@ -16,100 +18,146 @@ export async function assignRunnerInstances(params: {
   reservationId: string;
   runnerInstanceIds: string[];
 }): Promise<string[]> {
+  return await db().transaction(async (tx) => assignRunnerInstancesTx(tx, params));
+}
+
+export async function assignRunnerInstancesTx(
+  tx: Tx,
+  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);
+  await tx.execute(
+    sql`select pg_advisory_xact_lock(hashtext(${`runners_assignment:${params.provisionerId}:${params.reservationId}`}))`,
+  );
+  const runnerRows = await tx
+    .select({
+      id: providerRunners.id,
+      reservationId: providerRunners.reservationId,
+      intendedReservationId: providerRunners.intendedReservationId,
+      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 runnerRows = await tx
-      .select({
-        id: providerRunners.id,
-        reservationId: providerRunners.reservationId,
-        workspaceId: providerRunners.workspaceId,
-        providerRunnerId: providerRunners.providerRunnerId,
-        labels: providerRunners.labels,
-        state: providerRunners.state,
-      })
-      .from(providerRunners)
+  // The assignment outlives its short reservation row. A retry after maintenance
+  // cleanup is complete when every owned runner already carries the requested assignment.
+  if (
+    runnerInstanceIds.length > 0 &&
+    runnerRows.length === runnerInstanceIds.length &&
+    runnerRows.every((runner) => runner.reservationId === params.reservationId)
+  ) {
+    await tx
+      .update(providerRunners)
+      .set({intendedReservationId: null, updatedAt: sql`now()`})
       .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),
+    return params.runnerInstanceIds;
+  }
+
+  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);
+  if (runnerRows.length !== runnerInstanceIds.length)
+    throw new RunnerInstanceNotAssignableError(runnerInstanceIds[0] ?? '');
+
+  const activeControlSessions = await tx
+    .select({runnerInstanceId: runnerControlSessions.runnerInstanceId})
+    .from(runnerControlSessions)
+    .where(
+      and(
+        inArray(
+          runnerControlSessions.runnerInstanceId,
+          runnerRows.map((runner) => runner.id),
+        ),
+        isNull(runnerControlSessions.closedAt),
+      ),
     );
-    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 runnerInstanceIdsWithControlSession = new Set(
+    activeControlSessions.map((session) => session.runnerInstanceId),
+  );
+  const runners = runnerRows.map((runner) => ({
+    ...runner,
+    controlSessionId: runnerInstanceIdsWithControlSession.has(runner.id) ? 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 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),
+  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(
+      and(
+        eq(providerRunners.provisionerId, params.provisionerId),
+        or(
+          and(
+            eq(providerRunners.reservationId, reservation.id),
+            isNull(providerRunners.reservationReleasedAt),
           ),
-        );
-    }
-    return params.runnerInstanceIds;
-  });
+          and(
+            eq(providerRunners.intendedReservationId, reservation.id),
+            notInArray(providerRunners.id, runnerInstanceIds),
+            isNull(providerRunners.reservationReleasedAt),
+            notInArray(providerRunners.state, [...terminalStates]),
+          ),
+        ),
+      ),
+    );
+  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,
+        intendedReservationId: null,
+        assignedAt: sql`now()`,
+        updatedAt: sql`now()`,
+      })
+      .where(
+        inArray(
+          providerRunners.id,
+          newRunners.map((runner) => runner.id),
+        ),
+      );
+  }
+  return params.runnerInstanceIds;
 }

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

@@ -413,6 +413,134 @@ describe('reportRunnerInstances', () => {
     expect(providerRunnerRows[0]?.reservationReleasedAt).toBeInstanceOf(Date);
   });
 
+  it('releases an intended reservation for a runner that dies before enrollment', async () => {
+    const reservationId = await createReservation(2);
+    const [runner] = await db()
+      .insert(providerRunners)
+      .values({
+        provisionerId,
+        intendedReservationId: reservationId,
+        providerRunnerId: 'pre-enrollment-runner',
+        state: 'starting',
+        reportedAt: new Date(),
+      })
+      .returning();
+    if (!runner) throw new Error('Runner instance insert returned no row');
+    if (!runner.providerRunnerId) throw new Error('Runner instance provider id missing');
+
+    const result = await reportRunnerInstances({
+      workspaceId,
+      provisionerId,
+      events: [event({providerRunnerId: runner.providerRunnerId, state: 'failed'})],
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    const reservationRows = await reservationRowsFor({workspaceId, provisionerId});
+    expect(result).toEqual({accepted: 1, reservationsReleased: 1, terminateIntentsHonored: []});
+    expect(reservationRows[0]?.count).toBe(1);
+    expect(storedRunner?.intendedReservationId).toBeNull();
+    expect(storedRunner?.reservationReleasedAt).toBeInstanceOf(Date);
+  });
+
+  it('releases an intended reservation from an installation-scoped terminal report', async () => {
+    const installationProvisioner = await provisionerTokenFactory.create({
+      scope: 'installation',
+      workspaceId: null,
+    });
+    const reservationWorkspaceId = crypto.randomUUID();
+    await reservationFactory.create({
+      workspaceId: reservationWorkspaceId,
+      provisionerId: installationProvisioner.id,
+      requiredLabels: ['linux'],
+      count: 2,
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+    const [reservation] = await db()
+      .select()
+      .from(reservations)
+      .where(
+        and(
+          eq(reservations.workspaceId, reservationWorkspaceId),
+          eq(reservations.provisionerId, installationProvisioner.id),
+        ),
+      )
+      .orderBy(desc(reservations.id))
+      .limit(1);
+    if (!reservation) throw new Error('Expected reservation');
+    await db().insert(providerRunners).values({
+      provisionerId: installationProvisioner.id,
+      intendedReservationId: reservation.id,
+      providerRunnerId: 'installation-pre-enrollment-runner',
+      state: 'starting',
+      reportedAt: new Date(),
+    });
+
+    const result = await reportRunnerInstances({
+      workspaceId: null,
+      provisionerId: installationProvisioner.id,
+      events: [event({providerRunnerId: 'installation-pre-enrollment-runner', state: 'failed'})],
+    });
+
+    const [storedReservation] = await db()
+      .select()
+      .from(reservations)
+      .where(eq(reservations.id, reservation.id));
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.providerRunnerId, 'installation-pre-enrollment-runner'));
+    expect(result).toEqual({accepted: 1, reservationsReleased: 1, terminateIntentsHonored: []});
+    expect(storedReservation?.count).toBe(1);
+    expect(storedRunner?.intendedReservationId).toBeNull();
+    expect(storedRunner?.reservationReleasedAt).toBeInstanceOf(Date);
+  });
+
+  it('prefers the validated intended reservation when both reservation ids are set', async () => {
+    const installationProvisioner = await provisionerTokenFactory.create({
+      scope: 'installation',
+      workspaceId: null,
+    });
+    const staleReservationId = await createReservation(1, {
+      workspaceId: crypto.randomUUID(),
+      provisionerId: installationProvisioner.id,
+    });
+    const intendedReservationId = await createReservation(1, {
+      workspaceId: crypto.randomUUID(),
+      provisionerId: installationProvisioner.id,
+    });
+    await db().insert(providerRunners).values({
+      provisionerId: installationProvisioner.id,
+      reservationId: staleReservationId,
+      intendedReservationId,
+      providerRunnerId: 'installation-runner-with-stale-reservation',
+      state: 'starting',
+      reportedAt: new Date(),
+    });
+
+    const result = await reportRunnerInstances({
+      workspaceId: null,
+      provisionerId: installationProvisioner.id,
+      events: [
+        event({providerRunnerId: 'installation-runner-with-stale-reservation', state: 'failed'}),
+      ],
+    });
+
+    const [staleReservation] = await db()
+      .select()
+      .from(reservations)
+      .where(eq(reservations.id, staleReservationId));
+    const [intendedReservation] = await db()
+      .select()
+      .from(reservations)
+      .where(eq(reservations.id, intendedReservationId));
+    expect(result).toEqual({accepted: 1, reservationsReleased: 1, terminateIntentsHonored: []});
+    expect(staleReservation?.count).toBe(1);
+    expect(intendedReservation).toBeUndefined();
+  });
+
   it('does not release a reservation owned by another workspace or provisioner', async () => {
     const otherWorkspaceReservationId = await createReservation(1, {
       workspaceId: crypto.randomUUID(),
@@ -1045,6 +1173,60 @@ describe('reapStaleRunnerInstances', () => {
     expect(row).toMatchObject({workspaceId: null, providerRunnerId: null, state: 'failed'});
   });
 
+  it('releases an intended reservation while reaping a stale installation runner', async () => {
+    const installationProvisioner = await provisionerTokenFactory.create({
+      scope: 'installation',
+      workspaceId: null,
+    });
+    const reservationWorkspaceId = crypto.randomUUID();
+    await reservationFactory.create({
+      workspaceId: reservationWorkspaceId,
+      provisionerId: installationProvisioner.id,
+      requiredLabels: ['linux'],
+      count: 1,
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+    const [reservation] = await db()
+      .select()
+      .from(reservations)
+      .where(
+        and(
+          eq(reservations.workspaceId, reservationWorkspaceId),
+          eq(reservations.provisionerId, installationProvisioner.id),
+        ),
+      )
+      .orderBy(desc(reservations.id))
+      .limit(1);
+    if (!reservation) throw new Error('Expected reservation');
+    const [instance] = await db()
+      .insert(providerRunners)
+      .values({
+        provisionerId: installationProvisioner.id,
+        intendedReservationId: reservation.id,
+        providerRunnerId: 'stale-installation-runner-with-reservation',
+        state: 'starting',
+        reportedAt: staleAt(),
+        updatedAt: staleAt(),
+      })
+      .returning({id: providerRunners.id});
+    if (!instance) throw new Error('Runner instance insert returned no row');
+
+    const result = await reapStaleRunnerInstances({thresholdSeconds: 60, limit: 100});
+    const [row] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, instance.id));
+    const remainingReservations = await db()
+      .select()
+      .from(reservations)
+      .where(eq(reservations.id, reservation.id));
+
+    expect(result).toEqual({reaped: 1, reservationsReleased: 1});
+    expect(row?.intendedReservationId).toBeNull();
+    expect(row?.reservationReleasedAt).toBeInstanceOf(Date);
+    expect(remainingReservations).toHaveLength(0);
+  });
+
   it('skips fresh rows, live provisioners, terminal rows, running jobs, and fresh sessions', async () => {
     await createRunnerInstance({
       providerRunnerId: 'fresh-row',

package/src/db/runner-instances.ts

@@ -28,6 +28,7 @@ import {
 import {releaseReservationUnits} from './reservations.js';
 import {ephemeralRegistrationTokens} from './schema/ephemeral-registration-tokens.js';
 import {provisionerTokens} from './schema/provisioner-tokens.js';
+import {reservations} from './schema/reservations.js';
 import {providerRunners, toRunnerInstance} from './schema/runner-instances.js';
 import {runnerSessions} from './schema/runner-sessions.js';
 import {runningJobExecutions} from './schema/running-job-executions.js';
@@ -563,7 +564,7 @@ async function releaseTerminalRunnerInstanceReservations(
   events: RunnerInstanceReportEvent[],
 ): Promise<number> {
   const terminalEvents = events.filter((event) => isTerminalState(event.state));
-  if (terminalEvents.length === 0 || !params.workspaceId) return 0;
+  if (terminalEvents.length === 0) return 0;
 
   return await releaseTerminalRunnerInstanceReservationsByIds(tx, {
     workspaceId: params.workspaceId,
@@ -575,7 +576,7 @@ async function releaseTerminalRunnerInstanceReservations(
 async function releaseTerminalRunnerInstanceReservationsByIds(
   tx: Tx,
   params: {
-    workspaceId: string;
+    workspaceId: string | null;
     provisionerId: string;
     providerRunnerIds: string[];
     requireUnlinkedSession?: boolean;
@@ -583,31 +584,115 @@ async function releaseTerminalRunnerInstanceReservationsByIds(
 ): Promise<number> {
   if (params.providerRunnerIds.length === 0) return 0;
 
+  const reservationWorkspacePredicate =
+    params.workspaceId === null
+      ? sql``
+      : sql`and ${eq(reservations.workspaceId, params.workspaceId)}`;
+
   const rows = await tx
     .select({
       id: providerRunners.id,
-      reservationId: providerRunners.reservationId,
+      releaseReservationId: sql<string | null>`coalesce(
+        (select ${reservations.id}
+         from ${reservations}
+         where ${reservations.id} = ${providerRunners.intendedReservationId}
+           and ${reservations.provisionerId} = ${params.provisionerId}
+           ${reservationWorkspacePredicate}),
+        (select ${reservations.id}
+         from ${reservations}
+         where ${reservations.id} = ${providerRunners.reservationId}
+           and ${reservations.provisionerId} = ${params.provisionerId}
+           ${reservationWorkspacePredicate})
+      )`,
+      releaseReservationWorkspaceId: sql<string | null>`coalesce(
+        (select ${reservations.workspaceId}
+         from ${reservations}
+         where ${reservations.id} = ${providerRunners.intendedReservationId}
+           and ${reservations.provisionerId} = ${params.provisionerId}
+           ${reservationWorkspacePredicate}),
+        (select ${reservations.workspaceId}
+         from ${reservations}
+         where ${reservations.id} = ${providerRunners.reservationId}
+           and ${reservations.provisionerId} = ${params.provisionerId}
+           ${reservationWorkspacePredicate})
+      )`,
     })
     .from(providerRunners)
     .where(
       and(
-        eq(providerRunners.workspaceId, params.workspaceId),
         eq(providerRunners.provisionerId, params.provisionerId),
         inArray(providerRunners.providerRunnerId, params.providerRunnerIds),
         inArray(providerRunners.state, terminalStates),
-        isNotNull(providerRunners.reservationId),
+        or(
+          isNotNull(providerRunners.reservationId),
+          isNotNull(providerRunners.intendedReservationId),
+        ),
+        params.workspaceId === null
+          ? undefined
+          : or(
+              and(
+                eq(providerRunners.workspaceId, params.workspaceId),
+                or(
+                  exists(
+                    tx
+                      .select({id: reservations.id})
+                      .from(reservations)
+                      .where(
+                        and(
+                          eq(reservations.workspaceId, params.workspaceId),
+                          eq(reservations.provisionerId, params.provisionerId),
+                          eq(reservations.id, providerRunners.reservationId),
+                        ),
+                      ),
+                  ),
+                  exists(
+                    tx
+                      .select({id: reservations.id})
+                      .from(reservations)
+                      .where(
+                        and(
+                          eq(reservations.workspaceId, params.workspaceId),
+                          eq(reservations.provisionerId, params.provisionerId),
+                          eq(reservations.id, providerRunners.intendedReservationId),
+                        ),
+                      ),
+                  ),
+                ),
+              ),
+              and(
+                isNull(providerRunners.workspaceId),
+                isNotNull(providerRunners.intendedReservationId),
+                exists(
+                  tx
+                    .select({id: reservations.id})
+                    .from(reservations)
+                    .where(
+                      and(
+                        eq(reservations.workspaceId, params.workspaceId),
+                        eq(reservations.provisionerId, params.provisionerId),
+                        eq(reservations.id, providerRunners.intendedReservationId),
+                      ),
+                    ),
+                ),
+              ),
+            ),
         params.requireUnlinkedSession === false
           ? undefined
           : isNull(providerRunners.runnerSessionId),
         isNull(providerRunners.reservationReleasedAt),
       ),
-    );
+    )
+    .for('update');
 
   if (rows.length === 0) return 0;
 
   const updated = await tx
     .update(providerRunners)
-    .set({reservationReleasedAt: sql`now()`, updatedAt: sql`now()`})
+    .set({
+      intendedReservationId: null,
+      reservationReleasedAt: sql`now()`,
+      updatedAt: sql`now()`,
+    })
     .where(
       and(
         inArray(
@@ -620,27 +705,37 @@ async function releaseTerminalRunnerInstanceReservationsByIds(
         isNull(providerRunners.reservationReleasedAt),
       ),
     )
-    .returning({reservationId: providerRunners.reservationId});
-
-  const releasesByReservationId = new Map<string, number>();
-  for (const row of updated) {
-    if (!row.reservationId) continue;
-    releasesByReservationId.set(
-      row.reservationId,
-      (releasesByReservationId.get(row.reservationId) ?? 0) + 1,
-    );
+    .returning({id: providerRunners.id});
+
+  const releasesByReservationId = new Map<
+    string,
+    {workspaceId: string; reservationId: string; count: number}
+  >();
+  const updatedIds = new Set(updated.map((row) => row.id));
+  for (const row of rows) {
+    if (!updatedIds.has(row.id)) continue;
+    if (!row.releaseReservationId || !row.releaseReservationWorkspaceId) continue;
+    const key = `${row.releaseReservationWorkspaceId}:${row.releaseReservationId}`;
+    const release = releasesByReservationId.get(key) ?? {
+      workspaceId: row.releaseReservationWorkspaceId,
+      reservationId: row.releaseReservationId,
+      count: 0,
+    };
+    release.count += 1;
+    releasesByReservationId.set(key, release);
   }
 
   if (releasesByReservationId.size === 0) return 0;
 
-  return await releaseReservationUnits(tx, {
-    workspaceId: params.workspaceId,
-    provisionerId: params.provisionerId,
-    releases: [...releasesByReservationId].map(([reservationId, count]) => ({
-      reservationId,
-      count,
-    })),
-  });
+  let released = 0;
+  for (const release of releasesByReservationId.values()) {
+    released += await releaseReservationUnits(tx, {
+      workspaceId: release.workspaceId,
+      provisionerId: params.provisionerId,
+      releases: [{reservationId: release.reservationId, count: release.count}],
+    });
+  }
+  return released;
 }
 
 function staleRunnerInstanceCutoff(thresholdSeconds: number): SQL {
@@ -697,14 +792,14 @@ function groupRunnerInstanceIds(
     provisionerId: string;
     providerRunnerId: string | null;
   }>,
-): Array<{workspaceId: string; provisionerId: string; providerRunnerIds: string[]}> {
+): Array<{workspaceId: string | null; provisionerId: string; providerRunnerIds: string[]}> {
   const groups = new Map<
     string,
-    {workspaceId: string; provisionerId: string; providerRunnerIds: string[]}
+    {workspaceId: string | null; provisionerId: string; providerRunnerIds: string[]}
   >();
   for (const row of rows) {
-    if (!row.workspaceId || !row.providerRunnerId) continue;
-    const key = `${row.workspaceId}:${row.provisionerId}`;
+    if (!row.providerRunnerId) continue;
+    const key = `${row.workspaceId ?? ''}:${row.provisionerId}`;
     const group = groups.get(key) ?? {
       workspaceId: row.workspaceId,
       provisionerId: row.provisionerId,

package/src/db/schema/runner-instances.ts

@@ -23,6 +23,7 @@ export const providerRunners = pgTable(
     workspaceId: uuid('workspace_id'),
     provisionerId: uuid('provisioner_id').notNull(),
     providerRunnerId: text('provider_runner_id'),
+    intendedReservationId: uuid('intended_reservation_id'),
     reservationId: uuid('reservation_id'),
     assignedAt: timestamp('assigned_at', {withTimezone: true}),
     templateKey: text('template_key'),
@@ -53,6 +54,14 @@ export const providerRunners = pgTable(
       table.updatedAt,
       table.reportedAt,
     ),
+    index('runners_runner_instances_provisioner_intended_reservation_idx')
+      .on(table.provisionerId, table.intendedReservationId)
+      .where(
+        sql`${table.intendedReservationId} is not null and ${table.reservationReleasedAt} is null`,
+      ),
+    index('runners_runner_instances_provisioner_reservation_idx')
+      .on(table.provisionerId, table.reservationId)
+      .where(sql`${table.reservationId} is not null`),
     index('runners_runner_instances_active_template_counts_idx')
       .on(table.provisionerId, table.state, table.templateKey)
       .where(sql`"state" in ('starting', 'running') and "template_key" is not null`),