@shipfox/api-runners 12.1.0 → 12.3.0

package/src/db/reservations.test.ts

@@ -86,6 +86,350 @@ describe('pollDemandAndReserve', () => {
     });
   });
 
+  it('binds a runner whose intended reservation has passed its activation grace period', async () => {
+    const intendedReservation = await createIntendedReservation({
+      workspaceId,
+      expiresAt: new Date(Date.now() - 60_000),
+    });
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      intendedReservationId: intendedReservation.id,
+    });
+    await createPendingJobs(1, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(storedRunner).toMatchObject({
+      workspaceId,
+      reservationId: result.reservations[0]?.reservationId,
+      intendedReservationId: null,
+      assignedAt: expect.any(Date),
+    });
+  });
+
+  it('binds a runner whose intended reservation was deleted after its grace period', async () => {
+    const intendedReservation = await createIntendedReservation({
+      workspaceId,
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+    await db().delete(reservations).where(eq(reservations.id, intendedReservation.id));
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      intendedReservationId: intendedReservation.id,
+    });
+    await createPendingJobs(1, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(storedRunner).toMatchObject({
+      workspaceId,
+      reservationId: result.reservations[0]?.reservationId,
+      intendedReservationId: null,
+      assignedAt: expect.any(Date),
+    });
+  });
+
+  it('binds an assigned runner after its reservation grace period expires', async () => {
+    const staleReservation = await createIntendedReservation({
+      workspaceId,
+      expiresAt: new Date(Date.now() - 60_000),
+    });
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      workspaceId,
+      reservationId: staleReservation.id,
+      intendedReservationId: staleReservation.id,
+    });
+    const [activationToken] = await db()
+      .insert(runnerActivationTokens)
+      .values({
+        runnerInstanceId: runner.id,
+        hashedToken: crypto.randomUUID(),
+        prefix: 'test',
+        expiresAt: new Date(Date.now() + 60_000),
+      })
+      .returning();
+    if (!activationToken) throw new Error('Expected activation token');
+    await createPendingJobs(1, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    const [storedToken] = await db()
+      .select()
+      .from(runnerActivationTokens)
+      .where(eq(runnerActivationTokens.id, activationToken.id));
+    expect(storedRunner).toMatchObject({
+      workspaceId,
+      reservationId: result.reservations[0]?.reservationId,
+      intendedReservationId: null,
+      assignedAt: expect.any(Date),
+    });
+    expect(storedToken?.revokedAt).toBeInstanceOf(Date);
+  });
+
+  it('binds an assigned runner whose reservation was deleted', async () => {
+    const staleReservation = await createIntendedReservation({
+      workspaceId,
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+    await db().delete(reservations).where(eq(reservations.id, staleReservation.id));
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      workspaceId,
+      reservationId: staleReservation.id,
+      intendedReservationId: staleReservation.id,
+    });
+    await createPendingJobs(1, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(storedRunner).toMatchObject({
+      workspaceId,
+      reservationId: result.reservations[0]?.reservationId,
+      intendedReservationId: null,
+      assignedAt: expect.any(Date),
+    });
+  });
+
+  it('does not bind an assigned runner within its activation grace period', async () => {
+    const liveReservation = await createIntendedReservation({
+      workspaceId,
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      workspaceId,
+      reservationId: liveReservation.id,
+    });
+    await createPendingJobs(2, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 2,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 2)],
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(result.reservations[0]?.count).toBe(1);
+    expect(storedRunner).toMatchObject({
+      workspaceId,
+      reservationId: liveReservation.id,
+      intendedReservationId: null,
+      assignedAt: null,
+    });
+  });
+
+  it('does not bind a runner whose intended reservation is within its activation grace period', async () => {
+    const intendedReservation = await createIntendedReservation({
+      workspaceId: crypto.randomUUID(),
+      expiresAt: new Date(Date.now() + 60_000),
+    });
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      intendedReservationId: intendedReservation.id,
+    });
+    await createPendingJobs(1, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(result.reservations[0]?.count).toBe(1);
+    expect(storedRunner).toMatchObject({
+      workspaceId: null,
+      reservationId: null,
+      intendedReservationId: intendedReservation.id,
+      assignedAt: null,
+    });
+  });
+
+  it('does not rebind a stale runner from another workspace in a workspace-scoped poll', async () => {
+    const previousWorkspaceId = crypto.randomUUID();
+    const staleReservation = await createIntendedReservation({
+      workspaceId: previousWorkspaceId,
+      expiresAt: new Date(Date.now() - 60_000),
+    });
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      workspaceId: previousWorkspaceId,
+      reservationId: staleReservation.id,
+      intendedReservationId: staleReservation.id,
+    });
+    await createPendingJobs(1, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(result.reservations[0]?.count).toBe(1);
+    expect(storedRunner).toMatchObject({
+      workspaceId: previousWorkspaceId,
+      reservationId: staleReservation.id,
+      intendedReservationId: staleReservation.id,
+    });
+  });
+
+  it('rebinds an unowned runner whose reservation was deleted', async () => {
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      reservationId: crypto.randomUUID(),
+    });
+    await createPendingJobs(1, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(storedRunner).toMatchObject({
+      workspaceId,
+      reservationId: result.reservations[0]?.reservationId,
+      intendedReservationId: null,
+      assignedAt: expect.any(Date),
+    });
+  });
+
+  it('rebinds a stale runner across workspaces for an installation-scoped poll', async () => {
+    const previousWorkspaceId = crypto.randomUUID();
+    const targetWorkspaceId = crypto.randomUUID();
+    const staleReservation = await createIntendedReservation({
+      workspaceId: previousWorkspaceId,
+      expiresAt: new Date(Date.now() - 60_000),
+    });
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      workspaceId: previousWorkspaceId,
+      reservationId: staleReservation.id,
+      intendedReservationId: staleReservation.id,
+    });
+    await pendingJobFactory.create({workspaceId: targetWorkspaceId, requiredLabels: ['linux']});
+
+    const result = await pollInstallationDemandAndReserve({
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+      capabilityWindowSeconds: 60,
+      eligibleWorkspaceIds: new Set([targetWorkspaceId]),
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(result.reservations[0]).toMatchObject({workspaceId: targetWorkspaceId, count: 1});
+    expect(storedRunner).toMatchObject({
+      workspaceId: targetWorkspaceId,
+      reservationId: result.reservations[0]?.reservationId,
+      intendedReservationId: null,
+      assignedAt: expect.any(Date),
+    });
+  });
+
+  it('does not rebind a released runner after its reservation expires', async () => {
+    const staleReservation = await createIntendedReservation({
+      workspaceId,
+      expiresAt: new Date(Date.now() - 60_000),
+    });
+    const releasedAt = new Date(Date.now() - 30_000);
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      workspaceId,
+      reservationId: staleReservation.id,
+      intendedReservationId: staleReservation.id,
+      reservationReleasedAt: releasedAt,
+    });
+    await createPendingJobs(1, ['linux']);
+
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(result.reservations[0]?.count).toBe(1);
+    expect(storedRunner).toMatchObject({
+      workspaceId,
+      reservationId: staleReservation.id,
+      intendedReservationId: staleReservation.id,
+      reservationReleasedAt: releasedAt,
+    });
+  });
+
   it('does not bind a running runner without an active control session', async () => {
     const [runner] = await db()
       .insert(providerRunners)
@@ -461,6 +805,64 @@ describe('pollDemandAndReserve', () => {
     expect(storedToken?.revokedAt).toBeInstanceOf(Date);
   });
 
+  it('preserves activation tokens for active runners with an intended reservation', async () => {
+    const [reservation] = await db()
+      .insert(reservations)
+      .values({
+        workspaceId,
+        provisionerId,
+        requiredLabels: ['linux'],
+        count: 1,
+        expiresAt: new Date(Date.now() + 60_000),
+      })
+      .returning();
+    if (!reservation) throw new Error('Expected reservation');
+    const [runner] = await db()
+      .insert(providerRunners)
+      .values({
+        workspaceId,
+        provisionerId,
+        providerRunnerId: crypto.randomUUID(),
+        reservationId: reservation.id,
+        intendedReservationId: reservation.id,
+        runnerSessionId: crypto.randomUUID(),
+        labels: ['linux'],
+        state: 'running',
+        reportedAt: new Date(),
+      })
+      .returning();
+    if (!runner) throw new Error('Expected runner instance');
+    const [activationToken] = await db()
+      .insert(runnerActivationTokens)
+      .values({
+        runnerInstanceId: runner.id,
+        hashedToken: crypto.randomUUID(),
+        prefix: 'test',
+        expiresAt: new Date(Date.now() + 60_000),
+      })
+      .returning();
+    if (!activationToken) throw new Error('Expected activation token');
+
+    const deleted = await deleteReservationsByIds([reservation.id]);
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    const [storedToken] = await db()
+      .select()
+      .from(runnerActivationTokens)
+      .where(eq(runnerActivationTokens.id, activationToken.id));
+    expect(deleted).toBe(1);
+    expect(storedRunner).toMatchObject({
+      reservationId: reservation.id,
+      intendedReservationId: null,
+      runnerSessionId: expect.any(String),
+      reservationReleasedAt: null,
+    });
+    expect(storedToken?.revokedAt).toBeNull();
+  });
+
   it('unbinds prewarmed runners and revokes activation tokens when deleting expired reservations', async () => {
     const runner = await createIdleRunner({labels: ['linux']});
     await createPendingJobs(1, ['linux']);
@@ -509,6 +911,115 @@ describe('pollDemandAndReserve', () => {
     expect(storedToken?.revokedAt).toBeInstanceOf(Date);
   });
 
+  it('clears intended reservations and rebinds runners when deleting a reservation by id', async () => {
+    const [reservation, survivingReservation] = await db()
+      .insert(reservations)
+      .values([
+        {
+          workspaceId,
+          provisionerId,
+          requiredLabels: ['linux'],
+          count: 1,
+          expiresAt: new Date(Date.now() + 60_000),
+        },
+        {
+          workspaceId: crypto.randomUUID(),
+          provisionerId,
+          requiredLabels: ['linux'],
+          count: 1,
+          expiresAt: new Date(Date.now() + 60_000),
+        },
+      ])
+      .returning();
+    if (!reservation || !survivingReservation) throw new Error('Expected reservations');
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      intendedReservationId: reservation.id,
+    });
+    const survivingRunner = await createIdleRunner({
+      labels: ['linux'],
+      intendedReservationId: survivingReservation.id,
+    });
+
+    const deleted = await deleteReservationsByIds([reservation.id]);
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(deleted).toBe(1);
+    expect(storedRunner).toMatchObject({
+      intendedReservationId: null,
+      workspaceId: null,
+      reservationId: null,
+      assignedAt: null,
+      state: 'running',
+      reservationReleasedAt: null,
+    });
+    const [storedSurvivingRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, survivingRunner.id));
+    expect(storedSurvivingRunner?.intendedReservationId).toBe(survivingReservation.id);
+
+    await createPendingJobs(1, ['linux']);
+    const result = await pollDemandAndReserve({
+      workspaceId,
+      provisionerId,
+      maxReservations: 1,
+      ttlSeconds: 60,
+      templates: [template('linux', ['linux'], 1)],
+    });
+    const reboundReservationId = result.reservations[0]?.reservationId;
+    const [reboundRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(reboundReservationId).toEqual(expect.any(String));
+    expect(reboundRunner).toMatchObject({
+      intendedReservationId: null,
+      workspaceId,
+      reservationId: reboundReservationId,
+      assignedAt: expect.any(Date),
+      state: 'running',
+      reservationReleasedAt: null,
+    });
+  });
+
+  it('clears intended reservations when deleting expired reservations', async () => {
+    const [reservation] = await db()
+      .insert(reservations)
+      .values({
+        workspaceId,
+        provisionerId,
+        requiredLabels: ['linux'],
+        count: 1,
+        expiresAt: new Date(Date.now() - 60_000),
+      })
+      .returning();
+    if (!reservation) throw new Error('Expected reservation');
+    const runner = await createIdleRunner({
+      labels: ['linux'],
+      intendedReservationId: reservation.id,
+    });
+
+    const deleted = await deleteExpiredReservations();
+
+    const [storedRunner] = await db()
+      .select()
+      .from(providerRunners)
+      .where(eq(providerRunners.id, runner.id));
+    expect(deleted).toBeGreaterThanOrEqual(1);
+    expect(storedRunner).toMatchObject({
+      intendedReservationId: null,
+      workspaceId: null,
+      reservationId: null,
+      assignedAt: null,
+      state: 'running',
+      reservationReleasedAt: null,
+    });
+  });
+
   it('decrements reservation units inside a caller transaction', async () => {
     await reservationFactory.create({
       workspaceId,
@@ -641,17 +1152,43 @@ describe('pollDemandAndReserve', () => {
     }
   }
 
+  async function createIntendedReservation(params: {
+    workspaceId: string;
+    expiresAt: Date;
+  }): Promise<{id: string}> {
+    const [reservation] = await db()
+      .insert(reservations)
+      .values({
+        workspaceId: params.workspaceId,
+        provisionerId,
+        requiredLabels: ['linux'],
+        count: 1,
+        expiresAt: params.expiresAt,
+      })
+      .returning({id: reservations.id});
+    if (!reservation) throw new Error('Expected reservation');
+    return reservation;
+  }
+
   async function createIdleRunner(params: {
     labels: string[];
     createdAt?: Date;
     controlSessionExpiresAt?: Date;
+    workspaceId?: string | null;
+    reservationId?: string | null;
+    intendedReservationId?: string | null;
+    reservationReleasedAt?: Date | null;
   }) {
     const createdAt = params.createdAt ?? new Date();
     const [runner] = await db()
       .insert(providerRunners)
       .values({
         provisionerId,
+        workspaceId: params.workspaceId,
+        reservationId: params.reservationId,
         providerRunnerId: crypto.randomUUID(),
+        intendedReservationId: params.intendedReservationId,
+        reservationReleasedAt: params.reservationReleasedAt,
         labels: params.labels,
         state: 'running',
         reportedAt: createdAt,