feat: add standby pool functionality for EC2 runners with start/stop capabilities

Author: Brend-Smits

lambdas/functions/control-plane/src/aws/runners.d.ts

@@ -26,6 +26,7 @@ export interface ListRunnerFilters {
   environment?: string;
   orphan?: boolean;
   statuses?: string[];
+  standby?: boolean;
 }
 
 export interface RunnerInputParameters {

lambdas/functions/control-plane/src/aws/runners.test.ts

@@ -11,6 +11,8 @@ import {
   type DescribeInstancesResult,
   EC2Client,
   SpotAllocationStrategy,
+  StartInstancesCommand,
+  StopInstancesCommand,
   TerminateInstancesCommand,
 } from '@aws-sdk/client-ec2';
 import { GetParameterCommand, type GetParameterResult, PutParameterCommand, SSMClient } from '@aws-sdk/client-ssm';
@@ -19,7 +21,7 @@ import 'aws-sdk-client-mock-jest/vitest';
 
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 import ScaleError from './../scale-runners/ScaleError';
-import { createRunner, listEC2Runners, tag, terminateRunner, untag } from './runners';
+import { createRunner, listEC2Runners, tag, terminateRunner, untag, stopRunner, startRunner } from './runners';
 import type { RunnerInfo, RunnerInputParameters, RunnerType } from './runners.d';
 
 process.env.AWS_REGION = 'eu-east-1';
@@ -245,6 +247,19 @@ describe('list instances', () => {
       ],
     });
   });
+
+  it('Filter instances for standby state', async () => {
+    mockEC2Client.on(DescribeInstancesCommand).resolves(mockRunningInstances);
+    await listEC2Runners({ statuses: ['stopped'], standby: true, environment: ENVIRONMENT });
+    expect(mockEC2Client).toHaveReceivedCommandWith(DescribeInstancesCommand, {
+      Filters: [
+        { Name: 'instance-state-name', Values: ['stopped'] },
+        { Name: 'tag:ghr:environment', Values: [ENVIRONMENT] },
+        { Name: 'tag:ghr:state', Values: ['standby'] },
+        { Name: 'tag:ghr:Application', Values: ['github-action-runner'] },
+      ],
+    });
+  });
 });
 
 describe('terminate runner', () => {
@@ -266,6 +281,57 @@ describe('terminate runner', () => {
   });
 });
 
+describe('stop runner', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+  it('calls stop instances and tags with standby state', async () => {
+    mockEC2Client.on(StopInstancesCommand).resolves({});
+    mockEC2Client.on(CreateTagsCommand).resolves({});
+    const runner: RunnerInfo = {
+      instanceId: 'instance-3',
+      owner: 'owner-3',
+      type: 'Org',
+    };
+    await stopRunner(runner.instanceId);
+
+    expect(mockEC2Client).toHaveReceivedCommandWith(StopInstancesCommand, {
+      InstanceIds: [runner.instanceId],
+    });
+    expect(mockEC2Client).toHaveReceivedCommandWith(CreateTagsCommand, {
+      Resources: [runner.instanceId],
+      Tags: [
+        { Key: 'ghr:state', Value: 'standby' },
+        { Key: 'ghr:stopped_at', Value: expect.any(String) },
+      ],
+    });
+  });
+});
+
+describe('start runner', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+  it('calls start instances and removes standby tags', async () => {
+    mockEC2Client.on(StartInstancesCommand).resolves({});
+    mockEC2Client.on(DeleteTagsCommand).resolves({});
+    const runner: RunnerInfo = {
+      instanceId: 'instance-4',
+      owner: 'owner-4',
+      type: 'Org',
+    };
+    await startRunner(runner.instanceId);
+
+    expect(mockEC2Client).toHaveReceivedCommandWith(StartInstancesCommand, {
+      InstanceIds: [runner.instanceId],
+    });
+    expect(mockEC2Client).toHaveReceivedCommandWith(DeleteTagsCommand, {
+      Resources: [runner.instanceId],
+      Tags: [{ Key: 'ghr:state' }, { Key: 'ghr:stopped_at' }],
+    });
+  });
+});
+
 describe('tag runner', () => {
   beforeEach(() => {
     vi.clearAllMocks();

lambdas/functions/control-plane/src/aws/runners.ts

@@ -7,6 +7,8 @@ import {
   DescribeInstancesResult,
   EC2Client,
   FleetLaunchTemplateOverridesRequest,
+  StartInstancesCommand,
+  StopInstancesCommand,
   Tag,
   TerminateInstancesCommand,
   _InstanceType,
@@ -52,6 +54,9 @@ function constructFilters(filters?: Runners.ListRunnerFilters): Ec2Filter[][] {
     if (filters.orphan) {
       ec2FiltersBase.push({ Name: 'tag:ghr:orphan', Values: ['true'] });
     }
+    if (filters.standby) {
+      ec2FiltersBase.push({ Name: 'tag:ghr:state', Values: ['standby'] });
+    }
   }
 
   for (const key of ['tag:ghr:Application']) {
@@ -108,6 +113,25 @@ export async function terminateRunner(instanceId: string): Promise<void> {
   logger.debug(`Runner ${instanceId} has been terminated.`);
 }
 
+export async function stopRunner(instanceId: string): Promise<void> {
+  logger.debug(`Runner '${instanceId}' will be stopped.`);
+  const ec2 = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION }));
+  await ec2.send(new StopInstancesCommand({ InstanceIds: [instanceId] }));
+  await tag(instanceId, [
+    { Key: 'ghr:state', Value: 'standby' },
+    { Key: 'ghr:stopped_at', Value: new Date().toISOString() },
+  ]);
+  logger.debug(`Runner ${instanceId} has been stopped.`);
+}
+
+export async function startRunner(instanceId: string): Promise<void> {
+  logger.debug(`Runner '${instanceId}' will be started.`);
+  const ec2 = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION }));
+  await ec2.send(new StartInstancesCommand({ InstanceIds: [instanceId] }));
+  await untag(instanceId, [{ Key: 'ghr:state' }, { Key: 'ghr:stopped_at' }]);
+  logger.debug(`Runner ${instanceId} has been started.`);
+}
+
 export async function tag(instanceId: string, tags: Tag[]): Promise<void> {
   logger.debug(`Tagging '${instanceId}'`, { tags });
   const ec2 = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION }));

lambdas/functions/control-plane/src/scale-runners/scale-down.test.ts

@@ -5,7 +5,7 @@ import nock from 'nock';
 
 import { RunnerInfo, RunnerList } from '../aws/runners.d';
 import * as ghAuth from '../github/auth';
-import { listEC2Runners, terminateRunner, tag, untag } from './../aws/runners';
+import { listEC2Runners, terminateRunner, tag, untag, stopRunner } from './../aws/runners';
 import { githubCache } from './cache';
 import { newestFirstStrategy, oldestFirstStrategy, scaleDown } from './scale-down';
 import { describe, it, expect, beforeEach, vi } from 'vitest';
@@ -32,12 +32,13 @@ vi.mock('@octokit/rest', () => ({
 }));
 
 vi.mock('./../aws/runners', async (importOriginal) => {
-  const actual = await importOriginal();
+  const actual = await importOriginal<typeof import('./../aws/runners')>();
   return {
     ...actual,
     tag: vi.fn(),
     untag: vi.fn(),
     terminateRunner: vi.fn(),
+    stopRunner: vi.fn(),
     listEC2Runners: vi.fn(),
   };
 });
@@ -60,14 +61,14 @@ vi.mock('./cache', async () => ({
   },
 }));
 
-const mocktokit = Octokit as vi.MockedClass<typeof Octokit>;
 const mockedAppAuth = vi.mocked(ghAuth.createGithubAppAuth);
 const mockedInstallationAuth = vi.mocked(ghAuth.createGithubInstallationAuth);
 const mockCreateClient = vi.mocked(ghAuth.createOctokitClient);
 const mockListRunners = vi.mocked(listEC2Runners);
 const mockTagRunners = vi.mocked(tag);
 const mockUntagRunners = vi.mocked(untag);
 const mockTerminateRunners = vi.mocked(terminateRunner);
+const mockStopRunners = vi.mocked(stopRunner);
 
 export interface TestData {
   repositoryName: string;
@@ -180,7 +181,7 @@ describe('Scale down runners', () => {
       repositorySelection: 'all',
       installationId: 0,
     });
-    mockCreateClient.mockResolvedValue(new mocktokit());
+    mockCreateClient.mockResolvedValue(mockOctokit as unknown as Octokit);
   });
 
   const endpoints = ['https://api.github.com', 'https://github.enterprise.something', 'https://companyname.ghe.com'];
@@ -286,6 +287,35 @@ describe('Scale down runners', () => {
         checkNonTerminated(runners);
       });
 
+      it(`Should stop idle runner to standby when standby pool enabled`, async () => {
+        process.env.STANDBY_POOL_SIZE = '2';
+        process.env.STANDBY_IDLE_TIME_MINUTES = '0';
+        
+        const runners = [
+          createRunnerTestData('idle-1', type, MINIMUM_TIME_RUNNING_IN_MINUTES + 1, true, false, false),
+          createRunnerTestData('idle-2', type, MINIMUM_TIME_RUNNING_IN_MINUTES + 1, true, false, false),
+        ];
+
+        mockGitHubRunners(runners);
+        mockAwsRunners(runners);
+        
+        const originalListRunners = mockListRunners.getMockImplementation();
+        mockListRunners.mockImplementation(async (filters) => {
+          if (filters?.statuses?.includes('stopped') && filters?.standby) {
+            return [];
+          }
+          return originalListRunners ? originalListRunners(filters) : runners;
+        });
+
+        await scaleDown();
+
+        expect(mockStopRunners).toHaveBeenCalledTimes(2);
+        expect(mockTerminateRunners).not.toHaveBeenCalled();
+        
+        delete process.env.STANDBY_POOL_SIZE;
+        delete process.env.STANDBY_IDLE_TIME_MINUTES;
+      });
+
       it(`Should not terminate a runner that became busy just before deregister runner.`, async () => {
         // setup
         const runners = [

lambdas/functions/control-plane/src/scale-runners/scale-down.ts

@@ -5,7 +5,7 @@ import { createChildLogger } from '@aws-github-runner/aws-powertools-util';
 import moment from 'moment';
 
 import { createGithubAppAuth, createGithubInstallationAuth, createOctokitClient } from '../github/auth';
-import { bootTimeExceeded, listEC2Runners, tag, untag, terminateRunner } from './../aws/runners';
+import { bootTimeExceeded, listEC2Runners, tag, untag, terminateRunner, stopRunner } from './../aws/runners';
 import { RunnerInfo, RunnerList } from './../aws/runners.d';
 import { GhRunners, githubCache } from './cache';
 import { ScalingDownConfig, getEvictionStrategy, getIdleRunnerCount } from './scale-down-config';
@@ -127,39 +127,52 @@ function runnerMinimumTimeExceeded(runner: RunnerInfo): boolean {
   return launchTimePlusMinimum < now;
 }
 
-async function removeRunner(ec2runner: RunnerInfo, ghRunnerIds: number[]): Promise<void> {
+function runnerIdleTimeExceeded(runner: RunnerInfo): boolean {
+  const idleTimeMinutes = parseInt(process.env.STANDBY_IDLE_TIME_MINUTES || '0');
+  if (idleTimeMinutes === 0) return true;
+  
+  const launchTimePlusIdle = moment(runner.launchTime).utc().add(idleTimeMinutes, 'minutes');
+  const now = moment(new Date()).utc();
+  return launchTimePlusIdle < now;
+}
+
+async function removeRunner(ec2runner: RunnerInfo, ghRunnerIds: number[], shouldStop: boolean = false): Promise<void> {
   const githubAppClient = await getOrCreateOctokit(ec2runner);
   try {
     const states = await Promise.all(
       ghRunnerIds.map(async (ghRunnerId) => {
-        // Get busy state instead of using the output of listGitHubRunners(...) to minimize to race condition.
         return await getGitHubRunnerBusyState(githubAppClient, ec2runner, ghRunnerId);
       }),
     );
 
     if (states.every((busy) => busy === false)) {
-      const statuses = await Promise.all(
-        ghRunnerIds.map(async (ghRunnerId) => {
-          return (
-            ec2runner.type === 'Org'
-              ? await githubAppClient.actions.deleteSelfHostedRunnerFromOrg({
-                  runner_id: ghRunnerId,
-                  org: ec2runner.owner,
-                })
-              : await githubAppClient.actions.deleteSelfHostedRunnerFromRepo({
-                  runner_id: ghRunnerId,
-                  owner: ec2runner.owner.split('/')[0],
-                  repo: ec2runner.owner.split('/')[1],
-                })
-          ).status;
-        }),
-      );
-
-      if (statuses.every((status) => status == 204)) {
-        await terminateRunner(ec2runner.instanceId);
-        logger.info(`AWS runner instance '${ec2runner.instanceId}' is terminated and GitHub runner is de-registered.`);
+      if (shouldStop) {
+        await stopRunner(ec2runner.instanceId);
+        logger.info(`AWS runner instance '${ec2runner.instanceId}' is stopped and moved to standby.`);
       } else {
-        logger.error(`Failed to de-register GitHub runner: ${statuses}`);
+        const statuses = await Promise.all(
+          ghRunnerIds.map(async (ghRunnerId) => {
+            return (
+              ec2runner.type === 'Org'
+                ? await githubAppClient.actions.deleteSelfHostedRunnerFromOrg({
+                    runner_id: ghRunnerId,
+                    org: ec2runner.owner,
+                  })
+                : await githubAppClient.actions.deleteSelfHostedRunnerFromRepo({
+                    runner_id: ghRunnerId,
+                    owner: ec2runner.owner.split('/')[0],
+                    repo: ec2runner.owner.split('/')[1],
+                  })
+            ).status;
+          }),
+        );
+
+        if (statuses.every((status) => status == 204)) {
+          await terminateRunner(ec2runner.instanceId);
+          logger.info(`AWS runner instance '${ec2runner.instanceId}' is terminated and GitHub runner is de-registered.`);
+        } else {
+          logger.error(`Failed to de-register GitHub runner: ${statuses}`);
+        }
       }
     } else {
       logger.info(`Runner '${ec2runner.instanceId}' cannot be de-registered, because it is still busy.`);
@@ -178,13 +191,25 @@ async function evaluateAndRemoveRunners(
   let idleCounter = getIdleRunnerCount(scaleDownConfigs);
   const evictionStrategy = getEvictionStrategy(scaleDownConfigs);
   const ownerTags = new Set(ec2Runners.map((runner) => runner.owner));
+  
+  const standbyPoolSize = parseInt(process.env.STANDBY_POOL_SIZE || '0');
+  const environment = process.env.ENVIRONMENT;
 
   for (const ownerTag of ownerTags) {
     const ec2RunnersFiltered = ec2Runners
       .filter((runner) => runner.owner === ownerTag)
       .sort(evictionStrategy === 'oldest_first' ? oldestFirstStrategy : newestFirstStrategy);
     logger.debug(`Found: '${ec2RunnersFiltered.length}' active GitHub runners with owner tag: '${ownerTag}'`);
     logger.debug(`Active GitHub runners with owner tag: '${ownerTag}': ${JSON.stringify(ec2RunnersFiltered)}`);
+    
+    const standbyRunners = standbyPoolSize > 0 
+      ? await listEC2Runners({ environment, runnerType: 'Org', runnerOwner: ownerTag, statuses: ['stopped'], standby: true })
+      : [];
+    const currentStandbyCount = standbyRunners.length;
+    let standbyCounter = Math.max(0, standbyPoolSize - currentStandbyCount);
+    
+    logger.debug(`Standby pool: target=${standbyPoolSize}, current=${currentStandbyCount}, needed=${standbyCounter}`);
+    
     for (const ec2Runner of ec2RunnersFiltered) {
       const ghRunners = await listGitHubRunners(ec2Runner);
       const ghRunnersFiltered = ghRunners.filter((runner: { name: string }) =>
@@ -201,6 +226,16 @@ async function evaluateAndRemoveRunners(
           if (idleCounter > 0) {
             idleCounter--;
             logger.info(`Runner '${ec2Runner.instanceId}' will be kept idle.`);
+          } else if (standbyCounter > 0 && runnerIdleTimeExceeded(ec2Runner)) {
+            standbyCounter--;
+            logger.info(`Runner '${ec2Runner.instanceId}' will be moved to standby.`);
+            await removeRunner(
+              ec2Runner,
+              ghRunnersFiltered.map((runner: { id: number }) => runner.id),
+              true,
+            );
+          } else if (standbyCounter > 0) {
+            logger.info(`Runner '${ec2Runner.instanceId}' waiting for idle time before moving to standby.`);
           } else {
             logger.info(`Terminating all non busy runners.`);
             await removeRunner(