@shipfox/api-runners 12.3.0 → 12.4.0

package/src/presentation/routes/late-runner-enrollment.test.ts

@@ -7,9 +7,11 @@ import {
   extractBearerToken,
 } from '@shipfox/node-fastify';
 import {vi} from '@shipfox/vitest/vi';
-import {eq} from 'drizzle-orm';
+import {and, eq} from 'drizzle-orm';
 import type {FastifyInstance, FastifyRequest} from 'fastify';
 import {db} from '#db/db.js';
+import {reservations} from '#db/schema/reservations.js';
+import {runnerControlSessions} from '#db/schema/runner-control-sessions.js';
 import {providerRunners} from '#db/schema/runner-instances.js';
 import {runnerSessions} from '#db/schema/runner-sessions.js';
 import {runnerReservationPromotionFailureCount} from '#metrics/instance.js';
@@ -160,7 +162,22 @@ describe('late runner enrollment recovery', () => {
     });
 
     expect(demand.statusCode).toBe(200);
-    const reboundReservationId = demand.json().reservations[0]?.reservation_id;
+    expect(demand.json().reservations).toEqual([]);
+    const reboundReservations = await db()
+      .select()
+      .from(reservations)
+      .where(
+        and(
+          eq(reservations.workspaceId, workspaceId),
+          eq(reservations.provisionerId, provisionerId),
+          eq(reservations.kind, 'bound'),
+        ),
+      );
+    const reboundReservation = reboundReservations.find(
+      (reservation) => reservation.id !== arrangement.reservationId,
+    );
+    const reboundReservationId = reboundReservation?.id;
+    if (!reboundReservationId) throw new Error('Expected rebound reservation');
     expect(reboundReservationId).toEqual(expect.any(String));
 
     const [reboundRunner] = await db()
@@ -216,6 +233,206 @@ describe('late runner enrollment recovery', () => {
     });
   });
 
+  it('preserves a rebound assignment across stale provider reports and releases it once', async () => {
+    const providerRunnerId = `stale-assignment-${crypto.randomUUID()}`;
+    const firstReportAt = new Date(Date.now() - 2_000);
+    const [staleReservation] = await db()
+      .insert(reservations)
+      .values({
+        workspaceId,
+        provisionerId,
+        requiredLabels: ['linux'],
+        count: 1,
+        expiresAt: new Date(Date.now() + 60_000),
+      })
+      .returning({id: reservations.id});
+    if (!staleReservation) throw new Error('Expected stale reservation');
+
+    const initialReport = await app.inject({
+      method: 'POST',
+      url: '/provisioners/runner-instances/report',
+      headers: {authorization: `Bearer ${workspaceToken}`},
+      payload: {
+        events: [
+          {
+            provider_runner_id: providerRunnerId,
+            reservation_id: staleReservation.id,
+            template_key: 'linux',
+            labels: ['linux'],
+            state: 'running',
+            reported_at: firstReportAt.toISOString(),
+            provider_kind: 'docker',
+          },
+        ],
+      },
+    });
+
+    expect(initialReport.statusCode).toBe(200);
+    expect(initialReport.json()).toEqual({accepted: 1, reservations_released: 0});
+
+    const [reportedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(
+        and(
+          eq(providerRunners.provisionerId, provisionerId),
+          eq(providerRunners.providerRunnerId, providerRunnerId),
+        ),
+      );
+    if (!reportedRunner) throw new Error('Expected reported runner');
+    expect(reportedRunner).toMatchObject({
+      workspaceId,
+      reservationId: staleReservation.id,
+      state: 'running',
+    });
+
+    await db()
+      .insert(runnerControlSessions)
+      .values({
+        runnerInstanceId: reportedRunner.id,
+        provisionerId,
+        hashedToken: crypto.randomUUID(),
+        prefix: 'test',
+        expiresAt: new Date(Date.now() + 60_000),
+      });
+    await db()
+      .update(reservations)
+      .set({expiresAt: new Date(Date.now() - 1_000)})
+      .where(eq(reservations.id, staleReservation.id));
+    await pendingJobFactory.create({workspaceId, requiredLabels: ['linux']});
+
+    const demand = await app.inject({
+      method: 'POST',
+      url: '/provisioners/demand/poll',
+      headers: {authorization: `Bearer ${workspaceToken}`},
+      payload: {
+        wait_seconds: 0,
+        max_reservations: 1,
+        reservation_ttl_seconds: 60,
+        templates: [
+          {
+            template_key: 'linux',
+            labels: ['linux'],
+            available_slots: 1,
+            starting: 0,
+            running: 1,
+          },
+        ],
+      },
+    });
+
+    expect(demand.statusCode).toBe(200);
+    expect(demand.json().reservations).toEqual([]);
+    const reboundReservations = await db()
+      .select()
+      .from(reservations)
+      .where(
+        and(
+          eq(reservations.workspaceId, workspaceId),
+          eq(reservations.provisionerId, provisionerId),
+          eq(reservations.kind, 'bound'),
+        ),
+      );
+    const reboundReservation = reboundReservations.find(
+      (reservation) => reservation.id !== staleReservation.id,
+    );
+    const reboundReservationId = reboundReservation?.id;
+    if (!reboundReservationId) throw new Error('Expected rebound reservation');
+    expect(reboundReservationId).toEqual(expect.any(String));
+    expect(reboundReservationId).not.toBe(staleReservation.id);
+
+    const staleReport = await app.inject({
+      method: 'POST',
+      url: '/provisioners/runner-instances/report',
+      headers: {authorization: `Bearer ${workspaceToken}`},
+      payload: {
+        events: [
+          {
+            provider_runner_id: providerRunnerId,
+            reservation_id: staleReservation.id,
+            template_key: 'linux',
+            labels: ['linux'],
+            state: 'running',
+            reported_at: new Date(firstReportAt.getTime() + 1_000).toISOString(),
+            provider_kind: 'docker',
+          },
+        ],
+      },
+    });
+
+    expect(staleReport.statusCode).toBe(200);
+    expect(staleReport.json()).toEqual({accepted: 1, reservations_released: 0});
+
+    const [reboundRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, reportedRunner.id));
+    expect(reboundRunner).toMatchObject({
+      workspaceId,
+      reservationId: reboundReservationId,
+      state: 'running',
+      intendedReservationId: null,
+    });
+
+    const terminalReport = await app.inject({
+      method: 'POST',
+      url: '/provisioners/runner-instances/report',
+      headers: {authorization: `Bearer ${workspaceToken}`},
+      payload: {
+        events: [
+          {
+            provider_runner_id: providerRunnerId,
+            reservation_id: staleReservation.id,
+            template_key: 'linux',
+            labels: ['linux'],
+            state: 'terminated',
+            reported_at: new Date(firstReportAt.getTime() + 2_000).toISOString(),
+            provider_kind: 'docker',
+          },
+        ],
+      },
+    });
+
+    expect(terminalReport.statusCode).toBe(200);
+    expect(terminalReport.json()).toEqual({accepted: 1, reservations_released: 1});
+
+    const [staleReservationAfterTerminal] = await db()
+      .select()
+      .from(reservations)
+      .where(eq(reservations.id, staleReservation.id));
+    const [reboundReservationAfterTerminal] = await db()
+      .select()
+      .from(reservations)
+      .where(eq(reservations.id, reboundReservationId));
+    expect(staleReservationAfterTerminal).toMatchObject({
+      id: staleReservation.id,
+      count: 1,
+    });
+    expect(reboundReservationAfterTerminal).toBeUndefined();
+
+    const repeatedTerminalReport = await app.inject({
+      method: 'POST',
+      url: '/provisioners/runner-instances/report',
+      headers: {authorization: `Bearer ${workspaceToken}`},
+      payload: {
+        events: [
+          {
+            provider_runner_id: providerRunnerId,
+            reservation_id: staleReservation.id,
+            template_key: 'linux',
+            labels: ['linux'],
+            state: 'terminated',
+            reported_at: new Date(firstReportAt.getTime() + 3_000).toISOString(),
+            provider_kind: 'docker',
+          },
+        ],
+      },
+    });
+
+    expect(repeatedTerminalReport.statusCode).toBe(200);
+    expect(repeatedTerminalReport.json()).toEqual({accepted: 1, reservations_released: 0});
+  });
+
   async function enrollRunner(controlSessionToken: string): Promise<void> {
     const enrolled = await app.inject({
       method: 'POST',

package/src/presentation/routes/poll-demand.test.ts

@@ -20,6 +20,7 @@ import {db} from '#db/db.js';
 import {provisionerCapabilitySnapshots} from '#db/schema/provisioner-capability-snapshots.js';
 import {runningJobExecutions} from '#db/schema/running-job-executions.js';
 import {
+  providerRunnerActivationOutcomeCount,
   providerRunnerCountDivergenceCount,
   providerRunnerTerminateIntentIssuedCount,
 } from '#metrics/instance.js';
@@ -99,6 +100,7 @@ describe('POST /provisioners/demand/poll', () => {
     expect(res.json()).toMatchObject({
       stats: [{labels: ['linux'], queued: 1, reserved: 1}],
       reservations: [{labels: ['linux'], count: 1}],
+      newly_reserved_count: 1,
       terminate_provider_runner_ids: [],
     });
     expect(res.json().reservations[0].reservation_id).toEqual(expect.any(String));
@@ -367,6 +369,60 @@ describe('POST /provisioners/demand/poll', () => {
     expect(intentCalls).toHaveLength(1);
   });
 
+  it('counts activation-timeout reaps once across terminate-intent retries', async () => {
+    const intentSpy = vi.spyOn(providerRunnerTerminateIntentIssuedCount, 'add');
+    const outcomeSpy = vi.spyOn(providerRunnerActivationOutcomeCount, 'add');
+    await providerRunnerFactory.create({
+      workspaceId,
+      provisionerId: provisionerTokenId,
+      providerRunnerId: 'stale-demand-runner',
+      launchKind: 'demand',
+      createdAt: new Date(Date.now() - 301_000),
+      templateKey: 'linux',
+      state: 'running',
+    });
+    const intentCallsBefore = intentSpy.mock.calls.length;
+    const outcomeCallsBefore = outcomeSpy.mock.calls.length;
+
+    const poll = () =>
+      app.inject({
+        method: 'POST',
+        url: '/provisioners/demand/poll',
+        headers: {authorization: `Bearer ${VALID_PROVISIONER_TOKEN}`},
+        payload: body({max_reservations: 0}),
+      });
+    const firstResponse = await poll();
+    const retryResponse = await poll();
+
+    expect(firstResponse.statusCode).toBe(200);
+    expect(retryResponse.statusCode).toBe(200);
+    expect(firstResponse.json()).toMatchObject({
+      reservations: [],
+      terminate_provider_runner_ids: ['stale-demand-runner'],
+    });
+    expect(retryResponse.json()).toMatchObject({
+      reservations: [],
+      terminate_provider_runner_ids: ['stale-demand-runner'],
+    });
+    const intentCalls = intentSpy.mock.calls
+      .slice(intentCallsBefore)
+      .filter(
+        ([value, attributes]) =>
+          value === 1 &&
+          JSON.stringify(attributes) ===
+            JSON.stringify({surface: 'poll-demand', reason: 'activation-timeout'}),
+      );
+    expect(intentCalls).toHaveLength(2);
+    expect(
+      outcomeSpy.mock.calls
+        .slice(outcomeCallsBefore)
+        .filter(
+          ([value, attributes]) =>
+            value === 1 && JSON.stringify(attributes) === JSON.stringify({outcome: 'reaped'}),
+        ),
+    ).toHaveLength(1);
+  });
+
   it('returns 400 for max reservations above the request bound', async () => {
     const res = await app.inject({
       method: 'POST',

package/src/presentation/routes/poll-demand.ts

@@ -98,6 +98,7 @@ export function createPollDemandRoute(options: CreateRunnersModuleOptions = {})
           provisionerId: provisionerContext.provisionerTokenId,
           maxReservations: request.body.max_reservations,
           ttlSeconds,
+          activationGraceSeconds: config.RESERVATION_TTL_SECONDS,
           templates,
           capabilityWindowSeconds: config.PROVISIONER_ACTIVE_WINDOW_SECONDS,
           eligibleWorkspaceIds,

package/src/presentation/routes/reconcile-runner-instances.test.ts

@@ -6,6 +6,10 @@ import {
   AUTH_USER,
   setProvisionerContext,
 } from '@shipfox/api-auth-context';
+import {
+  RECONCILE_RUNNER_INSTANCES_INTENDED_RESERVATION_HEADER,
+  RECONCILE_RUNNER_INSTANCES_INTENDED_RESERVATION_HEADER_VALUE,
+} from '@shipfox/api-runners-dto';
 import {
   type AuthMethod,
   ClientError,
@@ -86,7 +90,11 @@ describe('POST /provisioners/runner-instances/reconcile', () => {
     const jobId = crypto.randomUUID();
     const workflowRunId = crypto.randomUUID();
     const workflowRunAttemptId = crypto.randomUUID();
-    await createRunnerInstance({providerRunnerId: 'provisioned-runner-1'});
+    const intendedReservationId = crypto.randomUUID();
+    await createRunnerInstance({
+      providerRunnerId: 'provisioned-runner-1',
+      intendedReservationId,
+    });
     await insertRunningJob({
       jobId,
       workflowRunId,
@@ -98,7 +106,11 @@ describe('POST /provisioners/runner-instances/reconcile', () => {
     const res = await app.inject({
       method: 'POST',
       url: '/provisioners/runner-instances/reconcile',
-      headers: {authorization: `Bearer ${VALID_PROVISIONER_TOKEN}`},
+      headers: {
+        authorization: `Bearer ${VALID_PROVISIONER_TOKEN}`,
+        [RECONCILE_RUNNER_INSTANCES_INTENDED_RESERVATION_HEADER]:
+          RECONCILE_RUNNER_INSTANCES_INTENDED_RESERVATION_HEADER_VALUE,
+      },
       payload: {observed_provider_runner_ids: ['provisioned-runner-1']},
     });
 
@@ -108,6 +120,7 @@ describe('POST /provisioners/runner-instances/reconcile', () => {
         {
           provider_runner_id: 'provisioned-runner-1',
           state: 'running',
+          intended_reservation_id: intendedReservationId,
           reservation_id: null,
           runner_session_id: null,
           bound_job: {
@@ -123,6 +136,23 @@ describe('POST /provisioners/runner-instances/reconcile', () => {
     });
   });
 
+  it('omits the additive reservation field for legacy provisioners', async () => {
+    await createRunnerInstance({
+      providerRunnerId: 'provisioned-runner-1',
+      intendedReservationId: crypto.randomUUID(),
+    });
+
+    const res = await app.inject({
+      method: 'POST',
+      url: '/provisioners/runner-instances/reconcile',
+      headers: {authorization: `Bearer ${VALID_PROVISIONER_TOKEN}`},
+      payload: {observed_provider_runner_ids: ['provisioned-runner-1']},
+    });
+
+    expect(res.statusCode).toBe(200);
+    expect(Object.hasOwn(res.json().runners[0], 'intended_reservation_id')).toBe(false);
+  });
+
   it('returns terminate for an observed terminal runner', async () => {
     await createRunnerInstance({
       providerRunnerId: 'provisioned-runner-1',
@@ -340,6 +370,7 @@ describe('POST /provisioners/runner-instances/reconcile', () => {
   async function createRunnerInstance(params: {
     providerRunnerId: string;
     state?: 'starting' | 'running' | 'stopping' | 'stopped' | 'failed' | 'terminated';
+    intendedReservationId?: string | null;
     reservationId?: string | null;
     reportedAt?: Date;
   }) {
@@ -347,6 +378,7 @@ describe('POST /provisioners/runner-instances/reconcile', () => {
       workspaceId,
       provisionerId: provisionerTokenId,
       providerRunnerId: params.providerRunnerId,
+      intendedReservationId: params.intendedReservationId ?? null,
       reservationId: params.reservationId ?? null,
       state: params.state ?? 'running',
       reportedAt: params.reportedAt ?? new Date(),

package/src/presentation/routes/reconcile-runner-instances.ts

@@ -1,5 +1,7 @@
 import {requireProvisionerContext} from '@shipfox/api-auth-context';
 import {
+  RECONCILE_RUNNER_INSTANCES_INTENDED_RESERVATION_HEADER,
+  RECONCILE_RUNNER_INSTANCES_INTENDED_RESERVATION_HEADER_VALUE,
   reconcileRunnerInstancesBodySchema,
   reconcileRunnerInstancesResponseSchema,
 } from '@shipfox/api-runners-dto';
@@ -25,6 +27,10 @@ export const reconcileRunnerInstancesRoute = defineRoute({
       observedRunnerInstanceIds: request.body.observed_provider_runner_ids,
     });
 
-    return toReconcileRunnerInstancesResponseDto(result);
+    return toReconcileRunnerInstancesResponseDto(result, {
+      includeIntendedReservationId:
+        request.headers[RECONCILE_RUNNER_INSTANCES_INTENDED_RESERVATION_HEADER] ===
+        RECONCILE_RUNNER_INSTANCES_INTENDED_RESERVATION_HEADER_VALUE,
+    });
   },
 });