@shipfox/api-runners 12.3.0 → 12.4.1

package/src/db/ephemeral-registration-tokens.test.ts

@@ -3,6 +3,7 @@ import {and, eq, inArray} from 'drizzle-orm';
 import {
   ActiveEphemeralRegistrationTokensExistError,
   type RegistrationTokenBatchExceedsReservationError,
+  ReservationNotFoundError,
 } from '#core/errors.js';
 import {db} from '#db/db.js';
 import {
@@ -56,6 +57,20 @@ describe('createEphemeralRegistrationTokensBatch', () => {
     ).toEqual(expect.objectContaining({providerRunnerId: 'provisioned-runner-a'}));
   });
 
+  it('rejects bound reservations', async () => {
+    const boundReservationId = await createReservation({count: 1, kind: 'bound'});
+
+    await expect(
+      createEphemeralRegistrationTokensBatch({
+        workspaceId,
+        provisionerId,
+        reservationId: boundReservationId,
+        expiresAt: new Date(Date.now() + 300_000),
+        rows: [row('bound-runner', generateOpaqueToken('ephemeralRegistrationToken'))],
+      }),
+    ).rejects.toBeInstanceOf(ReservationNotFoundError);
+  });
+
   it('rejects the batch and inserts no rows when any provisioned runner already has an active token', async () => {
     await ephemeralRegistrationTokenFactory.create({
       workspaceId,
@@ -186,7 +201,10 @@ describe('createEphemeralRegistrationTokensBatch', () => {
     } satisfies Partial<RegistrationTokenBatchExceedsReservationError>);
   });
 
-  async function createReservation(params: {count: number}): Promise<string> {
+  async function createReservation(params: {
+    count: number;
+    kind?: 'bound' | 'launch';
+  }): Promise<string> {
     const [reservation] = await db()
       .insert(reservations)
       .values({
@@ -194,6 +212,7 @@ describe('createEphemeralRegistrationTokensBatch', () => {
         provisionerId,
         requiredLabels: ['linux'],
         count: params.count,
+        kind: params.kind ?? 'launch',
         expiresAt: new Date(Date.now() + 60_000),
       })
       .returning({id: reservations.id});

package/src/db/ephemeral-registration-tokens.ts

@@ -133,6 +133,7 @@ export async function createEphemeralRegistrationTokensBatch(
           eq(reservations.id, params.reservationId),
           eq(reservations.workspaceId, params.workspaceId),
           eq(reservations.provisionerId, params.provisionerId),
+          eq(reservations.kind, 'launch'),
         ),
       )
       .limit(1);

package/src/db/job-executions.test.ts

@@ -8,6 +8,7 @@ import {eq, sql} from 'drizzle-orm';
 import {EmptyRequiredLabelsError, RunnerSessionExhaustedError} from '#core/errors.js';
 import {claimJobExecution} from '#core/job-executions.js';
 import {detectAndExpireStuckJobs} from '#core/maintenance.js';
+import * as runnerMetrics from '#metrics/instance.js';
 import {
   getLeaseTokenClaims,
   pendingJobFactory,
@@ -242,20 +243,95 @@ describe('claimPendingJobExecution', () => {
     expect(new Date(payload.claimedAt).getTime()).toBe(running?.startedAt.getTime());
   });
 
+  it('records queue time from the pending row creation to the runner claim', async () => {
+    const created = await pendingJobFactory.create({workspaceId});
+    await db()
+      .update(pendingJobExecutions)
+      .set({createdAt: sql`now() - interval '60 seconds'`})
+      .where(eq(pendingJobExecutions.jobExecutionId, created.jobExecutionId));
+    const [pending] = await db()
+      .select({createdAt: pendingJobExecutions.createdAt})
+      .from(pendingJobExecutions)
+      .where(eq(pendingJobExecutions.jobExecutionId, created.jobExecutionId));
+    const recordQueueTime = vi.spyOn(runnerMetrics, 'recordJobExecutionQueueTime');
+
+    try {
+      await claimPendingJobExecution({workspaceId, runnerSessionId, maxClaims: null});
+
+      const [running] = await db()
+        .select({startedAt: runningJobExecutions.startedAt})
+        .from(runningJobExecutions)
+        .where(eq(runningJobExecutions.jobExecutionId, created.jobExecutionId));
+      const observation = recordQueueTime.mock.calls[0]?.[0];
+      if (!pending || !running) throw new Error('Expected pending and running job rows');
+      expect(recordQueueTime).toHaveBeenCalledTimes(1);
+      expect(observation).toMatchObject({provider: null, launchKind: 'manual'});
+      expect(observation?.durationMilliseconds).toBe(
+        running.startedAt.getTime() - pending.createdAt.getTime(),
+      );
+      expect(observation?.durationMilliseconds).toBeGreaterThanOrEqual(60_000);
+      expect(observation?.durationMilliseconds).toBeLessThan(120_000);
+    } finally {
+      recordQueueTime.mockRestore();
+    }
+  });
+
+  it('records bounded queue labels for a provisioned runner claim', async () => {
+    const provisionerId = crypto.randomUUID();
+    const providerRunnerId = `provisioned-runner-${crypto.randomUUID()}`;
+    await db()
+      .update(runnerSessions)
+      .set({registrationTokenKind: 'ephemeral', maxClaims: 1, provisionerId, providerRunnerId})
+      .where(eq(runnerSessions.id, runnerSessionId));
+    await db().insert(providerRunners).values({
+      workspaceId,
+      provisionerId,
+      providerRunnerId,
+      providerKind: 'ec2',
+      launchKind: 'demand',
+      labels: sessionLabels,
+      state: 'running',
+      reportedAt: new Date(),
+    });
+    await pendingJobFactory.create({workspaceId});
+    const recordQueueTime = vi.spyOn(runnerMetrics, 'recordJobExecutionQueueTime');
+
+    try {
+      await claimPendingJobExecution({workspaceId, runnerSessionId, maxClaims: 1});
+
+      expect(recordQueueTime).toHaveBeenCalledWith(
+        expect.objectContaining({provider: 'ec2', launchKind: 'demand'}),
+      );
+      expect(recordQueueTime.mock.calls[0]?.[0].durationMilliseconds).toBeGreaterThanOrEqual(0);
+    } finally {
+      recordQueueTime.mockRestore();
+    }
+  });
+
   it('emits no claimed event when there is nothing to claim', async () => {
+    const recordQueueTime = vi.spyOn(runnerMetrics, 'recordJobExecutionQueueTime');
     const before = await db()
       .select()
       .from(runnersOutbox)
       .where(eq(runnersOutbox.eventType, RUNNER_JOB_CLAIMED));
 
-    const claimed = await claimPendingJobExecution({workspaceId, runnerSessionId, maxClaims: null});
+    try {
+      const claimed = await claimPendingJobExecution({
+        workspaceId,
+        runnerSessionId,
+        maxClaims: null,
+      });
 
-    const after = await db()
-      .select()
-      .from(runnersOutbox)
-      .where(eq(runnersOutbox.eventType, RUNNER_JOB_CLAIMED));
-    expect(claimed).toBeNull();
-    expect(after).toHaveLength(before.length);
+      const after = await db()
+        .select()
+        .from(runnersOutbox)
+        .where(eq(runnersOutbox.eventType, RUNNER_JOB_CLAIMED));
+      expect(claimed).toBeNull();
+      expect(after).toHaveLength(before.length);
+      expect(recordQueueTime).not.toHaveBeenCalled();
+    } finally {
+      recordQueueTime.mockRestore();
+    }
   });
 
   it('emits no claimed event when dropping an orphan pending row', async () => {
@@ -273,13 +349,23 @@ describe('claimPendingJobExecution', () => {
     });
     // Clear the initial claim's events so this assertion only covers the orphan claim.
     const beforeOrphanClaim = await outboxEventsForJob(RUNNER_JOB_CLAIMED, created.jobId);
+    const recordQueueTime = vi.spyOn(runnerMetrics, 'recordJobExecutionQueueTime');
 
-    const second = await claimPendingJobExecution({workspaceId, runnerSessionId, maxClaims: null});
+    try {
+      const second = await claimPendingJobExecution({
+        workspaceId,
+        runnerSessionId,
+        maxClaims: null,
+      });
 
-    expect(second).toBeNull();
-    expect(await outboxEventsForJob(RUNNER_JOB_CLAIMED, created.jobId)).toHaveLength(
-      beforeOrphanClaim.length,
-    );
+      expect(second).toBeNull();
+      expect(await outboxEventsForJob(RUNNER_JOB_CLAIMED, created.jobId)).toHaveLength(
+        beforeOrphanClaim.length,
+      );
+      expect(recordQueueTime).not.toHaveBeenCalled();
+    } finally {
+      recordQueueTime.mockRestore();
+    }
   });
 
   it('returns the job ids when a job is available', async () => {

package/src/db/job-executions.ts

@@ -29,7 +29,14 @@ import {
   RunnerSessionExhaustedError,
   RunningJobExecutionNotFoundError,
 } from '#core/errors.js';
-import {jobExecutionEnqueuedCount, jobExecutionLeaseExpiredCount} from '#metrics/instance.js';
+import {
+  type JobExecutionQueueTimeObservation,
+  jobExecutionEnqueuedCount,
+  jobExecutionLeaseExpiredCount,
+  type ProviderRunnerLifecycleObservation,
+  recordJobExecutionQueueTime,
+  recordProviderRunnerActivationToFirstClaim,
+} from '#metrics/instance.js';
 import type {Tx} from './db.js';
 import {db} from './db.js';
 import {runnersOutbox} from './schema/outbox.js';
@@ -192,7 +199,9 @@ export async function claimPendingJobExecution(params: {
 
   if (params.sessionLabels.length === 0) return null;
 
-  return await db().transaction(async (tx) => {
+  let activationToFirstClaimObservation: ProviderRunnerLifecycleObservation | null = null;
+  let queueTimeObservation: JobExecutionQueueTimeObservation | null = null;
+  const result = await db().transaction(async (tx) => {
     let provisionerId: string | null = null;
     let providerRunnerId: string | null = null;
     let runnerInstanceId: string | null = null;
@@ -288,6 +297,12 @@ export async function claimPendingJobExecution(params: {
     const claimed = inserted[0];
     if (!claimed) return null;
 
+    queueTimeObservation = {
+      durationMilliseconds: claimed.claimedAt.getTime() - row.createdAt.getTime(),
+      provider: null,
+      launchKind: params.maxClaims === null ? 'manual' : 'unknown',
+    };
+
     const runnerInstanceCondition = runnerInstanceId
       ? eq(providerRunners.id, runnerInstanceId)
       : provisionerId && providerRunnerId
@@ -297,13 +312,42 @@ export async function claimPendingJobExecution(params: {
           )
         : null;
     if (runnerInstanceCondition) {
-      await tx
+      const [claimedRunner] = await tx
         .update(providerRunners)
         .set({
           firstClaimedAt: sql`coalesce(${providerRunners.firstClaimedAt}, ${claimed.claimedAt})`,
           updatedAt: sql`now()`,
         })
-        .where(runnerInstanceCondition);
+        .where(runnerInstanceCondition)
+        .returning({
+          firstClaimedAt: providerRunners.firstClaimedAt,
+          isFirstClaim: sql<boolean>`${providerRunners.firstClaimedAt} = ${claimed.claimedAt}`,
+          provider: providerRunners.providerKind,
+          launchKind: providerRunners.launchKind,
+          runnerInstanceId: providerRunners.id,
+          sessionCreatedAtEpochMs: sql<number | null>`(
+            select extract(epoch from ${runnerSessions.createdAt})::double precision * 1000
+            from ${runnerSessions}
+            where ${runnerSessions.id} = ${params.runnerSessionId}
+          )`,
+        });
+      if (claimedRunner && queueTimeObservation) {
+        queueTimeObservation.provider = claimedRunner.provider;
+        queueTimeObservation.launchKind = claimedRunner.launchKind;
+      }
+      if (
+        claimedRunner?.isFirstClaim &&
+        claimedRunner.firstClaimedAt !== null &&
+        claimedRunner.sessionCreatedAtEpochMs !== null &&
+        claimedRunner.sessionCreatedAtEpochMs !== undefined
+      )
+        activationToFirstClaimObservation = {
+          durationMilliseconds:
+            claimedRunner.firstClaimedAt.getTime() - claimedRunner.sessionCreatedAtEpochMs,
+          provider: claimedRunner.provider,
+          launchKind: claimedRunner.launchKind,
+          runnerInstanceId: claimedRunner.runnerInstanceId,
+        };
     }
 
     if (params.maxClaims !== null) {
@@ -335,6 +379,10 @@ export async function claimPendingJobExecution(params: {
       projectId: row.projectId,
     };
   });
+  if (queueTimeObservation) recordJobExecutionQueueTime(queueTimeObservation);
+  if (activationToFirstClaimObservation)
+    recordProviderRunnerActivationToFirstClaim(activationToFirstClaimObservation);
+  return result;
 }
 
 async function touchRunnerSessionLiveness(params: {

package/src/db/reservation-locks.ts

@@ -0,0 +1,22 @@
+import {sql} from 'drizzle-orm';
+import type {Tx} from './db.js';
+
+export async function lockRunnerReservationAdvisoryKeysTx(
+  tx: Tx,
+  params: {provisionerId: string; reservationIds: readonly string[]},
+): Promise<void> {
+  const reservationIds = [...new Set(params.reservationIds)].sort();
+  if (reservationIds.length === 0) return;
+
+  const lockKeys = sql.join(
+    reservationIds.map(
+      (reservationId) => sql`(${`runners_assignment:${params.provisionerId}:${reservationId}`})`,
+    ),
+    sql`, `,
+  );
+  await tx.execute(sql`
+    select pg_advisory_xact_lock(hashtext(lock_key))
+    from (values ${lockKeys}) as reservation_locks(lock_key)
+    order by lock_key
+  `);
+}