Reactor logic, add tests for booting and orphan runners

npalm · npalm · commit 1d42c4185727 · 2022-12-28T12:35:49.000+01:00
diff --git a/modules/runners/lambdas/runners/src/pool/pool.test.ts b/modules/runners/lambdas/runners/src/pool/pool.test.ts
@@ -1,10 +1,11 @@
 import { Octokit } from '@octokit/rest';
 import { mocked } from 'jest-mock';
+import moment from 'moment-timezone';
 import nock from 'nock';
 
 import { listEC2Runners } from '../aws/runners';
 import * as ghAuth from '../gh-auth/gh-auth';
-import * as scale from '../scale-runners/scale-up';
+import { createRunners } from '../scale-runners/scale-up';
 import { adjust } from './pool';
 
 const mockOctokit = {
@@ -24,6 +25,7 @@ jest.mock('@octokit/rest', () => ({
 
 jest.mock('./../aws/runners');
 jest.mock('./../gh-auth/gh-auth');
+jest.mock('./../scale-runners/scale-up');
 
 const mocktokit = Octokit as jest.MockedClass<typeof Octokit>;
 const mockedAppAuth = mocked(ghAuth.createGithubAppAuth, {
@@ -36,6 +38,28 @@ const mockListRunners = mocked(listEC2Runners);
 const cleanEnv = process.env;
 
 const ORG = 'my-org';
+const MINIMUM_TIME_RUNNING = 15;
+
+const ec2InstancesRegistered = [
+  {
+    instanceId: 'i-1',
+    launchTime: new Date(),
+    type: 'Org',
+    owner: ORG,
+  },
+  {
+    instanceId: 'i-2',
+    launchTime: new Date(),
+    type: 'Org',
+    owner: ORG,
+  },
+  {
+    instanceId: 'i-3',
+    launchTime: new Date(),
+    type: 'Org',
+    owner: ORG,
+  },
+];
 
 beforeEach(() => {
   nock.disableNetConnect();
@@ -54,6 +78,7 @@ beforeEach(() => {
   process.env.INSTANCE_TYPES = 'm5.large';
   process.env.INSTANCE_TARGET_CAPACITY_TYPE = 'spot';
   process.env.RUNNER_OWNER = ORG;
+  process.env.RUNNER_BOOT_TIME_IN_MINUTES = MINIMUM_TIME_RUNNING.toString();
 
   const mockTokenReturnValue = {
     data: {
@@ -87,44 +112,9 @@ beforeEach(() => {
       busy: false,
       labels: [],
     },
-    {
-      id: 11,
-      name: 'j-1', // some runner of another env
-      os: 'linux',
-      status: 'online',
-      busy: false,
-      labels: [],
-    },
-    {
-      id: 12,
-      name: 'j-2', // some runner of another env
-      os: 'linux',
-      status: 'online',
-      busy: true,
-      labels: [],
-    },
   ]);
 
-  mockListRunners.mockImplementation(async () => [
-    {
-      instanceId: 'i-1',
-      launchTime: new Date(),
-      type: 'Org',
-      owner: ORG,
-    },
-    {
-      instanceId: 'i-2',
-      launchTime: new Date(),
-      type: 'Org',
-      owner: ORG,
-    },
-    {
-      instanceId: 'i-3',
-      launchTime: new Date(),
-      type: 'Org',
-      owner: ORG,
-    },
-  ]);
+  mockListRunners.mockImplementation(async () => ec2InstancesRegistered);
 
   const mockInstallationIdReturnValueOrgs = {
     data: {
@@ -155,16 +145,64 @@ beforeEach(() => {
 
 describe('Test simple pool.', () => {
   describe('With GitHub Cloud', () => {
-    it('Top up pool with pool size 2.', async () => {
-      const spy = jest.spyOn(scale, 'createRunners');
-      await expect(adjust({ poolSize: 2 })).resolves;
-      expect(spy).toBeCalled;
+    it('Top up pool with pool size 2 registered.', async () => {
+      await expect(await adjust({ poolSize: 3 })).resolves;
+      expect(createRunners).toHaveBeenCalledTimes(1);
     });
 
     it('Should not top up if pool size is reached.', async () => {
-      const spy = jest.spyOn(scale, 'createRunners');
-      await expect(adjust({ poolSize: 1 })).resolves;
-      expect(spy).not.toHaveBeenCalled;
+      await expect(await adjust({ poolSize: 1 })).resolves;
+      expect(createRunners).not.toHaveBeenCalled();
+    });
+
+    it('Should top up if pool size is not reached including a booting instance.', async () => {
+      mockListRunners.mockImplementation(async () => [
+        ...ec2InstancesRegistered,
+        {
+          instanceId: 'i-4-still-booting',
+          launchTime: moment(new Date())
+            .subtract(MINIMUM_TIME_RUNNING - 3, 'minutes')
+            .toDate(),
+          type: 'Org',
+          owner: ORG,
+        },
+        {
+          instanceId: 'i-5-orphan',
+          launchTime: moment(new Date())
+            .subtract(MINIMUM_TIME_RUNNING + 3, 'minutes')
+            .toDate(),
+          type: 'Org',
+          owner: ORG,
+        },
+      ]);
+
+      await expect(await adjust({ poolSize: 5 })).resolves;
+      expect(createRunners).toHaveBeenCalledTimes(1);
+    });
+
+    it('Should not top up if pool size is reached including a booting instance.', async () => {
+      mockListRunners.mockImplementation(async () => [
+        ...ec2InstancesRegistered,
+        {
+          instanceId: 'i-4-still-booting',
+          launchTime: moment(new Date())
+            .subtract(MINIMUM_TIME_RUNNING - 3, 'minutes')
+            .toDate(),
+          type: 'Org',
+          owner: ORG,
+        },
+        {
+          instanceId: 'i-5-orphan',
+          launchTime: moment(new Date())
+            .subtract(MINIMUM_TIME_RUNNING + 3, 'minutes')
+            .toDate(),
+          type: 'Org',
+          owner: ORG,
+        },
+      ]);
+
+      await expect(await adjust({ poolSize: 2 })).resolves;
+      expect(createRunners).not.toHaveBeenCalled();
     });
   });
 
@@ -174,9 +212,8 @@ describe('Test simple pool.', () => {
     });
 
     it('Top up if the pool size is set to 5', async () => {
-      const spy = jest.spyOn(scale, 'createRunners');
-      await expect(adjust({ poolSize: 5 })).resolves;
-      expect(spy).toBeCalled;
+      await expect(await adjust({ poolSize: 3 })).resolves;
+      expect(createRunners).toHaveBeenCalledTimes(1);
     });
   });
 });
diff --git a/modules/runners/lambdas/runners/src/pool/pool.ts b/modules/runners/lambdas/runners/src/pool/pool.ts
@@ -12,6 +12,11 @@ export interface PoolEvent {
   poolSize: number;
 }
 
+interface RunnerStatus {
+  busy: boolean;
+  status: string;
+}
+
 export async function adjust(event: PoolEvent): Promise<void> {
   logger.info(`Checking current pool size against pool of size: ${event.poolSize}`);
   const runnerExtraLabels = process.env.RUNNER_EXTRA_LABELS;
@@ -45,7 +50,10 @@ export async function adjust(event: PoolEvent): Promise<void> {
       per_page: 100,
     },
   );
-  const githubIdleRunners = runners.filter((r) => !r.busy && r.status === 'online');
+  const runnerStatus = new Map<string, RunnerStatus>();
+  for (const runner of runners) {
+    runnerStatus.set(runner.name, { busy: runner.busy, status: runner.status });
+  }
 
   // Look up the managed ec2 runners in AWS, but running does not mean idle
   const ec2runners = await listEC2Runners({
@@ -56,13 +64,27 @@ export async function adjust(event: PoolEvent): Promise<void> {
   });
 
   // Runner should be considered idle if it is still booting, or is idle in GitHub
-  const managedIdleRunners = ec2runners
-    .filter((r) => {
-      return githubIdleRunners.some((ghRunner) => ghRunner.name == r.instanceId) || !bootTimeExceeded(r);
-    })
-    .map((r) => r.instanceId);
+  let numberOfRunnersInPool = 0;
+  for (const ec2Instance of ec2runners) {
+    if (
+      runnerStatus.get(ec2Instance.instanceId)?.busy === false &&
+      runnerStatus.get(ec2Instance.instanceId)?.status === 'online'
+    ) {
+      numberOfRunnersInPool++;
+      logger.debug(`Runner ${ec2Instance.instanceId} is idle in GitHub and counted as part of the pool`);
+    } else if (runnerStatus.get(ec2Instance.instanceId) != null) {
+      logger.debug(`Runner ${ec2Instance.instanceId} is not idle in GitHub and NOT counted as part of the pool`);
+    } else if (!bootTimeExceeded(ec2Instance)) {
+      numberOfRunnersInPool++;
+      logger.info(`Runner ${ec2Instance.instanceId} is still booting and counted as part of the pool`);
+    } else {
+      logger.debug(
+        `Runner ${ec2Instance.instanceId} is not idle in GitHub nor booting and not counted as part of the pool`,
+      );
+    }
+  }
 
-  const topUp = event.poolSize - managedIdleRunners.length;
+  const topUp = event.poolSize - numberOfRunnersInPool;
   if (topUp > 0) {
     logger.info(`The pool will be topped up with ${topUp} runners.`);
     await createRunners(
@@ -90,7 +112,7 @@ export async function adjust(event: PoolEvent): Promise<void> {
       githubInstallationClient,
     );
   } else {
-    logger.info(`Pool will not be topped up. Found ${managedIdleRunners} managed idle runners.`);
+    logger.info(`Pool will not be topped up. Found ${numberOfRunnersInPool} managed idle runners.`);
   }
 }
 
