fix(pool): count registered offline runners within boot time

lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.test.ts

@@ -41,17 +41,41 @@ describe('createEc2PoolCapability.countAvailableRunners', () => {
     expect(mockBootTimeExceeded).not.toHaveBeenCalled();
   });
 
-  it('does not count registered busy or offline runners', () => {
-    const runners: RunnerInfo[] = [
-      { id: 'i-busy', owner: 'owner', type: 'Org' },
-      { id: 'i-offline', owner: 'owner', type: 'Org' },
-    ];
-    const runnerStatus = new Map([
-      ['i-busy', { busy: true, status: 'online' }],
-      ['i-offline', { busy: false, status: 'offline' }],
-    ]);
+  it('does not count registered busy runners', () => {
+    const runners: RunnerInfo[] = [{ id: 'i-busy', owner: 'owner', type: 'Org' }];
+    const runnerStatus = new Map([['i-busy', { busy: true, status: 'online' }]]);
+
+    expect(capability.countAvailableRunners(runners, runnerStatus)).toBe(0);
+    expect(mockBootTimeExceeded).not.toHaveBeenCalled();
+  });
+
+  it('counts registered offline runners that are still booting', () => {
+    // With JIT config the runner is registered in GitHub before the agent on the instance
+    // connects, so it can be reported as offline while it is still booting.
+    const runners: RunnerInfo[] = [{ id: 'i-jit-booting', owner: 'owner', type: 'Org' }];
+    const runnerStatus = new Map([['i-jit-booting', { busy: false, status: 'offline' }]]);
+    mockBootTimeExceeded.mockReturnValue(false);
+
+    expect(capability.countAvailableRunners(runners, runnerStatus)).toBe(1);
+    expect(mockBootTimeExceeded).toHaveBeenCalledWith(runners[0]);
+  });
+
+  it('does not count registered offline runners whose boot time expired', () => {
+    const runners: RunnerInfo[] = [{ id: 'i-offline', owner: 'owner', type: 'Org' }];
+    const runnerStatus = new Map([['i-offline', { busy: false, status: 'offline' }]]);
+    mockBootTimeExceeded.mockReturnValue(true);
+
+    expect(capability.countAvailableRunners(runners, runnerStatus)).toBe(0);
+    expect(mockBootTimeExceeded).toHaveBeenCalledWith(runners[0]);
+  });
+
+  it('does not count registered offline runners that are busy', () => {
+    const runners: RunnerInfo[] = [{ id: 'i-offline-busy', owner: 'owner', type: 'Org' }];
+    const runnerStatus = new Map([['i-offline-busy', { busy: true, status: 'offline' }]]);
+    mockBootTimeExceeded.mockReturnValue(false);
 
     expect(capability.countAvailableRunners(runners, runnerStatus)).toBe(0);
+    expect(capability.countAvailableRunners(runners, runnerStatus, true)).toBe(0);
     expect(mockBootTimeExceeded).not.toHaveBeenCalled();
   });
 

lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.ts

@@ -13,13 +13,22 @@ function countAvailableEc2PoolRunners(
   // Runner should be considered idle if it is still booting, or is idle in GitHub
   let numberOfRunnersInPool = 0;
   for (const ec2Instance of ec2runners) {
-    if (
-      (runnerStatus.get(ec2Instance.id)?.busy === false || includeBusyRunners) &&
-      runnerStatus.get(ec2Instance.id)?.status === 'online'
-    ) {
+    const status = runnerStatus.get(ec2Instance.id);
+    if ((status?.busy === false || includeBusyRunners) && status?.status === 'online') {
       numberOfRunnersInPool++;
       logger.debug(`Runner ${ec2Instance.id} is idle in GitHub and counted as part of the pool`);
-    } else if (runnerStatus.get(ec2Instance.id) != null) {
+    } else if (status?.status === 'offline' && status.busy === false) {
+      // A runner can be registered before its agent connects (for example with JIT config) and is
+      // then reported as offline. Count it as booting until its boot time expires.
+      if (!bootTimeExceeded(ec2Instance)) {
+        numberOfRunnersInPool++;
+        logger.info(`Runner ${ec2Instance.id} is registered offline, still booting and counted as part of the pool`);
+      } else {
+        logger.debug(
+          `Runner ${ec2Instance.id} is registered offline past its boot time and NOT counted as part of the pool`,
+        );
+      }
+    } else if (status != null) {
       logger.debug(`Runner ${ec2Instance.id} is not idle in GitHub and NOT counted as part of the pool`);
     } else if (!bootTimeExceeded(ec2Instance)) {
       numberOfRunnersInPool++;