feat(runner): add source parameter to distinguish between scale-up and pool lambda (#5054)

Author: stuartp44

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

@@ -1,4 +1,5 @@
 import { DefaultTargetCapacityType, SpotAllocationStrategy } from '@aws-sdk/client-ec2';
+import { LambdaRunnerSource } from '../scale-runners/scale-up';
 
 export type RunnerType = 'Org' | 'Repo';
 
@@ -42,6 +43,7 @@ export interface RunnerInputParameters {
     instanceAllocationStrategy: SpotAllocationStrategy;
   };
   numberOfRunners: number;
+  source: LambdaRunnerSource;
   amiIdSsmParameterName?: string;
   tracingEnabled?: boolean;
   onDemandFailoverOnError?: string[];

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

@@ -21,6 +21,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest';
 import ScaleError from './../scale-runners/ScaleError';
 import { createRunner, listEC2Runners, tag, terminateRunner, untag } from './runners';
 import type { RunnerInfo, RunnerInputParameters, RunnerType } from './runners.d';
+import { LambdaRunnerSource } from '../scale-runners/scale-up';
 
 process.env.AWS_REGION = 'eu-east-1';
 const mockEC2Client = mockClient(EC2Client);
@@ -318,13 +319,16 @@ describe('create runner', () => {
     allocationStrategy: SpotAllocationStrategy.CAPACITY_OPTIMIZED,
     capacityType: 'spot',
     type: 'Org',
+    scaleErrors: [],
+    source: 'scale-up-lambda',
   };
 
   const defaultExpectedFleetRequestValues: ExpectedFleetRequestValues = {
     type: 'Org',
     capacityType: 'spot',
     allocationStrategy: SpotAllocationStrategy.CAPACITY_OPTIMIZED,
     totalTargetCapacity: 1,
+    source: 'scale-up-lambda',
   };
 
   beforeEach(() => {
@@ -365,6 +369,25 @@ describe('create runner', () => {
     });
   });
 
+  it('calls create fleet of multiple instances with pool-lambda source when specified', async () => {
+    const instances = [{ InstanceIds: ['i-1234', 'i-5678', 'i-9012'] }];
+
+    mockEC2Client.on(CreateFleetCommand).resolves({ Instances: instances });
+
+    await createRunner({
+      ...createRunnerConfig({ ...defaultRunnerConfig, source: 'pool-lambda' }),
+      numberOfRunners: 3,
+    });
+
+    expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, {
+      ...expectedCreateFleetRequest({
+        ...defaultExpectedFleetRequestValues,
+        totalTargetCapacity: 3,
+        source: 'pool-lambda',
+      }),
+    });
+  });
+
   it('calls create fleet of 1 instance with the on-demand capacity', async () => {
     await createRunner(createRunnerConfig({ ...defaultRunnerConfig, capacityType: 'on-demand' }));
     expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, {
@@ -425,6 +448,28 @@ describe('create runner', () => {
       }),
     });
   });
+
+  it('calls create fleet with source set to scale-up-lambda when source is specified', async () => {
+    await createRunner(createRunnerConfig({ ...defaultRunnerConfig, source: 'scale-up-lambda' }));
+
+    expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, {
+      ...expectedCreateFleetRequest({
+        ...defaultExpectedFleetRequestValues,
+        source: 'scale-up-lambda',
+      }),
+    });
+  });
+
+  it('calls create fleet with source set to pool-lambda when source is specified', async () => {
+    await createRunner(createRunnerConfig({ ...defaultRunnerConfig, source: 'pool-lambda' }));
+
+    expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, {
+      ...expectedCreateFleetRequest({
+        ...defaultExpectedFleetRequestValues,
+        source: 'pool-lambda',
+      }),
+    });
+  });
 });
 
 describe('create runner with errors', () => {
@@ -433,12 +478,14 @@ describe('create runner with errors', () => {
     capacityType: 'spot',
     type: 'Repo',
     scaleErrors: ['UnfulfillableCapacity', 'MaxSpotInstanceCountExceeded'],
+    source: 'scale-up-lambda',
   };
   const defaultExpectedFleetRequestValues: ExpectedFleetRequestValues = {
     type: 'Repo',
     capacityType: 'spot',
     allocationStrategy: SpotAllocationStrategy.CAPACITY_OPTIMIZED,
     totalTargetCapacity: 1,
+    source: 'scale-up-lambda',
   };
   beforeEach(() => {
     vi.clearAllMocks();
@@ -546,12 +593,15 @@ describe('create runner with errors fail over to OnDemand', () => {
     capacityType: 'spot',
     type: 'Repo',
     onDemandFailoverOnError: ['InsufficientInstanceCapacity'],
+    scaleErrors: [],
+    source: 'scale-up-lambda',
   };
   const defaultExpectedFleetRequestValues: ExpectedFleetRequestValues = {
     type: 'Repo',
     capacityType: 'spot',
     allocationStrategy: SpotAllocationStrategy.CAPACITY_OPTIMIZED,
     totalTargetCapacity: 1,
+    source: 'scale-up-lambda',
   };
   beforeEach(() => {
     vi.clearAllMocks();
@@ -704,6 +754,7 @@ interface RunnerConfig {
   tracingEnabled?: boolean;
   onDemandFailoverOnError?: string[];
   scaleErrors: string[];
+  source: LambdaRunnerSource;
 }
 
 function createRunnerConfig(runnerConfig: RunnerConfig): RunnerInputParameters {
@@ -724,6 +775,7 @@ function createRunnerConfig(runnerConfig: RunnerConfig): RunnerInputParameters {
     tracingEnabled: runnerConfig.tracingEnabled,
     onDemandFailoverOnError: runnerConfig.onDemandFailoverOnError,
     scaleErrors: runnerConfig.scaleErrors,
+    source: runnerConfig.source,
   };
 }
 
@@ -735,14 +787,15 @@ interface ExpectedFleetRequestValues {
   totalTargetCapacity: number;
   imageId?: string;
   tracingEnabled?: boolean;
+  source: LambdaRunnerSource;
 }
 
 function expectedCreateFleetRequest(expectedValues: ExpectedFleetRequestValues): CreateFleetCommandInput {
   const tags = [
     { Key: 'ghr:Application', Value: 'github-action-runner' },
     {
       Key: 'ghr:created_by',
-      Value: expectedValues.totalTargetCapacity > 1 ? 'pool-lambda' : 'scale-up-lambda',
+      Value: expectedValues.source,
     },
     { Key: 'ghr:Type', Value: expectedValues.type },
     { Key: 'ghr:Owner', Value: REPO_NAME },

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

@@ -241,7 +241,7 @@ async function createInstances(
 ) {
   const tags = [
     { Key: 'ghr:Application', Value: 'github-action-runner' },
-    { Key: 'ghr:created_by', Value: runnerParameters.numberOfRunners === 1 ? 'scale-up-lambda' : 'pool-lambda' },
+    { Key: 'ghr:created_by', Value: runnerParameters.source },
     { Key: 'ghr:Type', Value: runnerParameters.runnerType },
     { Key: 'ghr:Owner', Value: runnerParameters.runnerOwner },
   ];

lambdas/functions/control-plane/src/pool/pool.test.ts

@@ -192,7 +192,13 @@ describe('Test simple pool.', () => {
     it('Top up pool with pool size 2 registered.', async () => {
       await adjust({ poolSize: 3 });
       expect(createRunners).toHaveBeenCalledTimes(1);
-      expect(createRunners).toHaveBeenCalledWith(expect.anything(), expect.anything(), 1, expect.anything());
+      expect(createRunners).toHaveBeenCalledWith(
+        expect.anything(),
+        expect.anything(),
+        1,
+        expect.anything(),
+        'pool-lambda',
+      );
     });
 
     it('Should not top up if pool size is reached.', async () => {
@@ -268,7 +274,13 @@ describe('Test simple pool.', () => {
     it('Top up if the pool size is set to 5', async () => {
       await adjust({ poolSize: 5 });
       // 2 idle, top up with 3 to match a pool of 5
-      expect(createRunners).toHaveBeenCalledWith(expect.anything(), expect.anything(), 3, expect.anything());
+      expect(createRunners).toHaveBeenCalledWith(
+        expect.anything(),
+        expect.anything(),
+        3,
+        expect.anything(),
+        'pool-lambda',
+      );
     });
   });
 
@@ -283,7 +295,13 @@ describe('Test simple pool.', () => {
     it('Top up if the pool size is set to 5', async () => {
       await adjust({ poolSize: 5 });
       // 2 idle, top up with 3 to match a pool of 5
-      expect(createRunners).toHaveBeenCalledWith(expect.anything(), expect.anything(), 3, expect.anything());
+      expect(createRunners).toHaveBeenCalledWith(
+        expect.anything(),
+        expect.anything(),
+        3,
+        expect.anything(),
+        'pool-lambda',
+      );
     });
   });
 
@@ -333,7 +351,13 @@ describe('Test simple pool.', () => {
 
       await adjust({ poolSize: 5 });
       // 2 idle, 2 prefixed idle top up with 1 to match a pool of 5
-      expect(createRunners).toHaveBeenCalledWith(expect.anything(), expect.anything(), 1, expect.anything());
+      expect(createRunners).toHaveBeenCalledWith(
+        expect.anything(),
+        expect.anything(),
+        1,
+        expect.anything(),
+        'pool-lambda',
+      );
     });
   });
 });

lambdas/functions/control-plane/src/pool/pool.ts

@@ -106,6 +106,7 @@ export async function adjust(event: PoolEvent): Promise<void> {
       },
       topUp,
       githubInstallationClient,
+      'pool-lambda',
     );
   } else {
     logger.info(`Pool will not be topped up. Found ${numberOfRunnersInPool} managed idle runners.`);

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

@@ -113,6 +113,7 @@ const EXPECTED_RUNNER_PARAMS: RunnerInputParameters = {
   tracingEnabled: false,
   onDemandFailoverOnError: [],
   scaleErrors: ['UnfulfillableCapacity', 'MaxSpotInstanceCountExceeded', 'TargetCapacityLimitExceededException'],
+  source: 'scale-up-lambda',
 };
 let expectedRunnerParams: RunnerInputParameters;
 

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

@@ -11,6 +11,8 @@ import { publishRetryMessage } from './job-retry';
 
 const logger = createChildLogger('scale-up');
 
+export type LambdaRunnerSource = 'scale-up-lambda' | 'pool-lambda';
+
 export interface RunnerGroup {
   name: string;
   id: number;
@@ -248,11 +250,13 @@ export async function createRunners(
   ec2RunnerConfig: CreateEC2RunnerConfig,
   numberOfRunners: number,
   ghClient: Octokit,
+  source: LambdaRunnerSource = 'scale-up-lambda',
 ): Promise<string[]> {
   const instances = await createRunner({
     runnerType: githubRunnerConfig.runnerType,
     runnerOwner: githubRunnerConfig.runnerOwner,
     numberOfRunners,
+    source,
     ...ec2RunnerConfig,
   });
   if (instances.length !== 0) {
@@ -507,6 +511,7 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise<stri
       },
       newRunners,
       githubInstallationClient,
+      'scale-up-lambda',
     );
 
     // Not all runners we wanted were created, let's reject enough items so that