@shipfox/api-runners 12.6.0 → 12.7.0

package/src/db/runner-instances.ts

@@ -28,9 +28,9 @@ import {
   type RunnerInstanceBoundJobExecution,
 } from './job-executions.js';
 import {lockRunnerReservationAdvisoryKeysTx} from './reservation-locks.js';
-import {releaseReservationUnits} from './reservations.js';
+import {releaseTerminalRunnerInstanceReservationsByIds} from './reservations.js';
 import {validateRunnerReservationCapacityTx} from './runner-assignments.js';
-import {terminalStates} from './runner-states.js';
+import {activeStates, terminalStates} from './runner-states.js';
 import {ephemeralRegistrationTokens} from './schema/ephemeral-registration-tokens.js';
 import {provisionerTokens} from './schema/provisioner-tokens.js';
 import {reservations} from './schema/reservations.js';
@@ -39,11 +39,6 @@ import {providerRunners, toRunnerInstance} from './schema/runner-instances.js';
 import {runnerSessions} from './schema/runner-sessions.js';
 import {runningJobExecutions} from './schema/running-job-executions.js';
 
-export const activeStates = [
-  'starting',
-  'running',
-  'stopping',
-] as const satisfies readonly RunnerInstanceState[];
 export const divergenceCountStates = ['starting', 'running'] as const satisfies readonly Extract<
   RunnerInstanceState,
   'starting' | 'running'
@@ -157,6 +152,14 @@ export async function reportRunnerInstances(params: ReportRunnerInstancesParams)
 
     const events = await hydrateRunnerSessionIdsFromConsumedTokens(tx, params, aggregatedEvents);
     const reservationSafeEvents = await guardReportedReservationIdsTx(tx, params, events);
+    const existingReportedAtByProviderRunnerId = hasTerminalEvent
+      ? await listExistingProviderRunnerReportedAtTx(tx, params, reservationSafeEvents)
+      : new Map<string, Date>();
+    const freshTerminalEvents = reservationSafeEvents.filter((event) => {
+      if (!isTerminalState(event.state)) return false;
+      const existingReportedAt = existingReportedAtByProviderRunnerId.get(event.providerRunnerId);
+      return existingReportedAt === undefined || event.reportedAt >= existingReportedAt;
+    });
     const terminateIntentsHonored = await listTerminateIntentsHonoredByTerminatedReportsTx(
       tx,
       params,
@@ -214,9 +217,10 @@ export async function reportRunnerInstances(params: ReportRunnerInstancesParams)
         `,
       });
 
-    const reservationsReleased = hasTerminalEvent
-      ? await releaseTerminalRunnerInstanceReservations(tx, params, reservationSafeEvents)
-      : 0;
+    const reservationsReleased =
+      freshTerminalEvents.length > 0
+        ? await releaseTerminalRunnerInstanceReservations(tx, params, freshTerminalEvents)
+        : 0;
 
     return {
       accepted: reservationSafeEvents.length,
@@ -226,21 +230,73 @@ export async function reportRunnerInstances(params: ReportRunnerInstancesParams)
   });
 }
 
+async function listExistingProviderRunnerReportedAtTx(
+  tx: Tx,
+  params: ReportRunnerInstancesParams,
+  events: RunnerInstanceReportRow[],
+): Promise<Map<string, Date>> {
+  const providerRunnerIds = [...new Set(events.map((event) => event.providerRunnerId))];
+  if (providerRunnerIds.length === 0) return new Map();
+
+  const rows = await tx
+    .select({
+      providerRunnerId: providerRunners.providerRunnerId,
+      reportedAt: providerRunners.reportedAt,
+    })
+    .from(providerRunners)
+    .where(
+      and(
+        eq(providerRunners.provisionerId, params.provisionerId),
+        inArray(providerRunners.providerRunnerId, providerRunnerIds),
+      ),
+    );
+  return new Map(
+    rows.flatMap((row) =>
+      row.providerRunnerId ? [[row.providerRunnerId, row.reportedAt] as const] : [],
+    ),
+  );
+}
+
 async function guardReportedReservationIdsTx(
   tx: Tx,
   params: ReportRunnerInstancesParams,
   events: RunnerInstanceReportRow[],
 ): Promise<RunnerInstanceReportRow[]> {
-  const reservationIds = [
+  const reportedReservationIds = [
     ...new Set(events.flatMap((event) => (event.reservationId ? [event.reservationId] : []))),
   ].sort();
-  if (reservationIds.length === 0) return events;
+  const providerRunnerIds = [...new Set(events.map((event) => event.providerRunnerId))];
+  if (reportedReservationIds.length === 0 && !events.some((event) => isTerminalState(event.state)))
+    return events;
+
+  // Lock every reservation currently associated with the reported runners before taking a
+  // runner row lock. Terminal reports may omit reservationId, or carry a stale one; discovering
+  // the stored IDs here keeps report projection in the same reservation-then-runner order as
+  // assignment and terminal cleanup.
+  const existingReservationRows = await tx
+    .select({
+      reservationId: providerRunners.reservationId,
+      intendedReservationId: providerRunners.intendedReservationId,
+    })
+    .from(providerRunners)
+    .where(
+      and(
+        eq(providerRunners.provisionerId, params.provisionerId),
+        inArray(providerRunners.providerRunnerId, providerRunnerIds),
+      ),
+    );
+  const storedReservationIds = existingReservationRows.flatMap((row) =>
+    [row.reservationId, row.intendedReservationId].filter(
+      (reservationId): reservationId is string => reservationId !== null,
+    ),
+  );
+  const reservationIds = [...new Set([...reportedReservationIds, ...storedReservationIds])].sort();
+  if (reservationIds.length > 0)
+    await lockRunnerReservationAdvisoryKeysTx(tx, {
+      provisionerId: params.provisionerId,
+      reservationIds,
+    });
 
-  await lockRunnerReservationAdvisoryKeysTx(tx, {
-    provisionerId: params.provisionerId,
-    reservationIds,
-  });
-  const providerRunnerIds = events.map((event) => event.providerRunnerId);
   const existingRows = await tx
     .select({
       providerRunnerId: providerRunners.providerRunnerId,
@@ -713,16 +769,15 @@ export async function reconcileRunnerInstances(
               ),
             ),
           )
-          .returning({providerRunnerId: providerRunners.providerRunnerId});
+          .returning({id: providerRunners.id, providerRunnerId: providerRunners.providerRunnerId});
 
         absentIds = updated.flatMap((row) => (row.providerRunnerId ? [row.providerRunnerId] : []));
-        if (params.workspaceId) {
-          reservationsReleased = await releaseTerminalRunnerInstanceReservationsByIds(tx, {
-            workspaceId: params.workspaceId,
-            provisionerId: params.provisionerId,
-            providerRunnerIds: absentIds,
-          });
-        }
+        reservationsReleased = await releaseTerminalRunnerInstanceReservationsByIds(tx, {
+          workspaceId: params.workspaceId,
+          provisionerId: params.provisionerId,
+          runnerInstanceIds: updated.map((row) => row.id),
+          requireUnlinkedSession: false,
+        });
       }
     }
 
@@ -782,10 +837,41 @@ export async function reapStaleRunnerInstances(params: {
       .orderBy(asc(providerRunners.updatedAt), asc(providerRunners.id))
       .limit(params.limit);
 
-    if (candidateRows.length === 0) return {reaped: 0, reservationsReleased: 0};
+    const terminalReservationRows =
+      candidateRows.length === 0
+        ? await tx
+            .select({
+              id: providerRunners.id,
+              workspaceId: providerRunners.workspaceId,
+              provisionerId: providerRunners.provisionerId,
+              providerRunnerId: providerRunners.providerRunnerId,
+            })
+            .from(providerRunners)
+            .where(
+              and(
+                inArray(providerRunners.state, terminalStates),
+                isNotNull(providerRunners.runnerSessionId),
+                lt(providerRunners.updatedAt, cutoff),
+                isNull(providerRunners.reservationReleasedAt),
+                or(
+                  isNotNull(providerRunners.reservationId),
+                  isNotNull(providerRunners.intendedReservationId),
+                ),
+              ),
+            )
+            .orderBy(asc(providerRunners.updatedAt), asc(providerRunners.id))
+            .limit(params.limit)
+        : [];
+
+    if (candidateRows.length === 0 && terminalReservationRows.length === 0)
+      return {reaped: 0, reservationsReleased: 0};
 
     const workspaceIds = [
-      ...new Set(candidateRows.flatMap((row) => (row.workspaceId ? [row.workspaceId] : []))),
+      ...new Set(
+        [...candidateRows, ...terminalReservationRows].flatMap((row) =>
+          row.workspaceId ? [row.workspaceId] : [],
+        ),
+      ),
     ].sort();
     for (const workspaceId of workspaceIds) {
       await tx.execute(sql`select pg_advisory_xact_lock(hashtext(${workspaceId}))`);
@@ -809,6 +895,7 @@ export async function reapStaleRunnerInstances(params: {
         ),
       )
       .returning({
+        id: providerRunners.id,
         workspaceId: providerRunners.workspaceId,
         provisionerId: providerRunners.provisionerId,
         providerRunnerId: providerRunners.providerRunnerId,
@@ -819,7 +906,15 @@ export async function reapStaleRunnerInstances(params: {
       reservationsReleased += await releaseTerminalRunnerInstanceReservationsByIds(tx, {
         workspaceId: group.workspaceId,
         provisionerId: group.provisionerId,
-        providerRunnerIds: group.providerRunnerIds,
+        runnerInstanceIds: group.runnerInstanceIds,
+        requireUnlinkedSession: false,
+      });
+    }
+    for (const group of groupRunnerInstanceIds(terminalReservationRows)) {
+      reservationsReleased += await releaseTerminalRunnerInstanceReservationsByIds(tx, {
+        workspaceId: group.workspaceId,
+        provisionerId: group.provisionerId,
+        runnerInstanceIds: group.runnerInstanceIds,
         requireUnlinkedSession: false,
       });
     }
@@ -840,174 +935,13 @@ async function releaseTerminalRunnerInstanceReservations(
     workspaceId: params.workspaceId,
     provisionerId: params.provisionerId,
     providerRunnerIds: terminalEvents.map((event) => event.providerRunnerId),
+    reportedAtByProviderRunnerId: new Map(
+      terminalEvents.map((event) => [event.providerRunnerId, event.reportedAt]),
+    ),
+    requireUnlinkedSession: false,
   });
 }
 
-async function releaseTerminalRunnerInstanceReservationsByIds(
-  tx: Tx,
-  params: {
-    workspaceId: string | null;
-    provisionerId: string;
-    providerRunnerIds: string[];
-    requireUnlinkedSession?: boolean;
-  },
-): 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,
-      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.provisionerId, params.provisionerId),
-        inArray(providerRunners.providerRunnerId, params.providerRunnerIds),
-        inArray(providerRunners.state, terminalStates),
-        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({
-      intendedReservationId: null,
-      reservationReleasedAt: sql`now()`,
-      updatedAt: sql`now()`,
-    })
-    .where(
-      and(
-        inArray(
-          providerRunners.id,
-          rows.map((row) => row.id),
-        ),
-        params.requireUnlinkedSession === false
-          ? undefined
-          : isNull(providerRunners.runnerSessionId),
-        isNull(providerRunners.reservationReleasedAt),
-      ),
-    )
-    .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;
-
-  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 {
   return sql`now() - (${thresholdSeconds} || ' seconds')::interval`;
 }
@@ -1058,24 +992,28 @@ function staleRunnerInstanceWhere(tx: Tx, cutoff: SQL): SQL<boolean> {
 
 function groupRunnerInstanceIds(
   rows: Array<{
+    id: string;
     workspaceId: string | null;
     provisionerId: string;
     providerRunnerId: string | null;
   }>,
-): Array<{workspaceId: string | null; provisionerId: string; providerRunnerIds: string[]}> {
+): Array<{
+  workspaceId: string | null;
+  provisionerId: string;
+  runnerInstanceIds: string[];
+}> {
   const groups = new Map<
     string,
-    {workspaceId: string | null; provisionerId: string; providerRunnerIds: string[]}
+    {workspaceId: string | null; provisionerId: string; runnerInstanceIds: string[]}
   >();
   for (const row of rows) {
-    if (!row.providerRunnerId) continue;
     const key = `${row.workspaceId ?? ''}:${row.provisionerId}`;
     const group = groups.get(key) ?? {
       workspaceId: row.workspaceId,
       provisionerId: row.provisionerId,
-      providerRunnerIds: [],
+      runnerInstanceIds: [],
     };
-    group.providerRunnerIds.push(row.providerRunnerId);
+    group.runnerInstanceIds.push(row.id);
     groups.set(key, group);
   }
   return [...groups.values()];

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

@@ -0,0 +1,16 @@
+import {providerRunnerStateSchema} from '@shipfox/api-runners-dto';
+import {activeStates, terminalStates} from '#db/runner-states.js';
+
+describe('runner state partitions', () => {
+  it('covers every provider state exactly once', () => {
+    const states = [...activeStates, ...terminalStates];
+
+    expect(new Set(states).size).toBe(states.length);
+    expect([...states].sort()).toEqual([...providerRunnerStateSchema.options].sort());
+    expect([...terminalStates]).toEqual(
+      providerRunnerStateSchema.options.filter(
+        (state) => !activeStates.includes(state as (typeof activeStates)[number]),
+      ),
+    );
+  });
+});

package/src/db/runner-states.ts

@@ -1,7 +1,12 @@
+import {providerRunnerStateSchema} from '@shipfox/api-runners-dto';
 import type {RunnerInstanceState} from '#core/entities/runner-instance.js';
 
-export const terminalStates = [
-  'stopped',
-  'failed',
-  'terminated',
+export const activeStates = [
+  'starting',
+  'running',
+  'stopping',
 ] as const satisfies readonly RunnerInstanceState[];
+
+export const terminalStates = providerRunnerStateSchema.options.filter(
+  (state) => !activeStates.includes(state as (typeof activeStates)[number]),
+) as readonly Extract<RunnerInstanceState, 'stopped' | 'failed' | 'terminated'>[];