@shipfox/api-runners 12.3.0 → 12.4.1

package/dist/db/reservations.js

@@ -1,6 +1,9 @@
 import { canonicalizeLabels } from '@shipfox/runner-labels';
-import { and, arrayContains, asc, eq, exists, gt, inArray, isNotNull, isNull, lt, notExists, or, sql } from 'drizzle-orm';
+import { and, arrayContains, asc, eq, exists, gt, inArray, isNotNull, isNull, lt, notExists, notInArray, or, sql } from 'drizzle-orm';
+import { recordProviderRunnerActivationOutcome } from '#metrics/instance.js';
 import { db } from './db.js';
+import { lockRunnerReservationAdvisoryKeysTx } from './reservation-locks.js';
+import { terminalStates } from './runner-states.js';
 import { pendingJobExecutions } from './schema/pending-job-executions.js';
 import { provisionerCapabilitySnapshots } from './schema/provisioner-capability-snapshots.js';
 import { provisionerTokens } from './schema/provisioner-tokens.js';
@@ -9,9 +12,13 @@ import { runnerActivationTokens } from './schema/runner-activation-tokens.js';
 import { runnerControlSessions } from './schema/runner-control-sessions.js';
 import { providerRunners } from './schema/runner-instances.js';
 export async function pollDemandAndReserve(params) {
-    return await db().transaction(async (tx)=>{
+    const result = await db().transaction(async (tx)=>{
         return await pollDemandAndReserveTx(tx, params);
     });
+    return {
+        stats: result.stats,
+        reservations: result.reservations
+    };
 }
 export async function pollDemandAndReserveTx(tx, params) {
     await tx.execute(sql`select pg_advisory_xact_lock(hashtext(${params.workspaceId}))`);
@@ -38,13 +45,17 @@ export async function pollInstallationDemandAndReserve(params) {
             const locked = lockResult.rows[0];
             if (!locked?.locked) return {
                 stats: [],
-                reservations: []
+                reservations: [],
+                newlyReservedUnits: []
             };
             return await pollDemandAndReserveLockedTx(tx, {
                 workspaceId,
                 provisionerId: params.provisionerId,
                 maxReservations: remainingMaxReservations,
                 ttlSeconds: params.ttlSeconds,
+                ...params.activationGraceSeconds !== undefined ? {
+                    activationGraceSeconds: params.activationGraceSeconds
+                } : {},
                 templates: remainingTemplates.map((template)=>({
                         ...template,
                         availableSlots: template.remainingSlots
@@ -56,15 +67,21 @@ export async function pollInstallationDemandAndReserve(params) {
         results.push(result);
         consumeInstallationTemplateSlots(remainingTemplates, result.reservations);
         params.onReservations?.(result.reservations);
-        remainingMaxReservations -= result.reservations.reduce((total, reservation)=>total + reservation.count, 0);
+        remainingMaxReservations -= result.newlyReservedUnits.reduce((total, reservation)=>total + reservation.count, 0);
     }
+    const newlyReservedCount = results.reduce((total, result)=>total + result.newlyReservedUnits.reduce((subtotal, reservation)=>subtotal + reservation.count, 0), 0);
     return {
         stats: results.flatMap((result)=>result.stats),
-        reservations: results.flatMap((result)=>result.reservations)
+        reservations: results.flatMap((result)=>result.reservations),
+        ...newlyReservedCount > 0 ? {
+            newlyReservedCount
+        } : {}
     };
 }
-function consumeInstallationTemplateSlots(templates, reservations) {
-    for (const reservation of reservations){
+function consumeInstallationTemplateSlots(templates, launchGrants) {
+    // Adopted runners are already included in the running count behind availableSlots.
+    // Only units that still need a launch consume advertised template capacity.
+    for (const reservation of launchGrants){
         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);
     }
@@ -93,10 +110,10 @@ async function pollDemandAndReserveLockedTx(tx, params) {
             labelKey(row.requiredLabels),
             row.reserved
         ]));
-    const activeProvisionerReservationRows = await tx.select({
-        requiredLabels: reservations.requiredLabels,
-        reserved: sql`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 activeProvisionerReservationRows = await listActiveProvisionerReservationRowsTx(tx, {
+        workspaceId: params.workspaceId,
+        provisionerId: params.provisionerId
+    });
     const templates = params.templates.map((template)=>({
             templateKey: template.templateKey,
             labels: [
@@ -107,6 +124,7 @@ async function pollDemandAndReserveLockedTx(tx, params) {
     deductProvisionerReservations(templates, activeProvisionerReservationRows);
     const stats = [];
     const grants = [];
+    const newlyReservedUnits = [];
     let remainingMaxReservations = params.maxReservations;
     const workspaceIdField = params.capabilityWindowSeconds === undefined ? {} : {
         workspaceId: params.workspaceId
@@ -120,35 +138,68 @@ async function pollDemandAndReserveLockedTx(tx, params) {
         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);
-            await bindIdleRunnerInstancesTx(tx, {
+            const idleRunners = await listIdleRunnerInstancesTx(tx, {
                 provisionerId: params.provisionerId,
-                reservationId: inserted.id,
                 requiredLabels: demand.requiredLabels,
                 count: grant,
                 workspaceId: params.workspaceId,
                 scope: params.scope
             });
-            reservedAfterGrant += grant;
-            grants.push({
-                reservationId: inserted.id,
-                ...workspaceIdField,
+            remainingMaxReservations -= grant;
+            const boundCount = idleRunners.length;
+            if (boundCount > 0) {
+                const [boundReservation] = await tx.insert(reservations).values({
+                    workspaceId: params.workspaceId,
+                    provisionerId: params.provisionerId,
+                    requiredLabels: demand.requiredLabels,
+                    count: boundCount,
+                    kind: 'bound',
+                    expiresAt: sql`now() + (${params.activationGraceSeconds ?? params.ttlSeconds} || ' seconds')::interval`
+                }).returning({
+                    id: reservations.id
+                });
+                if (!boundReservation) throw new Error('Insert returned no rows');
+                await bindIdleRunnerInstancesTx(tx, {
+                    provisionerId: params.provisionerId,
+                    reservationId: boundReservation.id,
+                    workspaceId: params.workspaceId,
+                    requiredLabels: demand.requiredLabels,
+                    scope: params.scope,
+                    idleRunners
+                });
+            }
+            const launchCount = grant - boundCount;
+            drawSlots(satisfyingTemplates, launchCount);
+            let launchReservation;
+            if (launchCount > 0) {
+                const [inserted] = await tx.insert(reservations).values({
+                    workspaceId: params.workspaceId,
+                    provisionerId: params.provisionerId,
+                    requiredLabels: demand.requiredLabels,
+                    count: launchCount,
+                    kind: 'launch',
+                    expiresAt: sql`now() + (${params.ttlSeconds} || ' seconds')::interval`
+                }).returning({
+                    id: reservations.id,
+                    expiresAt: reservations.expiresAt
+                });
+                if (!inserted) throw new Error('Insert returned no rows');
+                launchReservation = inserted;
+            }
+            newlyReservedUnits.push({
                 labels: demand.requiredLabels,
-                count: grant,
-                expiresAt: inserted.expiresAt
+                count: grant
             });
+            reservedAfterGrant += grant;
+            if (launchReservation) {
+                grants.push({
+                    reservationId: launchReservation.id,
+                    ...workspaceIdField,
+                    labels: demand.requiredLabels,
+                    count: launchCount,
+                    expiresAt: launchReservation.expiresAt
+                });
+            }
         }
         stats.push({
             ...workspaceIdField,
@@ -160,33 +211,160 @@ async function pollDemandAndReserveLockedTx(tx, params) {
     }
     return {
         stats,
-        reservations: grants
+        reservations: grants,
+        newlyReservedUnits
     };
 }
-async function bindIdleRunnerInstancesTx(tx, params) {
-    // 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), canBindRunner(), isNotNull(providerRunners.providerRunnerId), eq(providerRunners.state, 'running'), arrayContains(providerRunners.labels, params.requiredLabels), exists(tx.select({
+async function listActiveProvisionerReservationRowsTx(tx, params) {
+    const candidateRows = await tx.select({
+        id: reservations.id
+    }).from(reservations).where(and(eq(reservations.workspaceId, params.workspaceId), eq(reservations.provisionerId, params.provisionerId), gt(reservations.expiresAt, sql`now()`)));
+    const candidateIds = candidateRows.map((row)=>row.id);
+    if (candidateIds.length === 0) return [];
+    // Assignment and provider-report transactions use this key before changing runner
+    // links. Lock it before counting so an in-flight assignment cannot expose its unit twice.
+    await lockRunnerReservationAdvisoryKeysTx(tx, {
+        provisionerId: params.provisionerId,
+        reservationIds: candidateIds
+    });
+    const activeRows = await tx.select({
+        id: reservations.id,
+        requiredLabels: reservations.requiredLabels,
+        count: reservations.count
+    }).from(reservations).where(and(eq(reservations.workspaceId, params.workspaceId), eq(reservations.provisionerId, params.provisionerId), gt(reservations.expiresAt, sql`now()`))).orderBy(asc(reservations.id)).for('update');
+    if (activeRows.length === 0) return [];
+    const activeIds = activeRows.map((row)=>row.id);
+    const activeIdSet = new Set(activeIds);
+    const usedRunnerRows = await tx.select({
+        id: providerRunners.id,
+        reservationId: providerRunners.reservationId,
+        intendedReservationId: providerRunners.intendedReservationId,
+        state: providerRunners.state
+    }).from(providerRunners).where(and(eq(providerRunners.provisionerId, params.provisionerId), isNull(providerRunners.reservationReleasedAt), or(inArray(providerRunners.reservationId, activeIds), and(inArray(providerRunners.intendedReservationId, activeIds), notInArray(providerRunners.state, [
+        ...terminalStates
+    ])))));
+    const usedByReservation = new Map();
+    for (const runner of usedRunnerRows){
+        if (runner.reservationId && activeIdSet.has(runner.reservationId)) {
+            addUsedRunner(usedByReservation, runner.reservationId, runner.id);
+        }
+        if (runner.intendedReservationId && activeIdSet.has(runner.intendedReservationId) && !terminalStates.includes(runner.state)) {
+            addUsedRunner(usedByReservation, runner.intendedReservationId, runner.id);
+        }
+    }
+    return activeRows.flatMap((reservation)=>{
+        const used = usedByReservation.get(reservation.id)?.size ?? 0;
+        const pending = Math.max(0, reservation.count - used);
+        return pending > 0 ? [
+            {
+                requiredLabels: reservation.requiredLabels,
+                reserved: pending
+            }
+        ] : [];
+    });
+}
+function addUsedRunner(usedByReservation, reservationId, runnerId) {
+    const runnerIds = usedByReservation.get(reservationId);
+    if (runnerIds) runnerIds.add(runnerId);
+    else usedByReservation.set(reservationId, new Set([
+        runnerId
+    ]));
+}
+// An expired or missing reservation is stale. A live reservation protects a runner that is
+// still booting or waiting for its activation grace period.
+function canBindRunner(tx, params) {
+    return 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));
+}
+function isBindableRunner(tx, params) {
+    return and(eq(providerRunners.provisionerId, params.provisionerId), canBindRunner(tx, params), isNotNull(providerRunners.providerRunnerId), eq(providerRunners.state, 'running'), arrayContains(providerRunners.labels, params.requiredLabels), exists(tx.select({
         id: runnerControlSessions.id
-    }).from(runnerControlSessions).where(and(eq(runnerControlSessions.runnerInstanceId, providerRunners.id), eq(runnerControlSessions.provisionerId, params.provisionerId), isNull(runnerControlSessions.closedAt), gt(runnerControlSessions.expiresAt, sql`now()`)))))).orderBy(asc(providerRunners.createdAt), asc(providerRunners.id)).limit(params.count).for('update');
-    if (idleRunners.length === 0) return;
+    }).from(runnerControlSessions).where(and(eq(runnerControlSessions.runnerInstanceId, providerRunners.id), eq(runnerControlSessions.provisionerId, params.provisionerId), isNull(runnerControlSessions.closedAt), gt(runnerControlSessions.expiresAt, sql`now()`)))));
+}
+async function listIdleRunnerInstancesTx(tx, params) {
+    const lockedRunners = await selectIdleRunnerInstancesTx(tx, params);
+    if (lockedRunners.length === 0) return lockedRunners;
+    // A SELECT ... FOR UPDATE can wake with a fresh target row but the original
+    // statement snapshot for reservation subqueries. Re-check the full predicate
+    // in a new statement while retaining the row locks from the first query.
+    return await selectIdleRunnerInstancesTx(tx, {
+        ...params,
+        runnerIds: lockedRunners.map((runner)=>runner.id),
+        count: lockedRunners.length
+    });
+}
+async function selectIdleRunnerInstancesTx(tx, params) {
+    if (params.runnerIds) {
+        // The first pass retains the row locks. Recheck the full predicate in a fresh statement so
+        // reservation subqueries see the state that was current when the lock was acquired.
+        return await tx.select({
+            id: providerRunners.id,
+            launchKind: providerRunners.launchKind
+        }).from(providerRunners).where(and(inArray(providerRunners.id, params.runnerIds), isBindableRunner(tx, params))).orderBy(asc(providerRunners.createdAt), asc(providerRunners.id)).limit(params.count);
+    }
+    // Select the oldest candidates first, but leave row locking to a second pass. The cleanup
+    // path locks runner rows by id, so the locking pass must use that same order. A skipped row
+    // rolls back only this nested savepoint and retries the bounded candidate scan, allowing the
+    // next-oldest eligible runner to refill the grant without retaining a partial lock set.
+    const retrySelection = Symbol('retry bindable runner selection');
+    const idleRunners = await (async ()=>{
+        while(true){
+            try {
+                return await tx.transaction(async (lockTx)=>{
+                    const candidateRunners = await lockTx.select({
+                        id: providerRunners.id,
+                        launchKind: providerRunners.launchKind
+                    }).from(providerRunners).where(isBindableRunner(lockTx, params)).orderBy(asc(providerRunners.createdAt), asc(providerRunners.id)).limit(params.count);
+                    if (candidateRunners.length === 0) return candidateRunners;
+                    // PostgreSQL can acquire row locks before sorting a multi-row FOR UPDATE result. Lock
+                    // each selected candidate individually so both polling and cleanup use id order.
+                    const lockedRunners = [];
+                    for (const candidate of [
+                        ...candidateRunners
+                    ].sort(compareRunnerIds)){
+                        const [runner] = await lockTx.select({
+                            id: providerRunners.id,
+                            launchKind: providerRunners.launchKind
+                        }).from(providerRunners).where(and(eq(providerRunners.id, candidate.id), isBindableRunner(lockTx, params))).for('update');
+                        if (runner) lockedRunners.push(runner);
+                    }
+                    if (lockedRunners.length !== candidateRunners.length) throw retrySelection;
+                    return lockedRunners;
+                });
+            } catch (error) {
+                if (error !== retrySelection) throw error;
+            }
+        }
+    })();
+    return idleRunners;
+}
+async function bindIdleRunnerInstancesTx(tx, params) {
+    const idleRunnerIds = params.idleRunners.map((runner)=>runner.id);
     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({
+    }).where(and(inArray(runnerActivationTokens.runnerInstanceId, idleRunnerIds), isNull(runnerActivationTokens.consumedAt), isNull(runnerActivationTokens.revokedAt)));
+    // listIdleRunnerInstancesTx re-checks the full eligibility predicate in a fresh statement
+    // after acquiring row locks. Those locks prevent the validated set from changing before this
+    // update runs.
+    const boundRunners = await tx.update(providerRunners).set({
         workspaceId: params.workspaceId,
         reservationId: params.reservationId,
         intendedReservationId: null,
         assignedAt: sql`now()`,
         updatedAt: sql`now()`
-    }).where(and(inArray(providerRunners.id, idleRunners.map((runner)=>runner.id)), canBindRunner()));
+    }).where(and(inArray(providerRunners.id, idleRunnerIds), isBindableRunner(tx, params))).returning({
+        id: providerRunners.id,
+        launchKind: providerRunners.launchKind
+    });
+    if (boundRunners.length !== params.idleRunners.length) throw new Error('Locked idle runner set changed before binding');
+    const reboundCount = boundRunners.filter((runner)=>runner.launchKind === 'demand').length;
+    if (reboundCount > 0) recordProviderRunnerActivationOutcome({
+        outcome: 'rebound',
+        count: reboundCount
+    });
 }
 async function listInstallationDemandWorkspaceIds(eligibleWorkspaceIds) {
     if (eligibleWorkspaceIds.size === 0) return [];
@@ -228,13 +406,27 @@ export async function deleteReservationsByIds(ids) {
 }
 async function deleteReservationsWithCleanupTx(tx, ids, params = {}) {
     if (ids.length === 0) return 0;
-    const affectedRunners = await tx.select({
+    const isAffectedRunner = ()=>or(and(inArray(providerRunners.reservationId, ids), isNull(providerRunners.runnerSessionId), isNull(providerRunners.reservationReleasedAt)), inArray(providerRunners.intendedReservationId, ids));
+    const candidateRunners = await tx.select({
         id: providerRunners.id,
         reservationId: providerRunners.reservationId,
         intendedReservationId: providerRunners.intendedReservationId,
         runnerSessionId: providerRunners.runnerSessionId,
         reservationReleasedAt: providerRunners.reservationReleasedAt
-    }).from(providerRunners).where(or(and(inArray(providerRunners.reservationId, ids), isNull(providerRunners.runnerSessionId), isNull(providerRunners.reservationReleasedAt)), inArray(providerRunners.intendedReservationId, ids))).orderBy(asc(providerRunners.id)).for('update');
+    }).from(providerRunners).where(isAffectedRunner()).orderBy(asc(providerRunners.id));
+    const affectedRunners = [];
+    for (const candidate of [
+        ...candidateRunners
+    ].sort(compareRunnerIds)){
+        const [runner] = await tx.select({
+            id: providerRunners.id,
+            reservationId: providerRunners.reservationId,
+            intendedReservationId: providerRunners.intendedReservationId,
+            runnerSessionId: providerRunners.runnerSessionId,
+            reservationReleasedAt: providerRunners.reservationReleasedAt
+        }).from(providerRunners).where(and(eq(providerRunners.id, candidate.id), isAffectedRunner())).for('update');
+        if (runner) affectedRunners.push(runner);
+    }
     const reservationRows = await tx.select({
         id: reservations.id
     }).from(reservations).where(and(inArray(reservations.id, ids), params.expiredOnly ? lt(reservations.expiresAt, sql`now()`) : undefined)).for('update');
@@ -290,6 +482,9 @@ function deductProvisionerReservations(templates, activeReservations) {
         drawSlots(satisfyingTemplates, reservation.reserved);
     }
 }
+function compareRunnerIds(left, right) {
+    return left.id < right.id ? -1 : left.id > right.id ? 1 : 0;
+}
 function sortDemandRows(rows) {
     return [
         ...rows