@shipfox/api-runners 12.2.0 → 12.4.0

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/index.ts

@@ -90,6 +90,7 @@ export type {
 } from './runner-instances.js';
 export {
   attachRunnerInstanceProviderId,
+  countStaleEnrolledRunnerInstances,
   isTerminalState,
   listActiveRunnerInstanceCountsByTemplateTx,
   listActiveRunnerInstances,

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
+  `);
+}