@shipfox/api-runners 12.2.0 → 12.3.0

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

@@ -0,0 +1,230 @@
+import {AUTH_PROVISIONER_TOKEN, AUTH_USER, setProvisionerContext} from '@shipfox/api-auth-context';
+import {
+  type AuthMethod,
+  ClientError,
+  closeApp,
+  createApp,
+  extractBearerToken,
+} from '@shipfox/node-fastify';
+import {vi} from '@shipfox/vitest/vi';
+import {eq} from 'drizzle-orm';
+import type {FastifyInstance, FastifyRequest} from 'fastify';
+import {db} from '#db/db.js';
+import {providerRunners} from '#db/schema/runner-instances.js';
+import {runnerSessions} from '#db/schema/runner-sessions.js';
+import {runnerReservationPromotionFailureCount} from '#metrics/instance.js';
+import {
+  createRunnerControlSessionAuthMethod,
+  createRunnerRegistrationTokenAuthMethod,
+} from '#presentation/auth/index.js';
+import {
+  arrangeDeletedRunnerEnrollment,
+  arrangeExpiredRunnerEnrollment,
+  fakeLeaseTokenAuthMethod,
+  fakeRunnerSessionAuthMethod,
+  pendingJobFactory,
+  provisionerTokenFactory,
+  runnersTestAuthClient,
+} from '#test/index.js';
+import {createRunnerRoutes} from './index.js';
+
+let workspaceToken: string;
+const fakeUserAuth: AuthMethod = {name: AUTH_USER, authenticate: () => Promise.resolve()};
+
+describe('late runner enrollment recovery', () => {
+  let app: FastifyInstance;
+  let provisionerId: string;
+  let workspaceId: string;
+
+  const provisionerAuth: AuthMethod = {
+    name: AUTH_PROVISIONER_TOKEN,
+    authenticate: (request: FastifyRequest) => {
+      if (extractBearerToken(request.headers.authorization) !== workspaceToken)
+        throw new ClientError('Invalid provisioner token', 'unauthorized', {status: 401});
+      setProvisionerContext(request, {
+        scope: 'workspace',
+        workspaceId,
+        provisionerTokenId: provisionerId,
+      });
+      return Promise.resolve();
+    },
+  };
+
+  beforeAll(async () => {
+    app = await createApp({
+      auth: [
+        fakeUserAuth,
+        provisionerAuth,
+        createRunnerRegistrationTokenAuthMethod(),
+        createRunnerControlSessionAuthMethod(),
+        fakeRunnerSessionAuthMethod,
+        fakeLeaseTokenAuthMethod,
+      ],
+      routes: createRunnerRoutes(runnersTestAuthClient),
+      swagger: false,
+    });
+    await app.ready();
+  });
+
+  afterAll(async () => {
+    await closeApp();
+  });
+
+  beforeEach(async () => {
+    workspaceId = crypto.randomUUID();
+    workspaceToken = `late-runner-workspace-provisioner-${crypto.randomUUID()}`;
+    const provisioner = await provisionerTokenFactory.create(
+      {scope: 'workspace', workspaceId},
+      {transient: {rawToken: workspaceToken}},
+    );
+    provisionerId = provisioner.id;
+  });
+
+  it('keeps an enrolled runner recoverable when its reservation expires first', async () => {
+    const arrangement = await arrangeExpiredRunnerEnrollment({provisionerId, workspaceId});
+    const failureSpy = vi.spyOn(runnerReservationPromotionFailureCount, 'add');
+
+    try {
+      await enrollRunner(arrangement.controlSessionToken);
+
+      expect(failureSpy).toHaveBeenCalledWith(1, {reason: 'reservation-expired'});
+    } finally {
+      failureSpy.mockRestore();
+    }
+
+    const [runner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, arrangement.runnerInstanceId));
+    expect(runner).toMatchObject({
+      workspaceId: null,
+      reservationId: null,
+      intendedReservationId: arrangement.reservationId,
+      assignedAt: null,
+      runnerSessionId: null,
+      state: 'running',
+      providerRunnerId: expect.any(String),
+    });
+  });
+
+  it('keeps an enrolled runner recoverable when its reservation was swept away', async () => {
+    const arrangement = await arrangeDeletedRunnerEnrollment({provisionerId, workspaceId});
+    const failureSpy = vi.spyOn(runnerReservationPromotionFailureCount, 'add');
+
+    try {
+      await enrollRunner(arrangement.controlSessionToken);
+
+      expect(failureSpy).toHaveBeenCalledWith(1, {reason: 'reservation-not-found'});
+    } finally {
+      failureSpy.mockRestore();
+    }
+
+    const [runner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, arrangement.runnerInstanceId));
+    expect(runner).toMatchObject({
+      workspaceId: null,
+      reservationId: null,
+      intendedReservationId: arrangement.reservationId,
+      assignedAt: null,
+      runnerSessionId: null,
+      state: 'running',
+      providerRunnerId: expect.any(String),
+    });
+  });
+
+  it('rebinds a stranded runner through activation and claims its pending job', async () => {
+    const arrangement = await arrangeExpiredRunnerEnrollment({provisionerId, workspaceId});
+    await enrollRunner(arrangement.controlSessionToken);
+    const pendingJob = 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);
+    const reboundReservationId = demand.json().reservations[0]?.reservation_id;
+    expect(reboundReservationId).toEqual(expect.any(String));
+
+    const [reboundRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, arrangement.runnerInstanceId));
+    expect(reboundRunner).toMatchObject({
+      workspaceId,
+      reservationId: reboundReservationId,
+      intendedReservationId: null,
+      assignedAt: expect.any(Date),
+      runnerSessionId: null,
+      state: 'running',
+    });
+
+    const assignment = await app.inject({
+      method: 'GET',
+      url: '/runner-control/assignment?wait_seconds=1',
+      headers: {authorization: `Bearer ${arrangement.controlSessionToken}`},
+    });
+    expect(assignment.statusCode).toBe(200);
+    const activationToken = assignment.json().activation_token;
+    expect(activationToken).toEqual(expect.any(String));
+
+    const registered = await app.inject({
+      method: 'POST',
+      url: '/runners/register',
+      headers: {authorization: `Bearer ${activationToken}`},
+      payload: {labels: ['linux']},
+    });
+    expect(registered.statusCode).toBe(200);
+    expect(registered.json()).toMatchObject({mode: 'activation', max_claims: 1});
+
+    const claimed = await app.inject({
+      method: 'POST',
+      url: '/runners/jobs/request',
+      headers: {authorization: `Bearer ${registered.json().session_token}`},
+    });
+    expect(claimed.statusCode).toBe(200);
+    expect(claimed.json()).toMatchObject({job_id: pendingJob.jobId});
+
+    const [session] = await db()
+      .select()
+      .from(runnerSessions)
+      .where(eq(runnerSessions.id, registered.json().session_id));
+    expect(session).toMatchObject({
+      registrationTokenKind: 'activation',
+      runnerInstanceId: arrangement.runnerInstanceId,
+      provisionerId,
+      providerRunnerId: reboundRunner?.providerRunnerId,
+      maxClaims: 1,
+      claimsUsed: 1,
+    });
+  });
+
+  async function enrollRunner(controlSessionToken: string): Promise<void> {
+    const enrolled = await app.inject({
+      method: 'POST',
+      url: '/runner-control/enrollment',
+      headers: {authorization: `Bearer ${controlSessionToken}`},
+      payload: {labels: ['linux'], provider_kind: 'docker', protocol_version: '1'},
+    });
+
+    expect(enrolled.statusCode).toBe(200);
+    expect(enrolled.json()).toEqual({activation_token: null});
+  }
+});

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

@@ -116,13 +116,30 @@ describe('POST /provisioners/demand/poll', () => {
       method: 'POST',
       url: '/provisioners/demand/poll',
       headers: {authorization: `Bearer ${VALID_PROVISIONER_TOKEN}`},
-      payload: body({max_reservations: 1, reservation_ttl_seconds: 600}),
+      payload: body({max_reservations: 1, reservation_ttl_seconds: 601}),
     });
 
     expect(res.statusCode).toBe(200);
     const expiresAt = Date.parse(res.json().reservations[0].expires_at);
-    expect(expiresAt).toBeGreaterThanOrEqual(requestStartedAt + 295_000);
-    expect(expiresAt).toBeLessThan(requestStartedAt + 305_000);
+    expect(expiresAt).toBeGreaterThanOrEqual(requestStartedAt + 595_000);
+    expect(expiresAt).toBeLessThan(requestStartedAt + 605_000);
+  });
+
+  it('accepts a slow-boot reservation TTL below the raised ceiling', async () => {
+    await pendingJobFactory.create({workspaceId, requiredLabels: ['linux']});
+
+    const requestStartedAt = Date.now();
+    const res = await app.inject({
+      method: 'POST',
+      url: '/provisioners/demand/poll',
+      headers: {authorization: `Bearer ${VALID_PROVISIONER_TOKEN}`},
+      payload: body({max_reservations: 1, reservation_ttl_seconds: 540}),
+    });
+
+    expect(res.statusCode).toBe(200);
+    const expiresAt = Date.parse(res.json().reservations[0].expires_at);
+    expect(expiresAt).toBeGreaterThanOrEqual(requestStartedAt + 535_000);
+    expect(expiresAt).toBeLessThan(requestStartedAt + 545_000);
   });
 
   it('applies a requested reservation TTL below the ceiling', async () => {
@@ -528,13 +545,13 @@ describe('POST /provisioners/demand/poll with installation provisioning configur
       method: 'POST',
       url: '/provisioners/demand/poll',
       headers: {authorization: `Bearer ${INSTALLATION_PROVISIONER_TOKEN}`},
-      payload: body({max_reservations: 1, reservation_ttl_seconds: 600}),
+      payload: body({max_reservations: 1, reservation_ttl_seconds: 601}),
     });
 
     expect(res.statusCode).toBe(200);
     const expiresAt = Date.parse(res.json().reservations[0].expires_at);
-    expect(expiresAt).toBeGreaterThanOrEqual(requestStartedAt + 295_000);
-    expect(expiresAt).toBeLessThan(requestStartedAt + 305_000);
+    expect(expiresAt).toBeGreaterThanOrEqual(requestStartedAt + 595_000);
+    expect(expiresAt).toBeLessThan(requestStartedAt + 605_000);
   });
 
   function body(params: {max_reservations: number; reservation_ttl_seconds?: number}) {

package/test/fixtures/late-runner-enrollment.ts

@@ -0,0 +1,103 @@
+import {extractDisplayPrefix, generateOpaqueToken, hashOpaqueToken} from '@shipfox/node-tokens';
+import {db} from '#db/db.js';
+import {deleteReservationsByIds} from '#db/reservations.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';
+
+export interface LateRunnerEnrollmentArrangement {
+  controlSessionToken: string;
+  provisionerId: string;
+  reservationId: string;
+  runnerInstanceId: string;
+  workspaceId: string;
+}
+
+export async function arrangeExpiredRunnerEnrollment(params: {
+  provisionerId: string;
+  workspaceId: string;
+}): Promise<LateRunnerEnrollmentArrangement> {
+  const reservationId = await insertReservation({
+    ...params,
+    expiresAt: new Date(Date.now() - 1_000),
+  });
+
+  return await createRunnerEnrollmentArrangement({...params, reservationId});
+}
+
+export async function arrangeDeletedRunnerEnrollment(params: {
+  provisionerId: string;
+  workspaceId: string;
+}): Promise<LateRunnerEnrollmentArrangement> {
+  const reservationId = await insertReservation({
+    ...params,
+    expiresAt: new Date(Date.now() + 60_000),
+  });
+  const deleted = await deleteReservationsByIds([reservationId]);
+  if (deleted !== 1) throw new Error('Expected reservation cleanup to delete one row');
+
+  // The provider can still enroll from a launch payload carrying the deleted reservation id.
+  return await createRunnerEnrollmentArrangement({...params, reservationId});
+}
+
+async function insertReservation(params: {
+  expiresAt: Date;
+  provisionerId: string;
+  workspaceId: string;
+}): Promise<string> {
+  const [reservation] = await db()
+    .insert(reservations)
+    .values({
+      workspaceId: params.workspaceId,
+      provisionerId: params.provisionerId,
+      requiredLabels: ['linux'],
+      count: 1,
+      expiresAt: params.expiresAt,
+    })
+    .returning({id: reservations.id});
+  if (!reservation) throw new Error('Expected reservation');
+
+  return reservation.id;
+}
+
+async function createRunnerEnrollmentArrangement(params: {
+  provisionerId: string;
+  reservationId: string;
+  workspaceId: string;
+}): Promise<LateRunnerEnrollmentArrangement> {
+  const providerRunnerId = `late-runner-${crypto.randomUUID()}`;
+  const controlSessionToken = generateOpaqueToken('runnerControlSession');
+  const [runner] = await db()
+    .insert(providerRunners)
+    .values({
+      provisionerId: params.provisionerId,
+      providerRunnerId,
+      intendedReservationId: params.reservationId,
+      templateKey: 'linux',
+      labels: ['linux'],
+      state: 'starting',
+      providerKind: 'docker',
+      protocolVersion: '1',
+      reportedAt: new Date(),
+    })
+    .returning({id: providerRunners.id});
+  if (!runner) throw new Error('Expected runner instance');
+
+  await db()
+    .insert(runnerControlSessions)
+    .values({
+      runnerInstanceId: runner.id,
+      provisionerId: params.provisionerId,
+      hashedToken: hashOpaqueToken(controlSessionToken),
+      prefix: extractDisplayPrefix(controlSessionToken),
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+
+  return {
+    controlSessionToken,
+    provisionerId: params.provisionerId,
+    reservationId: params.reservationId,
+    runnerInstanceId: runner.id,
+    workspaceId: params.workspaceId,
+  };
+}

package/test/index.ts

@@ -23,3 +23,8 @@ export {
   mintRunnerSessionToken,
 } from './fixtures/auth.js';
 export {runnersTestAuthClient} from './fixtures/auth-inter-module.js';
+export {
+  arrangeDeletedRunnerEnrollment,
+  arrangeExpiredRunnerEnrollment,
+  type LateRunnerEnrollmentArrangement,
+} from './fixtures/late-runner-enrollment.js';