From fe3ac43fbc1d3bc1e765a3e1c5368843d21fd52a Mon Sep 17 00:00:00 2001 From: edersonbrilhante Date: Wed, 26 Aug 2026 14:31:08 +0200 Subject: [PATCH 1/5] refactor(compute-providers): reuse EC2 runner clients --- .../aws/ec2/control-plane.ts | 21 ++- .../ec2/src/control-plane/dynamic-labels.ts | 9 +- .../aws/ec2/src/control-plane/pool.test.ts | 21 ++- .../aws/ec2/src/control-plane/pool.ts | 20 ++- .../ec2/src/control-plane/runner-config.ts | 25 +-- .../ec2/src/control-plane/scale-down.test.ts | 12 +- .../aws/ec2/src/control-plane/scale-down.ts | 28 ++-- .../ec2/src/control-plane/scale-up.test.ts | 10 +- .../aws/ec2/src/control-plane/scale-up.ts | 23 ++- .../aws/ec2/src/runners.test.ts | 29 +++- .../compute-providers/aws/ec2/src/runners.ts | 155 +++++++++++++++--- 11 files changed, 271 insertions(+), 82 deletions(-) diff --git a/lambdas/libs/compute-providers/aws/ec2/control-plane.ts b/lambdas/libs/compute-providers/aws/ec2/control-plane.ts index 6b4fb7c0bf..08d62a618e 100644 --- a/lambdas/libs/compute-providers/aws/ec2/control-plane.ts +++ b/lambdas/libs/compute-providers/aws/ec2/control-plane.ts @@ -1,25 +1,38 @@ import type { CreateStartRunnerConfig, ComputeProviderPlugin } from '../../core'; +import { getTracedAWSV3Client } from '@aws-github-runner/aws-powertools-util'; +import { EC2Client } from '@aws-sdk/client-ec2'; import type { ControlPlaneProviderCapabilities, ControlPlaneProviderModule } from '../../contracts'; import type {} from './src/environment'; import { createEc2PoolProvider } from './src/control-plane/pool'; import { createEc2ScaleDownProvider } from './src/control-plane/scale-down'; import { createEc2ScaleUpProvider } from './src/control-plane/scale-up'; +import { createEc2RunnerClient, type Ec2RunnerOperations } from './src/runners'; export function createEc2ControlPlanePlugin( createStartRunnerConfig: CreateStartRunnerConfig, + runners: Ec2RunnerOperations, + ec2Client: EC2Client, ): ComputeProviderPlugin { return { type: 'ec2', capabilities: { - pool: () => createEc2PoolProvider(createStartRunnerConfig), - scaleUp: () => createEc2ScaleUpProvider(createStartRunnerConfig), - scaleDown: createEc2ScaleDownProvider, + pool: () => createEc2PoolProvider(createStartRunnerConfig, runners), + scaleUp: () => createEc2ScaleUpProvider(createStartRunnerConfig, runners, ec2Client), + scaleDown: () => createEc2ScaleDownProvider(runners), }, }; } +function createProductionEc2ControlPlanePlugin( + createStartRunnerConfig: CreateStartRunnerConfig, +): ComputeProviderPlugin { + const ec2Client = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION })); + const runners = createEc2RunnerClient(ec2Client).forRequest({ signal: undefined }); + return createEc2ControlPlanePlugin(createStartRunnerConfig, runners, ec2Client); +} + export const provider = { type: 'ec2', - createPlugin: createEc2ControlPlanePlugin, + createPlugin: createProductionEc2ControlPlanePlugin, } satisfies ControlPlaneProviderModule<'ec2'>; diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/dynamic-labels.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/dynamic-labels.ts index 8de7882e24..a28c12ca97 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/dynamic-labels.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/dynamic-labels.ts @@ -12,7 +12,7 @@ import { CpuManufacturer, CpuPerformanceFactorRequest, DescribeLaunchTemplateVersionsCommand, - EC2Client, + type EC2Client, FleetBlockDeviceMappingRequest, FleetEbsBlockDeviceRequest, InstanceGeneration, @@ -30,7 +30,6 @@ import { VCpuCountRangeRequest, VolumeType, } from '@aws-sdk/client-ec2'; -import { getTracedAWSV3Client } from '@aws-github-runner/aws-powertools-util'; import { Ec2OverrideConfig } from '../runners.d'; @@ -366,8 +365,10 @@ export function shouldLoadLaunchTemplateBlockDeviceName(labels: string[]): boole return hasBlockDeviceOverride && !hasBlockDeviceName; } -export async function getDefaultBlockDeviceNameFromLaunchTemplate(launchTemplateName: string): Promise { - const ec2Client = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION })); +export async function getDefaultBlockDeviceNameFromLaunchTemplate( + ec2Client: EC2Client, + launchTemplateName: string, +): Promise { const launchTemplateVersions = await ec2Client.send( new DescribeLaunchTemplateVersionsCommand({ LaunchTemplateName: launchTemplateName, diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.test.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.test.ts index 5667cf9bc0..82c1728d04 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.test.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.test.ts @@ -1,13 +1,12 @@ import type { Octokit } from '@octokit/rest'; import type { CreateGitHubRunnerConfig, CreateStartRunnerConfig, RunnerInfo } from '../../../../core'; -import { bootTimeExceeded, listEC2Runners } from '../runners'; +import { bootTimeExceeded, type Ec2RunnerOperations } from '../runners'; import { calculateEc2PoolSize, createEc2PoolProvider } from './pool'; import { createRunners, type Ec2ProviderConfig, loadEc2ProviderConfig } from './runner-config'; import { beforeEach, describe, expect, it, vi } from 'vitest'; vi.mock('../runners', () => ({ bootTimeExceeded: vi.fn(), - listEC2Runners: vi.fn(), })); vi.mock('./runner-config', () => ({ @@ -16,10 +15,17 @@ vi.mock('./runner-config', () => ({ })); const mockBootTimeExceeded = vi.mocked(bootTimeExceeded); -const mockListRunners = vi.mocked(listEC2Runners); const mockCreateRunners = vi.mocked(createRunners); const mockLoadProviderConfig = vi.mocked(loadEc2ProviderConfig); +const runnerOperations = { + list: vi.fn(), + create: vi.fn(), + terminate: vi.fn(), + tag: vi.fn(), + untag: vi.fn(), +} satisfies Ec2RunnerOperations; + describe('calculateEc2PoolSize', () => { beforeEach(() => { vi.clearAllMocks(); @@ -107,8 +113,8 @@ describe('createEc2PoolProvider', () => { it('lists only running instances managed for the requested pool', async () => { const runners: RunnerInfo[] = [{ id: 'i-running', owner: 'owner', type: 'Org' }]; - mockListRunners.mockResolvedValue(runners); - const provider = createEc2PoolProvider(createStartRunnerConfig); + runnerOperations.list.mockResolvedValue(runners); + const provider = createEc2PoolProvider(createStartRunnerConfig, runnerOperations); await expect( provider.listRunners({ @@ -117,7 +123,7 @@ describe('createEc2PoolProvider', () => { runnerType: 'Org', }), ).resolves.toBe(runners); - expect(mockListRunners).toHaveBeenCalledWith({ + expect(runnerOperations.list).toHaveBeenCalledWith({ environment: 'test-environment', runnerOwner: 'owner', runnerType: 'Org', @@ -131,7 +137,7 @@ describe('createEc2PoolProvider', () => { retryableErrorCount: 0, nonRetryableErrorCount: 0, }); - const provider = createEc2PoolProvider(createStartRunnerConfig); + const provider = createEc2PoolProvider(createStartRunnerConfig, runnerOperations); await expect( provider.createRunners({ @@ -141,6 +147,7 @@ describe('createEc2PoolProvider', () => { }), ).resolves.toEqual(['i-created']); expect(mockCreateRunners).toHaveBeenCalledWith( + runnerOperations, githubRunnerConfig, providerConfig, 1, diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.ts index d1d243ca51..498a7ddbd1 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.ts @@ -7,17 +7,16 @@ import type { RunnerInfo, RunnerStatus, } from '../../../../core'; +import { bootTimeExceeded, type Ec2RunnerOperations } from '../runners'; import { createRunners, loadEc2ProviderConfig } from './runner-config'; -import { bootTimeExceeded, listEC2Runners } from '../runners'; const logger = createChildLogger('pool'); -async function listEc2PoolRunners({ - environment, - runnerOwner, - runnerType, -}: ListPoolRunnersInput): Promise { - return await listEC2Runners({ +async function listEc2PoolRunners( + runners: Ec2RunnerOperations, + { environment, runnerOwner, runnerType }: ListPoolRunnersInput, +): Promise { + return await runners.list({ environment, runnerOwner, runnerType, @@ -28,10 +27,12 @@ async function listEc2PoolRunners({ async function createEc2PoolRunners( { githubRunnerConfig, numberOfRunners, githubInstallationClient }: CreatePoolRunnersInput, createStartRunnerConfig: CreateStartRunnerConfig, + runners: Ec2RunnerOperations, ): Promise { const config = loadEc2ProviderConfig(); const { instances } = await createRunners( + runners, githubRunnerConfig, { ec2instanceCriteria: config.ec2instanceCriteria, @@ -53,11 +54,12 @@ async function createEc2PoolRunners( export function createEc2PoolProvider( createStartRunnerConfig: CreateStartRunnerConfig, + runners: Ec2RunnerOperations, ): Omit, 'type'> { return { - listRunners: listEc2PoolRunners, + listRunners: (input) => listEc2PoolRunners(runners, input), countAvailableRunners: calculateEc2PoolSize, - createRunners: (input) => createEc2PoolRunners(input, createStartRunnerConfig), + createRunners: (input) => createEc2PoolRunners(input, createStartRunnerConfig, runners), }; } diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-config.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-config.ts index 34f95aaf97..6f009f9e5f 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-config.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-config.ts @@ -11,7 +11,7 @@ import { Octokit } from '@octokit/rest'; import type { Tag } from '@aws-sdk/client-ec2'; import yn from 'yn'; -import { createRunner, tag, terminateRunner } from '../runners'; +import type { Ec2RunnerOperations } from '../runners'; import type { RunnerInputParameters } from '../runners.d'; const logger = createChildLogger('ec2-runners'); @@ -61,6 +61,7 @@ export function loadEc2ProviderConfig(): Ec2ProviderConfig { } export async function createRunners( + runners: Ec2RunnerOperations, githubRunnerConfig: CreateGitHubRunnerConfig, ec2RunnerConfig: CreateEC2RunnerConfig, numberOfRunners: number, @@ -70,7 +71,7 @@ export async function createRunners( ): Promise { let result: CreateRunnerResult; try { - result = await createRunner({ + result = await runners.create({ runnerType: githubRunnerConfig.runnerType, runnerOwner: githubRunnerConfig.runnerOwner, numberOfRunners, @@ -93,7 +94,7 @@ export async function createRunners( githubRunnerConfig, result.instances, ghClient, - createEc2StartRunnerConfigOptions(), + createEc2StartRunnerConfigOptions(runners), ); } catch (error) { logger.error('Unexpected error while registering GitHub runners.', { @@ -113,7 +114,7 @@ export async function createRunners( retryable: true, }); - await terminateFailedInstances(failedInstances); + await terminateFailedInstances(runners, failedInstances); return { instances: result.instances.filter((id) => !failedInstances.includes(id)), @@ -126,10 +127,10 @@ export async function createRunners( return result; } -async function terminateFailedInstances(instanceIds: string[]): Promise { +async function terminateFailedInstances(runners: Ec2RunnerOperations, instanceIds: string[]): Promise { for (const instanceId of instanceIds) { try { - await terminateRunner(instanceId); + await runners.terminate(instanceId); } catch (error) { logger.error('Failed to terminate instance', { instanceId, @@ -139,21 +140,25 @@ async function terminateFailedInstances(instanceIds: string[]): Promise { } } -function createEc2StartRunnerConfigOptions(): StartRunnerConfigOptions { +function createEc2StartRunnerConfigOptions(runners: Ec2RunnerOperations): StartRunnerConfigOptions { return { getSsmParameterTags: (instanceId) => [{ Key: 'InstanceId', Value: instanceId }], - onJitConfigCreated: async (instanceId, metadata) => await tagEc2RunnerMetadata(instanceId, metadata), + onJitConfigCreated: async (instanceId, metadata) => await tagEc2RunnerMetadata(runners, instanceId, metadata), }; } -async function tagEc2RunnerMetadata(instanceId: string, metadata: GitHubRunnerMetadata): Promise { +async function tagEc2RunnerMetadata( + runners: Ec2RunnerOperations, + instanceId: string, + metadata: GitHubRunnerMetadata, +): Promise { const tags = [ { Key: 'ghr:github_runner_id', Value: metadata.githubRunnerId }, ...generateRunnerLabelsTags(metadata.runnerLabels), ]; try { - await tag(instanceId, tags); + await runners.tag(instanceId, tags); } catch (e) { logger.error(`Failed to mark EC2 runner '${instanceId}' with GitHub runner metadata.`, { error: e }); } diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.test.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.test.ts index f122f3b3a0..96f2115c3c 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.test.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.test.ts @@ -3,6 +3,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; import type { RunnerInfo, RunnerType } from '../../../../core'; import { createEc2ScaleDownProvider } from './scale-down'; import { listEC2Runners, tag, terminateRunner, untag } from '../runners'; +import type { Ec2RunnerOperations } from '../runners'; vi.mock('../runners', async (importOriginal) => { const actual = await importOriginal(); @@ -19,6 +20,13 @@ const mockListRunners = vi.mocked(listEC2Runners); const mockTagRunner = vi.mocked(tag); const mockTerminateRunner = vi.mocked(terminateRunner); const mockUntagRunner = vi.mocked(untag); +const runnerOperations: Ec2RunnerOperations = { + list: mockListRunners, + create: vi.fn(), + terminate: mockTerminateRunner, + tag: mockTagRunner, + untag: mockUntagRunner, +}; describe('Scale down runners', () => { beforeEach(() => { @@ -47,7 +55,7 @@ describe('Scale down runners', () => { mockListRunners.mockResolvedValueOnce([]).mockResolvedValueOnce([runner]); mockTagRunner.mockResolvedValue(); mockUntagRunner.mockResolvedValue(); - const provider = createEc2ScaleDownProvider(); + const provider = createEc2ScaleDownProvider(runnerOperations); await expect(provider.list('unit-test-environment')).resolves.toEqual([]); await expect(provider.list('unit-test-environment', true)).resolves.toEqual([runner]); @@ -71,7 +79,7 @@ describe('Scale down runners', () => { launchTime: new Date(), }; process.env.RUNNER_BOOT_TIME_IN_MINUTES = '5'; - const provider = createEc2ScaleDownProvider(); + const provider = createEc2ScaleDownProvider(runnerOperations); expect(provider.bootTimeExceeded(scaleDownRunner)).toBe(false); expect(mockTerminateRunner).not.toHaveBeenCalled(); diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.ts index 3084eebe21..15e2ba04b9 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.ts @@ -1,24 +1,28 @@ import type { RunnerInfo, ScaleDownComputeProvider } from '../../../../core'; -import { bootTimeExceeded, listEC2Runners, tag, terminateRunner, untag } from '../runners'; +import { bootTimeExceeded, type Ec2RunnerOperations } from '../runners'; -async function listEc2ScaleDownRunners(environment: string, orphan?: boolean): Promise { - return await listEC2Runners({ environment, orphan }); +async function listEc2ScaleDownRunners( + runners: Ec2RunnerOperations, + environment: string, + orphan?: boolean, +): Promise { + return await runners.list({ environment, orphan }); } -async function markEc2RunnerOrphan(id: string): Promise { - await tag(id, [{ Key: 'ghr:orphan', Value: 'true' }]); +async function markEc2RunnerOrphan(runners: Ec2RunnerOperations, id: string): Promise { + await runners.tag(id, [{ Key: 'ghr:orphan', Value: 'true' }]); } -async function unmarkEc2RunnerOrphan(id: string): Promise { - await untag(id, [{ Key: 'ghr:orphan', Value: 'true' }]); +async function unmarkEc2RunnerOrphan(runners: Ec2RunnerOperations, id: string): Promise { + await runners.untag(id, [{ Key: 'ghr:orphan', Value: 'true' }]); } -export function createEc2ScaleDownProvider(): Omit { +export function createEc2ScaleDownProvider(runners: Ec2RunnerOperations): Omit { return { - list: listEc2ScaleDownRunners, + list: (environment, orphan) => listEc2ScaleDownRunners(runners, environment, orphan), bootTimeExceeded, - markOrphan: markEc2RunnerOrphan, - unmarkOrphan: unmarkEc2RunnerOrphan, - terminate: terminateRunner, + markOrphan: (id) => markEc2RunnerOrphan(runners, id), + unmarkOrphan: (id) => unmarkEc2RunnerOrphan(runners, id), + terminate: (id) => runners.terminate(id), }; } diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts index 1a967fc4cc..a41f143241 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts @@ -8,6 +8,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; import { parseEc2OverrideConfig } from './dynamic-labels'; import { EC2_TAG_VALUE_MAX_LENGTH, RUNNER_LABELS_TAG_MAX_COUNT } from './runner-config'; import { createRunner, listEC2Runners, tag, terminateRunner } from '../runners'; +import type { Ec2RunnerOperations } from '../runners'; import type { RunnerInputParameters } from '../runners.d'; import { createEc2ScaleUpProvider } from './scale-up'; @@ -29,7 +30,14 @@ const githubClient = {} as Octokit; const runnerOwner = 'Codertocat'; const repositoryRunnerOwner = 'Codertocat/hello-world'; const cleanEnv = process.env; -const provider = createEc2ScaleUpProvider(mockCreateStartRunnerConfig); +const runnerOperations: Ec2RunnerOperations = { + list: mockListRunners, + create: mockCreateRunner, + terminate: mockTerminateRunner, + tag: mockTag, + untag: vi.fn(), +}; +const provider = createEc2ScaleUpProvider(mockCreateStartRunnerConfig, runnerOperations, new EC2Client({})); interface CreateProviderRunnersOptions { labels?: string[]; diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.ts index 6bebaa31ba..8f2c2457e0 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.ts @@ -1,4 +1,5 @@ import { createChildLogger } from '@aws-github-runner/aws-powertools-util'; +import type { EC2Client } from '@aws-sdk/client-ec2'; import type { CreateRunnerResult, CreateScaleUpRunnersInput, @@ -9,7 +10,7 @@ import type { } from '../../../../core'; import yn from 'yn'; -import { listEC2Runners } from '../runners'; +import type { Ec2RunnerOperations } from '../runners'; import type { Ec2OverrideConfig } from '../runners.d'; import { getDefaultBlockDeviceNameFromLaunchTemplate, @@ -32,7 +33,10 @@ function loadEc2ScaleUpProviderConfig(): CreateEC2RunnerConfig { }; } -async function resolveEc2LabelsForRunners(messageLabels: string[]): Promise> { +async function resolveEc2LabelsForRunners( + ec2Client: EC2Client, + messageLabels: string[], +): Promise> { const trimmedLabels = messageLabels.map((label) => label.trim()); const dynamicEC2Labels = trimmedLabels.filter((label) => label.startsWith('ghr-ec2-')); const nonEc2DynamicLabels = trimmedLabels.filter( @@ -43,7 +47,7 @@ async function resolveEc2LabelsForRunners(messageLabels: string[]): Promise 0) { const defaultBlockDeviceName = shouldLoadLaunchTemplateBlockDeviceName(dynamicEC2Labels) - ? await getDefaultBlockDeviceNameFromLaunchTemplate(process.env.LAUNCH_TEMPLATE_NAME) + ? await getDefaultBlockDeviceNameFromLaunchTemplate(ec2Client, process.env.LAUNCH_TEMPLATE_NAME) : undefined; ec2OverrideConfig = parseEc2OverrideConfig(dynamicEC2Labels, defaultBlockDeviceName); @@ -56,19 +60,22 @@ async function resolveEc2LabelsForRunners(messageLabels: string[]): Promise { - return (await listEC2Runners({ environment: process.env.ENVIRONMENT, runnerType, runnerOwner })).length; + return (await runners.list({ environment: process.env.ENVIRONMENT, runnerType, runnerOwner })).length; } async function createEc2ScaleUpRunners( { githubRunnerConfig, numberOfRunners, githubInstallationClient, state }: CreateScaleUpRunnersInput, createStartRunnerConfig: CreateStartRunnerConfig, + runners: Ec2RunnerOperations, ): Promise { const config = loadEc2ScaleUpProviderConfig(); return await createRunners( + runners, githubRunnerConfig, { ...config, @@ -83,10 +90,12 @@ async function createEc2ScaleUpRunners( export function createEc2ScaleUpProvider( createStartRunnerConfig: CreateStartRunnerConfig, + runners: Ec2RunnerOperations, + ec2Client: EC2Client, ): Omit, 'type'> { return { - resolveLabelsForRunners: resolveEc2LabelsForRunners, - getCurrentRunners: getCurrentEc2Runners, - createRunners: (input) => createEc2ScaleUpRunners(input, createStartRunnerConfig), + resolveLabelsForRunners: (labels) => resolveEc2LabelsForRunners(ec2Client, labels), + getCurrentRunners: (state, input) => getCurrentEc2Runners(runners, state, input), + createRunners: (input) => createEc2ScaleUpRunners(input, createStartRunnerConfig, runners), }; } diff --git a/lambdas/libs/compute-providers/aws/ec2/src/runners.test.ts b/lambdas/libs/compute-providers/aws/ec2/src/runners.test.ts index affd063c39..1520d86a22 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/runners.test.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/runners.test.ts @@ -22,7 +22,7 @@ import 'aws-sdk-client-mock-jest/vitest'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import type { RunnerInfo, RunnerSource, RunnerType } from '../../../core'; -import { createRunner, listEC2Runners, tag, terminateRunner, untag } from './runners'; +import { createEc2RunnerClient, createRunner, listEC2Runners, tag, terminateRunner, untag } from './runners'; import type { Ec2OverrideConfig, RunnerInputParameters } from './runners.d'; process.env.AWS_REGION = 'eu-east-1'; @@ -502,6 +502,33 @@ describe('create runner', () => { Name: 'my-ami-id-param', }); }); + + it('keeps cancellation request-scoped and rejects before calling AWS', async () => { + const abortController = new AbortController(); + const abortReason = new Error('service stopping'); + const runners = createEc2RunnerClient(new EC2Client({})).forRequest({ + signal: abortController.signal, + }); + abortController.abort(abortReason); + + await expect(runners.create(createRunnerConfig(defaultRunnerConfig))).rejects.toThrow('service stopping'); + expect(mockEC2Client).not.toHaveReceivedCommand(CreateFleetCommand); + expect(mockSSMClient).not.toHaveReceivedCommand(GetParameterCommand); + }); + + it('keeps another request usable when a request sharing the same clients is cancelled', async () => { + const abortController = new AbortController(); + const client = createEc2RunnerClient(new EC2Client({})); + const cancelledRequest = client.forRequest({ signal: abortController.signal }); + const activeRequest = client.forRequest({ signal: undefined }); + mockEC2Client.on(DescribeInstancesCommand).resolves({ Reservations: [] }); + abortController.abort(new Error('service stopping')); + + await expect(cancelledRequest.list()).rejects.toThrow('service stopping'); + await expect(activeRequest.list()).resolves.toEqual([]); + expect(mockEC2Client).toHaveReceivedCommandTimes(DescribeInstancesCommand, 1); + }); + it('calls create fleet of 1 instance with runner tracing enabled', async () => { tracer.getRootXrayTraceId = vi.fn().mockReturnValue('123'); diff --git a/lambdas/libs/compute-providers/aws/ec2/src/runners.ts b/lambdas/libs/compute-providers/aws/ec2/src/runners.ts index 4e187437ec..69f6eec76c 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/runners.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/runners.ts @@ -32,13 +32,68 @@ interface Ec2Filter { Values: string[]; } +export interface Ec2RunnerRequestContext { + readonly signal: AbortSignal | undefined; +} + +export interface Ec2RunnerOperations { + list(filters?: Ec2ListRunnerFilters): Promise; + create(runnerParameters: RunnerInputParameters): Promise; + terminate(instanceId: string): Promise; + tag(instanceId: string, tags: Tag[]): Promise; + untag(instanceId: string, tags: Tag[]): Promise; +} + +export interface Ec2RunnerClient { + forRequest(context: Ec2RunnerRequestContext): Ec2RunnerOperations; +} + +async function runWithRequestSignal( + signal: AbortSignal | undefined, + operation: () => Promise, +): Promise { + signal?.throwIfAborted(); + return await operation(); +} + +export function createEc2RunnerClient(ec2Client: EC2Client): Ec2RunnerClient { + return { + forRequest: ({ signal }) => ({ + list: (filters) => runWithRequestSignal(signal, () => listRunners(ec2Client, filters, signal)), + create: (runnerParameters) => + runWithRequestSignal(signal, () => createEc2Runner(runnerParameters, ec2Client, signal)), + terminate: (instanceId) => runWithRequestSignal(signal, () => terminateEc2Runner(ec2Client, instanceId, signal)), + tag: (instanceId, tags) => runWithRequestSignal(signal, () => tagEc2Runner(ec2Client, instanceId, tags, signal)), + untag: (instanceId, tags) => + runWithRequestSignal(signal, () => untagEc2Runner(ec2Client, instanceId, tags, signal)), + }), + }; +} + +let defaultRunnerOperations: Ec2RunnerOperations | undefined; + +function getDefaultRunnerOperations(): Ec2RunnerOperations { + defaultRunnerOperations ??= createEc2RunnerClient( + getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION })), + ).forRequest({ signal: undefined }); + return defaultRunnerOperations; +} + type FleetError = NonNullable[number]; export async function listEC2Runners(filters: Ec2ListRunnerFilters | undefined = undefined): Promise { + return await getDefaultRunnerOperations().list(filters); +} + +async function listRunners( + ec2Client: EC2Client, + filters: Ec2ListRunnerFilters | undefined, + signal: AbortSignal | undefined, +): Promise { const ec2Filters = constructFilters(filters); const runners: RunnerInfo[] = []; for (const filter of ec2Filters) { - runners.push(...(await getRunners(filter))); + runners.push(...(await getRunners(ec2Client, filter, signal))); } return runners; } @@ -68,14 +123,18 @@ function constructFilters(filters?: Ec2ListRunnerFilters): Ec2Filter[][] { return ec2Filters; } -async function getRunners(ec2Filters: Ec2Filter[]): Promise { - const ec2 = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION })); +async function getRunners( + ec2Client: EC2Client, + ec2Filters: Ec2Filter[], + signal: AbortSignal | undefined, +): Promise { const runners: RunnerInfo[] = []; let nextToken; let hasNext = true; while (hasNext) { - const instances: DescribeInstancesResult = await ec2.send( + const instances: DescribeInstancesResult = await ec2Client.send( new DescribeInstancesCommand({ Filters: ec2Filters, NextToken: nextToken }), + { abortSignal: signal }, ); hasNext = instances.NextToken ? true : false; nextToken = instances.NextToken; @@ -109,22 +168,45 @@ function getRunnerInfo(runningInstances: DescribeInstancesResult) { } export async function terminateRunner(instanceId: string): Promise { + await getDefaultRunnerOperations().terminate(instanceId); +} + +async function terminateEc2Runner( + ec2Client: EC2Client, + instanceId: string, + signal: AbortSignal | undefined, +): Promise { logger.debug(`Runner '${instanceId}' will be terminated.`); - const ec2 = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION })); - await ec2.send(new TerminateInstancesCommand({ InstanceIds: [instanceId] })); + await ec2Client.send(new TerminateInstancesCommand({ InstanceIds: [instanceId] }), { abortSignal: signal }); logger.debug(`Runner ${instanceId} has been terminated.`); } export async function tag(instanceId: string, tags: Tag[]): Promise { + await getDefaultRunnerOperations().tag(instanceId, tags); +} + +async function tagEc2Runner( + ec2Client: EC2Client, + instanceId: string, + tags: Tag[], + signal: AbortSignal | undefined, +): Promise { logger.debug(`Tagging '${instanceId}'`, { tags }); - const ec2 = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION })); - await ec2.send(new CreateTagsCommand({ Resources: [instanceId], Tags: tags })); + await ec2Client.send(new CreateTagsCommand({ Resources: [instanceId], Tags: tags }), { abortSignal: signal }); } export async function untag(instanceId: string, tags: Tag[]): Promise { + await getDefaultRunnerOperations().untag(instanceId, tags); +} + +async function untagEc2Runner( + ec2Client: EC2Client, + instanceId: string, + tags: Tag[], + signal: AbortSignal | undefined, +): Promise { logger.debug(`Untagging '${instanceId}'`, { tags }); - const ec2 = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION })); - await ec2.send(new DeleteTagsCommand({ Resources: [instanceId], Tags: tags })); + await ec2Client.send(new DeleteTagsCommand({ Resources: [instanceId], Tags: tags }), { abortSignal: signal }); } const SPOT_ALLOCATION_STRATEGIES = [ @@ -299,6 +381,14 @@ function buildRunInstancesOverrides( } export async function createRunner(runnerParameters: RunnerInputParameters): Promise { + return await getDefaultRunnerOperations().create(runnerParameters); +} + +async function createEc2Runner( + runnerParameters: RunnerInputParameters, + ec2Client: EC2Client, + signal: AbortSignal | undefined, +): Promise { logger.debug('Runner configuration.', { runner: { configuration: { @@ -307,11 +397,11 @@ export async function createRunner(runnerParameters: RunnerInputParameters): Pro }, }); - const ec2Client = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION })); let amiIdOverride: string | undefined; try { amiIdOverride = await getAmiIdOverride(runnerParameters); } catch (error) { + throwIfAborted(signal, error); const retryable = isRetryableAwsError(error, runnerParameters.scaleErrors); logger.warn('Runner creation failed before an EC2 request could be made.', { error: error as Error, @@ -325,21 +415,22 @@ export async function createRunner(runnerParameters: RunnerInputParameters): Pro // for instance types like mac*.metal. Use RunInstances directly instead. if (runnerParameters.useDedicatedHost) { logger.info('Using RunInstances for dedicated host placement (CreateFleet does not support dedicated hosts).'); - const result = await createInstancesWithRunInstances(runnerParameters, amiIdOverride, ec2Client); + const result = await createInstancesWithRunInstances(runnerParameters, amiIdOverride, ec2Client, signal); logger.info(`Created instance(s) via RunInstances: ${result.instances.join(',')}`); return result; } let fleet: CreateFleetResult; try { - fleet = await createInstances(runnerParameters, amiIdOverride, ec2Client); + fleet = await createInstances(runnerParameters, amiIdOverride, ec2Client, signal); } catch (error) { + throwIfAborted(signal, error); const retryable = isRetryableAwsError(error, runnerParameters.scaleErrors); logger.warn('Create fleet request failed.', { error: error as Error, retryable }); return failedCreateRunnerResult(runnerParameters.numberOfRunners, retryable); } - const result = await processFleetResult(fleet, runnerParameters); + const result = await processFleetResult(fleet, runnerParameters, ec2Client, signal); logger.info(`Created instance(s): ${result.instances.join(',')}`); @@ -349,6 +440,8 @@ export async function createRunner(runnerParameters: RunnerInputParameters): Pro async function processFleetResult( fleet: CreateFleetResult, runnerParameters: RunnerInputParameters, + ec2Client: EC2Client, + signal: AbortSignal | undefined, ): Promise { const instances: string[] = fleet.Instances?.flatMap((i) => i.InstanceIds?.flatMap((j) => j) || []) || []; @@ -376,16 +469,20 @@ async function processFleetResult( runnerParameters.ec2instanceCriteria.instanceAllocationStrategy, 'on-demand', ); - const onDemandResult = await createRunner({ - ...runnerParameters, - numberOfRunners: numberOfInstances, - onDemandFailoverOnError: ['InsufficientInstanceCapacity'], - ec2instanceCriteria: { - ...runnerParameters.ec2instanceCriteria, - targetCapacityType: 'on-demand', - instanceAllocationStrategy: failoverAllocationStrategy, + const onDemandResult = await createEc2Runner( + { + ...runnerParameters, + numberOfRunners: numberOfInstances, + onDemandFailoverOnError: ['InsufficientInstanceCapacity'], + ec2instanceCriteria: { + ...runnerParameters.ec2instanceCriteria, + targetCapacityType: 'on-demand', + instanceAllocationStrategy: failoverAllocationStrategy, + }, }, - }); + ec2Client, + signal, + ); instances.push(...onDemandResult.instances); return { instances, @@ -474,6 +571,11 @@ function failedCreateRunnerResult(failedInstanceCount: number, isRetryable: bool }; } +function throwIfAborted(signal: AbortSignal | undefined, error: unknown): void { + if (signal?.aborted) signal.throwIfAborted(); + if (error instanceof Error && error.name === 'AbortError') throw error; +} + async function getAmiIdOverride(runnerParameters: RunnerInputParameters): Promise { if (!runnerParameters.amiIdSsmParameterName) { return undefined; @@ -497,6 +599,7 @@ async function createInstances( runnerParameters: RunnerInputParameters, amiIdOverride: string | undefined, ec2Client: EC2Client, + signal: AbortSignal | undefined, ) { const tags = [ { Key: 'ghr:Application', Value: 'github-action-runner' }, @@ -569,7 +672,7 @@ async function createInstances( Type: 'instant', }); logger.debug('CreateFleet request payload.', { payload: createFleetCommand.input }); - fleet = await ec2Client.send(createFleetCommand); + fleet = await ec2Client.send(createFleetCommand, { abortSignal: signal }); } catch (e) { logger.warn('Create fleet request failed.', { error: e as Error }); throw e; @@ -581,6 +684,7 @@ async function createInstancesWithRunInstances( runnerParameters: RunnerInputParameters, amiIdOverride: string | undefined, ec2Client: EC2Client, + signal: AbortSignal | undefined, ): Promise { const tags = [ { Key: 'ghr:Application', Value: 'github-action-runner' }, @@ -627,9 +731,10 @@ async function createInstancesWithRunInstances( }); logger.debug('RunInstances request payload.', { payload: runInstancesCommand.input }); - const result = await ec2Client.send(runInstancesCommand); + const result = await ec2Client.send(runInstancesCommand, { abortSignal: signal }); return processRunInstanceResult(result, runnerParameters); } catch (error) { + throwIfAborted(signal, error); const retryable = isRetryableAwsError(error, runnerParameters.scaleErrors); logger.warn('RunInstances request failed for dedicated host.', { error: error as Error, retryable }); return failedCreateRunnerResult(runnerParameters.numberOfRunners, retryable); From 55397ca832482ff82356fd809bf1f1d9d63bf6a6 Mon Sep 17 00:00:00 2001 From: edersonbrilhante Date: Thu, 27 Aug 2026 12:36:20 +0200 Subject: [PATCH 2/5] refactor(ec2): clarify runner operations naming --- .../aws/ec2/control-plane.ts | 18 +++++++------- .../aws/ec2/src/control-plane/pool.ts | 14 +++++------ .../ec2/src/control-plane/runner-config.ts | 21 ++++++++-------- .../aws/ec2/src/control-plane/scale-down.ts | 24 ++++++++++--------- .../aws/ec2/src/control-plane/scale-up.ts | 14 +++++------ .../aws/ec2/src/runners.test.ts | 4 ++-- 6 files changed, 49 insertions(+), 46 deletions(-) diff --git a/lambdas/libs/compute-providers/aws/ec2/control-plane.ts b/lambdas/libs/compute-providers/aws/ec2/control-plane.ts index 08d62a618e..be24e0b241 100644 --- a/lambdas/libs/compute-providers/aws/ec2/control-plane.ts +++ b/lambdas/libs/compute-providers/aws/ec2/control-plane.ts @@ -9,30 +9,30 @@ import { createEc2ScaleDownProvider } from './src/control-plane/scale-down'; import { createEc2ScaleUpProvider } from './src/control-plane/scale-up'; import { createEc2RunnerClient, type Ec2RunnerOperations } from './src/runners'; -export function createEc2ControlPlanePlugin( +export function createEc2ControlPlanePluginWithDependencies( createStartRunnerConfig: CreateStartRunnerConfig, - runners: Ec2RunnerOperations, + runnerOperations: Ec2RunnerOperations, ec2Client: EC2Client, ): ComputeProviderPlugin { return { type: 'ec2', capabilities: { - pool: () => createEc2PoolProvider(createStartRunnerConfig, runners), - scaleUp: () => createEc2ScaleUpProvider(createStartRunnerConfig, runners, ec2Client), - scaleDown: () => createEc2ScaleDownProvider(runners), + pool: () => createEc2PoolProvider(createStartRunnerConfig, runnerOperations), + scaleUp: () => createEc2ScaleUpProvider(createStartRunnerConfig, runnerOperations, ec2Client), + scaleDown: () => createEc2ScaleDownProvider(runnerOperations), }, }; } -function createProductionEc2ControlPlanePlugin( +function createEc2ControlPlanePlugin( createStartRunnerConfig: CreateStartRunnerConfig, ): ComputeProviderPlugin { const ec2Client = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION })); - const runners = createEc2RunnerClient(ec2Client).forRequest({ signal: undefined }); - return createEc2ControlPlanePlugin(createStartRunnerConfig, runners, ec2Client); + const runnerOperations = createEc2RunnerClient(ec2Client).forRequest({ signal: undefined }); + return createEc2ControlPlanePluginWithDependencies(createStartRunnerConfig, runnerOperations, ec2Client); } export const provider = { type: 'ec2', - createPlugin: createProductionEc2ControlPlanePlugin, + createPlugin: createEc2ControlPlanePlugin, } satisfies ControlPlaneProviderModule<'ec2'>; diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.ts index 498a7ddbd1..b768efb3ec 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.ts @@ -13,10 +13,10 @@ import { createRunners, loadEc2ProviderConfig } from './runner-config'; const logger = createChildLogger('pool'); async function listEc2PoolRunners( - runners: Ec2RunnerOperations, + runnerOperations: Ec2RunnerOperations, { environment, runnerOwner, runnerType }: ListPoolRunnersInput, ): Promise { - return await runners.list({ + return await runnerOperations.list({ environment, runnerOwner, runnerType, @@ -27,12 +27,12 @@ async function listEc2PoolRunners( async function createEc2PoolRunners( { githubRunnerConfig, numberOfRunners, githubInstallationClient }: CreatePoolRunnersInput, createStartRunnerConfig: CreateStartRunnerConfig, - runners: Ec2RunnerOperations, + runnerOperations: Ec2RunnerOperations, ): Promise { const config = loadEc2ProviderConfig(); const { instances } = await createRunners( - runners, + runnerOperations, githubRunnerConfig, { ec2instanceCriteria: config.ec2instanceCriteria, @@ -54,12 +54,12 @@ async function createEc2PoolRunners( export function createEc2PoolProvider( createStartRunnerConfig: CreateStartRunnerConfig, - runners: Ec2RunnerOperations, + runnerOperations: Ec2RunnerOperations, ): Omit, 'type'> { return { - listRunners: (input) => listEc2PoolRunners(runners, input), + listRunners: (input) => listEc2PoolRunners(runnerOperations, input), countAvailableRunners: calculateEc2PoolSize, - createRunners: (input) => createEc2PoolRunners(input, createStartRunnerConfig, runners), + createRunners: (input) => createEc2PoolRunners(input, createStartRunnerConfig, runnerOperations), }; } diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-config.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-config.ts index 6f009f9e5f..083b7ab56d 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-config.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-config.ts @@ -61,7 +61,7 @@ export function loadEc2ProviderConfig(): Ec2ProviderConfig { } export async function createRunners( - runners: Ec2RunnerOperations, + runnerOperations: Ec2RunnerOperations, githubRunnerConfig: CreateGitHubRunnerConfig, ec2RunnerConfig: CreateEC2RunnerConfig, numberOfRunners: number, @@ -71,7 +71,7 @@ export async function createRunners( ): Promise { let result: CreateRunnerResult; try { - result = await runners.create({ + result = await runnerOperations.create({ runnerType: githubRunnerConfig.runnerType, runnerOwner: githubRunnerConfig.runnerOwner, numberOfRunners, @@ -94,7 +94,7 @@ export async function createRunners( githubRunnerConfig, result.instances, ghClient, - createEc2StartRunnerConfigOptions(runners), + createEc2StartRunnerConfigOptions(runnerOperations), ); } catch (error) { logger.error('Unexpected error while registering GitHub runners.', { @@ -114,7 +114,7 @@ export async function createRunners( retryable: true, }); - await terminateFailedInstances(runners, failedInstances); + await terminateFailedInstances(runnerOperations, failedInstances); return { instances: result.instances.filter((id) => !failedInstances.includes(id)), @@ -127,10 +127,10 @@ export async function createRunners( return result; } -async function terminateFailedInstances(runners: Ec2RunnerOperations, instanceIds: string[]): Promise { +async function terminateFailedInstances(runnerOperations: Ec2RunnerOperations, instanceIds: string[]): Promise { for (const instanceId of instanceIds) { try { - await runners.terminate(instanceId); + await runnerOperations.terminate(instanceId); } catch (error) { logger.error('Failed to terminate instance', { instanceId, @@ -140,15 +140,16 @@ async function terminateFailedInstances(runners: Ec2RunnerOperations, instanceId } } -function createEc2StartRunnerConfigOptions(runners: Ec2RunnerOperations): StartRunnerConfigOptions { +function createEc2StartRunnerConfigOptions(runnerOperations: Ec2RunnerOperations): StartRunnerConfigOptions { return { getSsmParameterTags: (instanceId) => [{ Key: 'InstanceId', Value: instanceId }], - onJitConfigCreated: async (instanceId, metadata) => await tagEc2RunnerMetadata(runners, instanceId, metadata), + onJitConfigCreated: async (instanceId, metadata) => + await tagEc2RunnerMetadata(runnerOperations, instanceId, metadata), }; } async function tagEc2RunnerMetadata( - runners: Ec2RunnerOperations, + runnerOperations: Ec2RunnerOperations, instanceId: string, metadata: GitHubRunnerMetadata, ): Promise { @@ -158,7 +159,7 @@ async function tagEc2RunnerMetadata( ]; try { - await runners.tag(instanceId, tags); + await runnerOperations.tag(instanceId, tags); } catch (e) { logger.error(`Failed to mark EC2 runner '${instanceId}' with GitHub runner metadata.`, { error: e }); } diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.ts index 15e2ba04b9..23381c69cd 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.ts @@ -2,27 +2,29 @@ import type { RunnerInfo, ScaleDownComputeProvider } from '../../../../core'; import { bootTimeExceeded, type Ec2RunnerOperations } from '../runners'; async function listEc2ScaleDownRunners( - runners: Ec2RunnerOperations, + runnerOperations: Ec2RunnerOperations, environment: string, orphan?: boolean, ): Promise { - return await runners.list({ environment, orphan }); + return await runnerOperations.list({ environment, orphan }); } -async function markEc2RunnerOrphan(runners: Ec2RunnerOperations, id: string): Promise { - await runners.tag(id, [{ Key: 'ghr:orphan', Value: 'true' }]); +async function markEc2RunnerOrphan(runnerOperations: Ec2RunnerOperations, id: string): Promise { + await runnerOperations.tag(id, [{ Key: 'ghr:orphan', Value: 'true' }]); } -async function unmarkEc2RunnerOrphan(runners: Ec2RunnerOperations, id: string): Promise { - await runners.untag(id, [{ Key: 'ghr:orphan', Value: 'true' }]); +async function unmarkEc2RunnerOrphan(runnerOperations: Ec2RunnerOperations, id: string): Promise { + await runnerOperations.untag(id, [{ Key: 'ghr:orphan', Value: 'true' }]); } -export function createEc2ScaleDownProvider(runners: Ec2RunnerOperations): Omit { +export function createEc2ScaleDownProvider( + runnerOperations: Ec2RunnerOperations, +): Omit { return { - list: (environment, orphan) => listEc2ScaleDownRunners(runners, environment, orphan), + list: (environment, orphan) => listEc2ScaleDownRunners(runnerOperations, environment, orphan), bootTimeExceeded, - markOrphan: (id) => markEc2RunnerOrphan(runners, id), - unmarkOrphan: (id) => unmarkEc2RunnerOrphan(runners, id), - terminate: (id) => runners.terminate(id), + markOrphan: (id) => markEc2RunnerOrphan(runnerOperations, id), + unmarkOrphan: (id) => unmarkEc2RunnerOrphan(runnerOperations, id), + terminate: (id) => runnerOperations.terminate(id), }; } diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.ts index 8f2c2457e0..4120a4f6e7 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.ts @@ -60,22 +60,22 @@ async function resolveEc2LabelsForRunners( } async function getCurrentEc2Runners( - runners: Ec2RunnerOperations, + runnerOperations: Ec2RunnerOperations, _state: Ec2ScaleUpState, { runnerType, runnerOwner }: CurrentRunnersInput, ): Promise { - return (await runners.list({ environment: process.env.ENVIRONMENT, runnerType, runnerOwner })).length; + return (await runnerOperations.list({ environment: process.env.ENVIRONMENT, runnerType, runnerOwner })).length; } async function createEc2ScaleUpRunners( { githubRunnerConfig, numberOfRunners, githubInstallationClient, state }: CreateScaleUpRunnersInput, createStartRunnerConfig: CreateStartRunnerConfig, - runners: Ec2RunnerOperations, + runnerOperations: Ec2RunnerOperations, ): Promise { const config = loadEc2ScaleUpProviderConfig(); return await createRunners( - runners, + runnerOperations, githubRunnerConfig, { ...config, @@ -90,12 +90,12 @@ async function createEc2ScaleUpRunners( export function createEc2ScaleUpProvider( createStartRunnerConfig: CreateStartRunnerConfig, - runners: Ec2RunnerOperations, + runnerOperations: Ec2RunnerOperations, ec2Client: EC2Client, ): Omit, 'type'> { return { resolveLabelsForRunners: (labels) => resolveEc2LabelsForRunners(ec2Client, labels), - getCurrentRunners: (state, input) => getCurrentEc2Runners(runners, state, input), - createRunners: (input) => createEc2ScaleUpRunners(input, createStartRunnerConfig, runners), + getCurrentRunners: (state, input) => getCurrentEc2Runners(runnerOperations, state, input), + createRunners: (input) => createEc2ScaleUpRunners(input, createStartRunnerConfig, runnerOperations), }; } diff --git a/lambdas/libs/compute-providers/aws/ec2/src/runners.test.ts b/lambdas/libs/compute-providers/aws/ec2/src/runners.test.ts index 1520d86a22..d5a0d13d0a 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/runners.test.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/runners.test.ts @@ -506,12 +506,12 @@ describe('create runner', () => { it('keeps cancellation request-scoped and rejects before calling AWS', async () => { const abortController = new AbortController(); const abortReason = new Error('service stopping'); - const runners = createEc2RunnerClient(new EC2Client({})).forRequest({ + const runnerOperations = createEc2RunnerClient(new EC2Client({})).forRequest({ signal: abortController.signal, }); abortController.abort(abortReason); - await expect(runners.create(createRunnerConfig(defaultRunnerConfig))).rejects.toThrow('service stopping'); + await expect(runnerOperations.create(createRunnerConfig(defaultRunnerConfig))).rejects.toThrow('service stopping'); expect(mockEC2Client).not.toHaveReceivedCommand(CreateFleetCommand); expect(mockSSMClient).not.toHaveReceivedCommand(GetParameterCommand); }); From f82b17ab71939d439fd833df9d47dcd2abfc759c Mon Sep 17 00:00:00 2001 From: edersonbrilhante Date: Thu, 27 Aug 2026 14:27:33 +0200 Subject: [PATCH 3/5] refactor(ec2): put EC2 client first --- lambdas/libs/compute-providers/aws/ec2/src/runners.ts | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/lambdas/libs/compute-providers/aws/ec2/src/runners.ts b/lambdas/libs/compute-providers/aws/ec2/src/runners.ts index 69f6eec76c..8f4aa5e2a7 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/runners.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/runners.ts @@ -61,7 +61,7 @@ export function createEc2RunnerClient(ec2Client: EC2Client): Ec2RunnerClient { forRequest: ({ signal }) => ({ list: (filters) => runWithRequestSignal(signal, () => listRunners(ec2Client, filters, signal)), create: (runnerParameters) => - runWithRequestSignal(signal, () => createEc2Runner(runnerParameters, ec2Client, signal)), + runWithRequestSignal(signal, () => createEc2Runner(ec2Client, runnerParameters, signal)), terminate: (instanceId) => runWithRequestSignal(signal, () => terminateEc2Runner(ec2Client, instanceId, signal)), tag: (instanceId, tags) => runWithRequestSignal(signal, () => tagEc2Runner(ec2Client, instanceId, tags, signal)), untag: (instanceId, tags) => @@ -385,8 +385,8 @@ export async function createRunner(runnerParameters: RunnerInputParameters): Pro } async function createEc2Runner( - runnerParameters: RunnerInputParameters, ec2Client: EC2Client, + runnerParameters: RunnerInputParameters, signal: AbortSignal | undefined, ): Promise { logger.debug('Runner configuration.', { @@ -470,6 +470,7 @@ async function processFleetResult( 'on-demand', ); const onDemandResult = await createEc2Runner( + ec2Client, { ...runnerParameters, numberOfRunners: numberOfInstances, @@ -480,7 +481,6 @@ async function processFleetResult( instanceAllocationStrategy: failoverAllocationStrategy, }, }, - ec2Client, signal, ); instances.push(...onDemandResult.instances); From 1ddd0351fc3b322c50059c974db38005c2a52881 Mon Sep 17 00:00:00 2001 From: edersonbrilhante Date: Thu, 27 Aug 2026 14:57:42 +0200 Subject: [PATCH 4/5] refactor(ec2): simplify runner resource boundary --- .../aws/ec2/control-plane.ts | 21 ++-- .../aws/ec2/src/control-plane/pool.test.ts | 18 +-- .../aws/ec2/src/control-plane/pool.ts | 83 +++++++------- .../ec2/src/control-plane/runner-config.ts | 13 ++- .../ec2/src/control-plane/scale-down.test.ts | 4 +- .../aws/ec2/src/control-plane/scale-down.ts | 28 +---- .../ec2/src/control-plane/scale-up.test.ts | 6 +- .../aws/ec2/src/control-plane/scale-up.ts | 107 ++++++++---------- .../compute-providers/aws/ec2/src/runners.ts | 8 +- 9 files changed, 128 insertions(+), 160 deletions(-) diff --git a/lambdas/libs/compute-providers/aws/ec2/control-plane.ts b/lambdas/libs/compute-providers/aws/ec2/control-plane.ts index be24e0b241..b26babcd37 100644 --- a/lambdas/libs/compute-providers/aws/ec2/control-plane.ts +++ b/lambdas/libs/compute-providers/aws/ec2/control-plane.ts @@ -7,31 +7,24 @@ import type {} from './src/environment'; import { createEc2PoolProvider } from './src/control-plane/pool'; import { createEc2ScaleDownProvider } from './src/control-plane/scale-down'; import { createEc2ScaleUpProvider } from './src/control-plane/scale-up'; -import { createEc2RunnerClient, type Ec2RunnerOperations } from './src/runners'; +import { createEc2RunnerClient } from './src/runners'; -export function createEc2ControlPlanePluginWithDependencies( +export function createEc2ControlPlanePlugin( createStartRunnerConfig: CreateStartRunnerConfig, - runnerOperations: Ec2RunnerOperations, - ec2Client: EC2Client, + ec2Client: EC2Client = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION })), ): ComputeProviderPlugin { + const runnerOperations = createEc2RunnerClient(ec2Client).forRequest({ signal: undefined }); + return { type: 'ec2', capabilities: { - pool: () => createEc2PoolProvider(createStartRunnerConfig, runnerOperations), - scaleUp: () => createEc2ScaleUpProvider(createStartRunnerConfig, runnerOperations, ec2Client), + pool: () => createEc2PoolProvider(runnerOperations, createStartRunnerConfig), + scaleUp: () => createEc2ScaleUpProvider(runnerOperations, ec2Client, createStartRunnerConfig), scaleDown: () => createEc2ScaleDownProvider(runnerOperations), }, }; } -function createEc2ControlPlanePlugin( - createStartRunnerConfig: CreateStartRunnerConfig, -): ComputeProviderPlugin { - const ec2Client = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION })); - const runnerOperations = createEc2RunnerClient(ec2Client).forRequest({ signal: undefined }); - return createEc2ControlPlanePluginWithDependencies(createStartRunnerConfig, runnerOperations, ec2Client); -} - export const provider = { type: 'ec2', createPlugin: createEc2ControlPlanePlugin, diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.test.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.test.ts index 82c1728d04..f23df5ec6e 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.test.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.test.ts @@ -1,6 +1,6 @@ import type { Octokit } from '@octokit/rest'; import type { CreateGitHubRunnerConfig, CreateStartRunnerConfig, RunnerInfo } from '../../../../core'; -import { bootTimeExceeded, type Ec2RunnerOperations } from '../runners'; +import { bootTimeExceeded, type Ec2RunnerResourceOperations } from '../runners'; import { calculateEc2PoolSize, createEc2PoolProvider } from './pool'; import { createRunners, type Ec2ProviderConfig, loadEc2ProviderConfig } from './runner-config'; import { beforeEach, describe, expect, it, vi } from 'vitest'; @@ -19,12 +19,12 @@ const mockCreateRunners = vi.mocked(createRunners); const mockLoadProviderConfig = vi.mocked(loadEc2ProviderConfig); const runnerOperations = { - list: vi.fn(), - create: vi.fn(), - terminate: vi.fn(), - tag: vi.fn(), - untag: vi.fn(), -} satisfies Ec2RunnerOperations; + list: vi.fn(), + create: vi.fn(), + terminate: vi.fn(), + tag: vi.fn(), + untag: vi.fn(), +} satisfies Ec2RunnerResourceOperations; describe('calculateEc2PoolSize', () => { beforeEach(() => { @@ -114,7 +114,7 @@ describe('createEc2PoolProvider', () => { it('lists only running instances managed for the requested pool', async () => { const runners: RunnerInfo[] = [{ id: 'i-running', owner: 'owner', type: 'Org' }]; runnerOperations.list.mockResolvedValue(runners); - const provider = createEc2PoolProvider(createStartRunnerConfig, runnerOperations); + const provider = createEc2PoolProvider(runnerOperations, createStartRunnerConfig); await expect( provider.listRunners({ @@ -137,7 +137,7 @@ describe('createEc2PoolProvider', () => { retryableErrorCount: 0, nonRetryableErrorCount: 0, }); - const provider = createEc2PoolProvider(createStartRunnerConfig, runnerOperations); + const provider = createEc2PoolProvider(runnerOperations, createStartRunnerConfig); await expect( provider.createRunners({ diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.ts index b768efb3ec..ed02d05159 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.ts @@ -7,59 +7,54 @@ import type { RunnerInfo, RunnerStatus, } from '../../../../core'; -import { bootTimeExceeded, type Ec2RunnerOperations } from '../runners'; +import { bootTimeExceeded, type Ec2RunnerResourceOperations } from '../runners'; import { createRunners, loadEc2ProviderConfig } from './runner-config'; const logger = createChildLogger('pool'); -async function listEc2PoolRunners( - runnerOperations: Ec2RunnerOperations, - { environment, runnerOwner, runnerType }: ListPoolRunnersInput, -): Promise { - return await runnerOperations.list({ - environment, - runnerOwner, - runnerType, - statuses: ['running'], - }); -} - -async function createEc2PoolRunners( - { githubRunnerConfig, numberOfRunners, githubInstallationClient }: CreatePoolRunnersInput, - createStartRunnerConfig: CreateStartRunnerConfig, - runnerOperations: Ec2RunnerOperations, -): Promise { - const config = loadEc2ProviderConfig(); - - const { instances } = await createRunners( - runnerOperations, - githubRunnerConfig, - { - ec2instanceCriteria: config.ec2instanceCriteria, - environment: config.environment, - launchTemplateName: config.launchTemplateName, - subnets: config.subnets, - amiIdSsmParameterName: config.amiIdSsmParameterName, - tracingEnabled: config.tracingEnabled, - onDemandFailoverOnError: config.onDemandFailoverOnError, - scaleErrors: config.scaleErrors, - }, - numberOfRunners, - githubInstallationClient, - createStartRunnerConfig, - 'pool-lambda', - ); - return instances; -} - export function createEc2PoolProvider( + runnerOperations: Ec2RunnerResourceOperations, createStartRunnerConfig: CreateStartRunnerConfig, - runnerOperations: Ec2RunnerOperations, ): Omit, 'type'> { return { - listRunners: (input) => listEc2PoolRunners(runnerOperations, input), + async listRunners({ environment, runnerOwner, runnerType }: ListPoolRunnersInput): Promise { + return await runnerOperations.list({ + environment, + runnerOwner, + runnerType, + statuses: ['running'], + }); + }, + countAvailableRunners: calculateEc2PoolSize, - createRunners: (input) => createEc2PoolRunners(input, createStartRunnerConfig, runnerOperations), + + async createRunners({ + githubRunnerConfig, + numberOfRunners, + githubInstallationClient, + }: CreatePoolRunnersInput): Promise { + const config = loadEc2ProviderConfig(); + + const { instances } = await createRunners( + runnerOperations, + githubRunnerConfig, + { + ec2instanceCriteria: config.ec2instanceCriteria, + environment: config.environment, + launchTemplateName: config.launchTemplateName, + subnets: config.subnets, + amiIdSsmParameterName: config.amiIdSsmParameterName, + tracingEnabled: config.tracingEnabled, + onDemandFailoverOnError: config.onDemandFailoverOnError, + scaleErrors: config.scaleErrors, + }, + numberOfRunners, + githubInstallationClient, + createStartRunnerConfig, + 'pool-lambda', + ); + return instances; + }, }; } diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-config.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-config.ts index 083b7ab56d..cbce200f42 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-config.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-config.ts @@ -11,7 +11,7 @@ import { Octokit } from '@octokit/rest'; import type { Tag } from '@aws-sdk/client-ec2'; import yn from 'yn'; -import type { Ec2RunnerOperations } from '../runners'; +import type { Ec2RunnerResourceOperations } from '../runners'; import type { RunnerInputParameters } from '../runners.d'; const logger = createChildLogger('ec2-runners'); @@ -61,7 +61,7 @@ export function loadEc2ProviderConfig(): Ec2ProviderConfig { } export async function createRunners( - runnerOperations: Ec2RunnerOperations, + runnerOperations: Ec2RunnerResourceOperations, githubRunnerConfig: CreateGitHubRunnerConfig, ec2RunnerConfig: CreateEC2RunnerConfig, numberOfRunners: number, @@ -127,7 +127,10 @@ export async function createRunners( return result; } -async function terminateFailedInstances(runnerOperations: Ec2RunnerOperations, instanceIds: string[]): Promise { +async function terminateFailedInstances( + runnerOperations: Ec2RunnerResourceOperations, + instanceIds: string[], +): Promise { for (const instanceId of instanceIds) { try { await runnerOperations.terminate(instanceId); @@ -140,7 +143,7 @@ async function terminateFailedInstances(runnerOperations: Ec2RunnerOperations, i } } -function createEc2StartRunnerConfigOptions(runnerOperations: Ec2RunnerOperations): StartRunnerConfigOptions { +function createEc2StartRunnerConfigOptions(runnerOperations: Ec2RunnerResourceOperations): StartRunnerConfigOptions { return { getSsmParameterTags: (instanceId) => [{ Key: 'InstanceId', Value: instanceId }], onJitConfigCreated: async (instanceId, metadata) => @@ -149,7 +152,7 @@ function createEc2StartRunnerConfigOptions(runnerOperations: Ec2RunnerOperations } async function tagEc2RunnerMetadata( - runnerOperations: Ec2RunnerOperations, + runnerOperations: Ec2RunnerResourceOperations, instanceId: string, metadata: GitHubRunnerMetadata, ): Promise { diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.test.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.test.ts index 96f2115c3c..fe07cfa056 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.test.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.test.ts @@ -3,7 +3,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; import type { RunnerInfo, RunnerType } from '../../../../core'; import { createEc2ScaleDownProvider } from './scale-down'; import { listEC2Runners, tag, terminateRunner, untag } from '../runners'; -import type { Ec2RunnerOperations } from '../runners'; +import type { Ec2RunnerResourceOperations } from '../runners'; vi.mock('../runners', async (importOriginal) => { const actual = await importOriginal(); @@ -20,7 +20,7 @@ const mockListRunners = vi.mocked(listEC2Runners); const mockTagRunner = vi.mocked(tag); const mockTerminateRunner = vi.mocked(terminateRunner); const mockUntagRunner = vi.mocked(untag); -const runnerOperations: Ec2RunnerOperations = { +const runnerOperations: Ec2RunnerResourceOperations = { list: mockListRunners, create: vi.fn(), terminate: mockTerminateRunner, diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.ts index 23381c69cd..310e6e3e30 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.ts @@ -1,30 +1,14 @@ -import type { RunnerInfo, ScaleDownComputeProvider } from '../../../../core'; -import { bootTimeExceeded, type Ec2RunnerOperations } from '../runners'; - -async function listEc2ScaleDownRunners( - runnerOperations: Ec2RunnerOperations, - environment: string, - orphan?: boolean, -): Promise { - return await runnerOperations.list({ environment, orphan }); -} - -async function markEc2RunnerOrphan(runnerOperations: Ec2RunnerOperations, id: string): Promise { - await runnerOperations.tag(id, [{ Key: 'ghr:orphan', Value: 'true' }]); -} - -async function unmarkEc2RunnerOrphan(runnerOperations: Ec2RunnerOperations, id: string): Promise { - await runnerOperations.untag(id, [{ Key: 'ghr:orphan', Value: 'true' }]); -} +import type { ScaleDownComputeProvider } from '../../../../core'; +import { bootTimeExceeded, type Ec2RunnerResourceOperations } from '../runners'; export function createEc2ScaleDownProvider( - runnerOperations: Ec2RunnerOperations, + runnerOperations: Ec2RunnerResourceOperations, ): Omit { return { - list: (environment, orphan) => listEc2ScaleDownRunners(runnerOperations, environment, orphan), + list: (environment, orphan) => runnerOperations.list({ environment, orphan }), bootTimeExceeded, - markOrphan: (id) => markEc2RunnerOrphan(runnerOperations, id), - unmarkOrphan: (id) => unmarkEc2RunnerOrphan(runnerOperations, id), + markOrphan: (id) => runnerOperations.tag(id, [{ Key: 'ghr:orphan', Value: 'true' }]), + unmarkOrphan: (id) => runnerOperations.untag(id, [{ Key: 'ghr:orphan', Value: 'true' }]), terminate: (id) => runnerOperations.terminate(id), }; } diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts index a41f143241..a11248a31b 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts @@ -8,7 +8,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; import { parseEc2OverrideConfig } from './dynamic-labels'; import { EC2_TAG_VALUE_MAX_LENGTH, RUNNER_LABELS_TAG_MAX_COUNT } from './runner-config'; import { createRunner, listEC2Runners, tag, terminateRunner } from '../runners'; -import type { Ec2RunnerOperations } from '../runners'; +import type { Ec2RunnerResourceOperations } from '../runners'; import type { RunnerInputParameters } from '../runners.d'; import { createEc2ScaleUpProvider } from './scale-up'; @@ -30,14 +30,14 @@ const githubClient = {} as Octokit; const runnerOwner = 'Codertocat'; const repositoryRunnerOwner = 'Codertocat/hello-world'; const cleanEnv = process.env; -const runnerOperations: Ec2RunnerOperations = { +const runnerOperations: Ec2RunnerResourceOperations = { list: mockListRunners, create: mockCreateRunner, terminate: mockTerminateRunner, tag: mockTag, untag: vi.fn(), }; -const provider = createEc2ScaleUpProvider(mockCreateStartRunnerConfig, runnerOperations, new EC2Client({})); +const provider = createEc2ScaleUpProvider(runnerOperations, new EC2Client({}), mockCreateStartRunnerConfig); interface CreateProviderRunnersOptions { labels?: string[]; diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.ts index 4120a4f6e7..d1bc416ca0 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.ts @@ -10,7 +10,7 @@ import type { } from '../../../../core'; import yn from 'yn'; -import type { Ec2RunnerOperations } from '../runners'; +import type { Ec2RunnerResourceOperations } from '../runners'; import type { Ec2OverrideConfig } from '../runners.d'; import { getDefaultBlockDeviceNameFromLaunchTemplate, @@ -33,69 +33,62 @@ function loadEc2ScaleUpProviderConfig(): CreateEC2RunnerConfig { }; } -async function resolveEc2LabelsForRunners( +export function createEc2ScaleUpProvider( + runnerOperations: Ec2RunnerResourceOperations, ec2Client: EC2Client, - messageLabels: string[], -): Promise> { - const trimmedLabels = messageLabels.map((label) => label.trim()); - const dynamicEC2Labels = trimmedLabels.filter((label) => label.startsWith('ghr-ec2-')); - const nonEc2DynamicLabels = trimmedLabels.filter( - (label) => label.startsWith('ghr-') && !label.startsWith('ghr-ec2-'), - ); - const runnerLabels = [...nonEc2DynamicLabels, ...dynamicEC2Labels]; - let ec2OverrideConfig: Ec2OverrideConfig | undefined; + createStartRunnerConfig: CreateStartRunnerConfig, +): Omit, 'type'> { + return { + async resolveLabelsForRunners(messageLabels: string[]): Promise> { + const trimmedLabels = messageLabels.map((label) => label.trim()); + const dynamicEC2Labels = trimmedLabels.filter((label) => label.startsWith('ghr-ec2-')); + const nonEc2DynamicLabels = trimmedLabels.filter( + (label) => label.startsWith('ghr-') && !label.startsWith('ghr-ec2-'), + ); + const runnerLabels = [...nonEc2DynamicLabels, ...dynamicEC2Labels]; + let ec2OverrideConfig: Ec2OverrideConfig | undefined; - if (dynamicEC2Labels.length > 0) { - const defaultBlockDeviceName = shouldLoadLaunchTemplateBlockDeviceName(dynamicEC2Labels) - ? await getDefaultBlockDeviceNameFromLaunchTemplate(ec2Client, process.env.LAUNCH_TEMPLATE_NAME) - : undefined; + if (dynamicEC2Labels.length > 0) { + const defaultBlockDeviceName = shouldLoadLaunchTemplateBlockDeviceName(dynamicEC2Labels) + ? await getDefaultBlockDeviceNameFromLaunchTemplate(ec2Client, process.env.LAUNCH_TEMPLATE_NAME) + : undefined; - ec2OverrideConfig = parseEc2OverrideConfig(dynamicEC2Labels, defaultBlockDeviceName); - if (ec2OverrideConfig) { - logger.debug('EC2 override config parsed from labels', { ec2OverrideConfig }); - } - } + ec2OverrideConfig = parseEc2OverrideConfig(dynamicEC2Labels, defaultBlockDeviceName); + if (ec2OverrideConfig) { + logger.debug('EC2 override config parsed from labels', { ec2OverrideConfig }); + } + } - return { runnerLabels, state: { ec2OverrideConfig } }; -} + return { runnerLabels, state: { ec2OverrideConfig } }; + }, -async function getCurrentEc2Runners( - runnerOperations: Ec2RunnerOperations, - _state: Ec2ScaleUpState, - { runnerType, runnerOwner }: CurrentRunnersInput, -): Promise { - return (await runnerOperations.list({ environment: process.env.ENVIRONMENT, runnerType, runnerOwner })).length; -} + async getCurrentRunners( + _state: Ec2ScaleUpState, + { runnerType, runnerOwner }: CurrentRunnersInput, + ): Promise { + return (await runnerOperations.list({ environment: process.env.ENVIRONMENT, runnerType, runnerOwner })).length; + }, -async function createEc2ScaleUpRunners( - { githubRunnerConfig, numberOfRunners, githubInstallationClient, state }: CreateScaleUpRunnersInput, - createStartRunnerConfig: CreateStartRunnerConfig, - runnerOperations: Ec2RunnerOperations, -): Promise { - const config = loadEc2ScaleUpProviderConfig(); + async createRunners({ + githubRunnerConfig, + numberOfRunners, + githubInstallationClient, + state, + }: CreateScaleUpRunnersInput): Promise { + const config = loadEc2ScaleUpProviderConfig(); - return await createRunners( - runnerOperations, - githubRunnerConfig, - { - ...config, - ec2OverrideConfig: state.ec2OverrideConfig, + return await createRunners( + runnerOperations, + githubRunnerConfig, + { + ...config, + ec2OverrideConfig: state.ec2OverrideConfig, + }, + numberOfRunners, + githubInstallationClient, + createStartRunnerConfig, + 'scale-up-lambda', + ); }, - numberOfRunners, - githubInstallationClient, - createStartRunnerConfig, - 'scale-up-lambda', - ); -} - -export function createEc2ScaleUpProvider( - createStartRunnerConfig: CreateStartRunnerConfig, - runnerOperations: Ec2RunnerOperations, - ec2Client: EC2Client, -): Omit, 'type'> { - return { - resolveLabelsForRunners: (labels) => resolveEc2LabelsForRunners(ec2Client, labels), - getCurrentRunners: (state, input) => getCurrentEc2Runners(runnerOperations, state, input), - createRunners: (input) => createEc2ScaleUpRunners(input, createStartRunnerConfig, runnerOperations), }; } diff --git a/lambdas/libs/compute-providers/aws/ec2/src/runners.ts b/lambdas/libs/compute-providers/aws/ec2/src/runners.ts index 8f4aa5e2a7..f0ef27c9fa 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/runners.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/runners.ts @@ -36,7 +36,7 @@ export interface Ec2RunnerRequestContext { readonly signal: AbortSignal | undefined; } -export interface Ec2RunnerOperations { +export interface Ec2RunnerResourceOperations { list(filters?: Ec2ListRunnerFilters): Promise; create(runnerParameters: RunnerInputParameters): Promise; terminate(instanceId: string): Promise; @@ -45,7 +45,7 @@ export interface Ec2RunnerOperations { } export interface Ec2RunnerClient { - forRequest(context: Ec2RunnerRequestContext): Ec2RunnerOperations; + forRequest(context: Ec2RunnerRequestContext): Ec2RunnerResourceOperations; } async function runWithRequestSignal( @@ -70,9 +70,9 @@ export function createEc2RunnerClient(ec2Client: EC2Client): Ec2RunnerClient { }; } -let defaultRunnerOperations: Ec2RunnerOperations | undefined; +let defaultRunnerOperations: Ec2RunnerResourceOperations | undefined; -function getDefaultRunnerOperations(): Ec2RunnerOperations { +function getDefaultRunnerOperations(): Ec2RunnerResourceOperations { defaultRunnerOperations ??= createEc2RunnerClient( getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION })), ).forRequest({ signal: undefined }); From 6857f950988a5596ac2ec18284e17c81cb6e5097 Mon Sep 17 00:00:00 2001 From: edersonbrilhante Date: Thu, 27 Aug 2026 15:35:29 +0200 Subject: [PATCH 5/5] refactor(compute-providers): clarify control-plane capabilities --- .../aws/ec2/control-plane.ts | 16 +- .../ec2/src/control-plane/dynamic-labels.ts | 25 --- .../aws/ec2/src/control-plane/pool.test.ts | 75 ++++---- .../aws/ec2/src/control-plane/pool.ts | 86 ++++----- .../{runner-config.ts => runner-creation.ts} | 21 +-- .../ec2/src/control-plane/scale-down.test.ts | 43 ++--- .../aws/ec2/src/control-plane/scale-down.ts | 12 +- .../ec2/src/control-plane/scale-up.test.ts | 80 +++----- .../aws/ec2/src/control-plane/scale-up.ts | 93 ++++----- .../aws/ec2/src/launch-template.ts | 26 +++ .../aws/ec2/src/runners.test.ts | 176 +++++++++++------- .../compute-providers/aws/ec2/src/runners.ts | 49 ++--- lambdas/libs/compute-providers/package.json | 2 +- .../templates/provider/control-plane.ts | 12 +- 14 files changed, 331 insertions(+), 385 deletions(-) rename lambdas/libs/compute-providers/aws/ec2/src/control-plane/{runner-config.ts => runner-creation.ts} (90%) create mode 100644 lambdas/libs/compute-providers/aws/ec2/src/launch-template.ts diff --git a/lambdas/libs/compute-providers/aws/ec2/control-plane.ts b/lambdas/libs/compute-providers/aws/ec2/control-plane.ts index b26babcd37..07b5ce1f1d 100644 --- a/lambdas/libs/compute-providers/aws/ec2/control-plane.ts +++ b/lambdas/libs/compute-providers/aws/ec2/control-plane.ts @@ -4,23 +4,23 @@ import { EC2Client } from '@aws-sdk/client-ec2'; import type { ControlPlaneProviderCapabilities, ControlPlaneProviderModule } from '../../contracts'; import type {} from './src/environment'; -import { createEc2PoolProvider } from './src/control-plane/pool'; -import { createEc2ScaleDownProvider } from './src/control-plane/scale-down'; -import { createEc2ScaleUpProvider } from './src/control-plane/scale-up'; +import { createEc2PoolCapability } from './src/control-plane/pool'; +import { createEc2ScaleDownCapability } from './src/control-plane/scale-down'; +import { createEc2ScaleUpCapability } from './src/control-plane/scale-up'; import { createEc2RunnerClient } from './src/runners'; export function createEc2ControlPlanePlugin( createStartRunnerConfig: CreateStartRunnerConfig, - ec2Client: EC2Client = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION })), ): ComputeProviderPlugin { - const runnerOperations = createEc2RunnerClient(ec2Client).forRequest({ signal: undefined }); + const ec2Client = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION })); + const ec2Operations = createEc2RunnerClient(ec2Client).forRequest({ signal: undefined }); return { type: 'ec2', capabilities: { - pool: () => createEc2PoolProvider(runnerOperations, createStartRunnerConfig), - scaleUp: () => createEc2ScaleUpProvider(runnerOperations, ec2Client, createStartRunnerConfig), - scaleDown: () => createEc2ScaleDownProvider(runnerOperations), + pool: () => createEc2PoolCapability(ec2Operations, createStartRunnerConfig), + scaleUp: () => createEc2ScaleUpCapability(ec2Operations, createStartRunnerConfig), + scaleDown: () => createEc2ScaleDownCapability(ec2Operations), }, }; } diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/dynamic-labels.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/dynamic-labels.ts index a28c12ca97..ec6ac2489d 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/dynamic-labels.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/dynamic-labels.ts @@ -11,8 +11,6 @@ import { BurstablePerformance, CpuManufacturer, CpuPerformanceFactorRequest, - DescribeLaunchTemplateVersionsCommand, - type EC2Client, FleetBlockDeviceMappingRequest, FleetEbsBlockDeviceRequest, InstanceGeneration, @@ -364,26 +362,3 @@ export function shouldLoadLaunchTemplateBlockDeviceName(labels: string[]): boole return hasBlockDeviceOverride && !hasBlockDeviceName; } - -export async function getDefaultBlockDeviceNameFromLaunchTemplate( - ec2Client: EC2Client, - launchTemplateName: string, -): Promise { - const launchTemplateVersions = await ec2Client.send( - new DescribeLaunchTemplateVersionsCommand({ - LaunchTemplateName: launchTemplateName, - Versions: ['$Default'], - }), - ); - const blockDeviceMappings = - launchTemplateVersions.LaunchTemplateVersions?.[0]?.LaunchTemplateData?.BlockDeviceMappings; - const blockDeviceName = - blockDeviceMappings?.find((blockDeviceMapping) => blockDeviceMapping.DeviceName && blockDeviceMapping.Ebs) - ?.DeviceName ?? blockDeviceMappings?.find((blockDeviceMapping) => blockDeviceMapping.DeviceName)?.DeviceName; - - if (!blockDeviceName) { - throw new Error(`Failed to determine block device name from launch template '${launchTemplateName}'.`); - } - - return blockDeviceName; -} diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.test.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.test.ts index f23df5ec6e..85b65ba7cf 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.test.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.test.ts @@ -1,15 +1,15 @@ import type { Octokit } from '@octokit/rest'; import type { CreateGitHubRunnerConfig, CreateStartRunnerConfig, RunnerInfo } from '../../../../core'; import { bootTimeExceeded, type Ec2RunnerResourceOperations } from '../runners'; -import { calculateEc2PoolSize, createEc2PoolProvider } from './pool'; -import { createRunners, type Ec2ProviderConfig, loadEc2ProviderConfig } from './runner-config'; +import { createEc2PoolCapability } from './pool'; +import { createRunners, type Ec2ProviderConfig, loadEc2ProviderConfig } from './runner-creation'; import { beforeEach, describe, expect, it, vi } from 'vitest'; vi.mock('../runners', () => ({ bootTimeExceeded: vi.fn(), })); -vi.mock('./runner-config', () => ({ +vi.mock('./runner-creation', () => ({ createRunners: vi.fn(), loadEc2ProviderConfig: vi.fn(), })); @@ -18,15 +18,17 @@ const mockBootTimeExceeded = vi.mocked(bootTimeExceeded); const mockCreateRunners = vi.mocked(createRunners); const mockLoadProviderConfig = vi.mocked(loadEc2ProviderConfig); -const runnerOperations = { +const ec2Operations = { list: vi.fn(), create: vi.fn(), terminate: vi.fn(), tag: vi.fn(), untag: vi.fn(), } satisfies Ec2RunnerResourceOperations; +const createStartRunnerConfig = vi.fn(); +const capability = createEc2PoolCapability(ec2Operations, createStartRunnerConfig); -describe('calculateEc2PoolSize', () => { +describe('createEc2PoolCapability.countAvailableRunners', () => { beforeEach(() => { vi.clearAllMocks(); }); @@ -35,7 +37,7 @@ describe('calculateEc2PoolSize', () => { const runners: RunnerInfo[] = [{ id: 'i-idle', owner: 'owner', type: 'Org' }]; const runnerStatus = new Map([['i-idle', { busy: false, status: 'online' }]]); - expect(calculateEc2PoolSize(runners, runnerStatus)).toBe(1); + expect(capability.countAvailableRunners(runners, runnerStatus)).toBe(1); expect(mockBootTimeExceeded).not.toHaveBeenCalled(); }); @@ -49,7 +51,7 @@ describe('calculateEc2PoolSize', () => { ['i-offline', { busy: false, status: 'offline' }], ]); - expect(calculateEc2PoolSize(runners, runnerStatus)).toBe(0); + expect(capability.countAvailableRunners(runners, runnerStatus)).toBe(0); expect(mockBootTimeExceeded).not.toHaveBeenCalled(); }); @@ -57,7 +59,7 @@ describe('calculateEc2PoolSize', () => { const runners: RunnerInfo[] = [{ id: 'i-busy', owner: 'owner', type: 'Org' }]; const runnerStatus = new Map([['i-busy', { busy: true, status: 'online' }]]); - expect(calculateEc2PoolSize(runners, runnerStatus, true)).toBe(1); + expect(capability.countAvailableRunners(runners, runnerStatus, true)).toBe(1); expect(mockBootTimeExceeded).not.toHaveBeenCalled(); }); @@ -65,19 +67,43 @@ describe('calculateEc2PoolSize', () => { const runners: RunnerInfo[] = [{ id: 'i-booting', owner: 'owner', type: 'Org' }]; mockBootTimeExceeded.mockReturnValue(false); - expect(calculateEc2PoolSize(runners, new Map())).toBe(1); + expect(capability.countAvailableRunners(runners, new Map())).toBe(1); }); it('does not count unregistered runners whose boot time expired', () => { const runners: RunnerInfo[] = [{ id: 'i-expired', owner: 'owner', type: 'Org' }]; mockBootTimeExceeded.mockReturnValue(true); - expect(calculateEc2PoolSize(runners, new Map())).toBe(0); + expect(capability.countAvailableRunners(runners, new Map())).toBe(0); }); }); -describe('createEc2PoolProvider', () => { - const createStartRunnerConfig = vi.fn(); +describe('createEc2PoolCapability.listRunners', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('lists only running instances managed for the requested pool', async () => { + const runners: RunnerInfo[] = [{ id: 'i-running', owner: 'owner', type: 'Org' }]; + ec2Operations.list.mockResolvedValue(runners); + + await expect( + capability.listRunners({ + environment: 'test-environment', + runnerOwner: 'owner', + runnerType: 'Org', + }), + ).resolves.toBe(runners); + expect(ec2Operations.list).toHaveBeenCalledWith({ + environment: 'test-environment', + runnerOwner: 'owner', + runnerType: 'Org', + statuses: ['running'], + }); + }); +}); + +describe('createEc2PoolCapability.createRunners', () => { const githubInstallationClient = {} as Octokit; const githubRunnerConfig: CreateGitHubRunnerConfig = { ephemeral: true, @@ -111,43 +137,22 @@ describe('createEc2PoolProvider', () => { mockLoadProviderConfig.mockReturnValue(providerConfig); }); - it('lists only running instances managed for the requested pool', async () => { - const runners: RunnerInfo[] = [{ id: 'i-running', owner: 'owner', type: 'Org' }]; - runnerOperations.list.mockResolvedValue(runners); - const provider = createEc2PoolProvider(runnerOperations, createStartRunnerConfig); - - await expect( - provider.listRunners({ - environment: 'test-environment', - runnerOwner: 'owner', - runnerType: 'Org', - }), - ).resolves.toBe(runners); - expect(runnerOperations.list).toHaveBeenCalledWith({ - environment: 'test-environment', - runnerOwner: 'owner', - runnerType: 'Org', - statuses: ['running'], - }); - }); - it('creates pool runners with the pool source and returns their instance IDs', async () => { mockCreateRunners.mockResolvedValue({ instances: ['i-created'], retryableErrorCount: 0, nonRetryableErrorCount: 0, }); - const provider = createEc2PoolProvider(runnerOperations, createStartRunnerConfig); await expect( - provider.createRunners({ + capability.createRunners({ githubRunnerConfig, numberOfRunners: 1, githubInstallationClient, }), ).resolves.toEqual(['i-created']); expect(mockCreateRunners).toHaveBeenCalledWith( - runnerOperations, + ec2Operations, githubRunnerConfig, providerConfig, 1, diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.ts index ed02d05159..f032a78e20 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.ts @@ -1,42 +1,54 @@ import { createChildLogger } from '@aws-github-runner/aws-powertools-util'; -import type { - CreateStartRunnerConfig, - CreatePoolRunnersInput, - ListPoolRunnersInput, - PoolComputeProvider, - RunnerInfo, - RunnerStatus, -} from '../../../../core'; +import type { CreateStartRunnerConfig, PoolComputeProvider, RunnerInfo, RunnerStatus } from '../../../../core'; import { bootTimeExceeded, type Ec2RunnerResourceOperations } from '../runners'; -import { createRunners, loadEc2ProviderConfig } from './runner-config'; +import { createRunners, loadEc2ProviderConfig } from './runner-creation'; const logger = createChildLogger('pool'); -export function createEc2PoolProvider( - runnerOperations: Ec2RunnerResourceOperations, +function countAvailableEc2PoolRunners( + ec2runners: RunnerInfo[], + runnerStatus: Map, + includeBusyRunners = false, +): number { + // 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' + ) { + numberOfRunnersInPool++; + logger.debug(`Runner ${ec2Instance.id} is idle in GitHub and counted as part of the pool`); + } else if (runnerStatus.get(ec2Instance.id) != null) { + logger.debug(`Runner ${ec2Instance.id} is not idle in GitHub and NOT counted as part of the pool`); + } else if (!bootTimeExceeded(ec2Instance)) { + numberOfRunnersInPool++; + logger.info(`Runner ${ec2Instance.id} is still booting and counted as part of the pool`); + } else { + logger.debug(`Runner ${ec2Instance.id} is not idle in GitHub nor booting and not counted as part of the pool`); + } + } + return numberOfRunnersInPool; +} + +export function createEc2PoolCapability( + ec2Operations: Ec2RunnerResourceOperations, createStartRunnerConfig: CreateStartRunnerConfig, ): Omit, 'type'> { return { - async listRunners({ environment, runnerOwner, runnerType }: ListPoolRunnersInput): Promise { - return await runnerOperations.list({ + listRunners: ({ environment, runnerOwner, runnerType }) => + ec2Operations.list({ environment, runnerOwner, runnerType, statuses: ['running'], - }); - }, - - countAvailableRunners: calculateEc2PoolSize, - - async createRunners({ - githubRunnerConfig, - numberOfRunners, - githubInstallationClient, - }: CreatePoolRunnersInput): Promise { + }), + countAvailableRunners: countAvailableEc2PoolRunners, + createRunners: async ({ githubRunnerConfig, numberOfRunners, githubInstallationClient }) => { const config = loadEc2ProviderConfig(); const { instances } = await createRunners( - runnerOperations, + ec2Operations, githubRunnerConfig, { ec2instanceCriteria: config.ec2instanceCriteria, @@ -57,29 +69,3 @@ export function createEc2PoolProvider( }, }; } - -export function calculateEc2PoolSize( - ec2runners: RunnerInfo[], - runnerStatus: Map, - includeBusyRunners = false, -): number { - // 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' - ) { - numberOfRunnersInPool++; - logger.debug(`Runner ${ec2Instance.id} is idle in GitHub and counted as part of the pool`); - } else if (runnerStatus.get(ec2Instance.id) != null) { - logger.debug(`Runner ${ec2Instance.id} is not idle in GitHub and NOT counted as part of the pool`); - } else if (!bootTimeExceeded(ec2Instance)) { - numberOfRunnersInPool++; - logger.info(`Runner ${ec2Instance.id} is still booting and counted as part of the pool`); - } else { - logger.debug(`Runner ${ec2Instance.id} is not idle in GitHub nor booting and not counted as part of the pool`); - } - } - return numberOfRunnersInPool; -} diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-config.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-creation.ts similarity index 90% rename from lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-config.ts rename to lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-creation.ts index cbce200f42..d3863db8ad 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-config.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/runner-creation.ts @@ -61,7 +61,7 @@ export function loadEc2ProviderConfig(): Ec2ProviderConfig { } export async function createRunners( - runnerOperations: Ec2RunnerResourceOperations, + ec2Operations: Ec2RunnerResourceOperations, githubRunnerConfig: CreateGitHubRunnerConfig, ec2RunnerConfig: CreateEC2RunnerConfig, numberOfRunners: number, @@ -71,7 +71,7 @@ export async function createRunners( ): Promise { let result: CreateRunnerResult; try { - result = await runnerOperations.create({ + result = await ec2Operations.create({ runnerType: githubRunnerConfig.runnerType, runnerOwner: githubRunnerConfig.runnerOwner, numberOfRunners, @@ -94,7 +94,7 @@ export async function createRunners( githubRunnerConfig, result.instances, ghClient, - createEc2StartRunnerConfigOptions(runnerOperations), + createEc2StartRunnerConfigOptions(ec2Operations), ); } catch (error) { logger.error('Unexpected error while registering GitHub runners.', { @@ -114,7 +114,7 @@ export async function createRunners( retryable: true, }); - await terminateFailedInstances(runnerOperations, failedInstances); + await terminateFailedInstances(ec2Operations, failedInstances); return { instances: result.instances.filter((id) => !failedInstances.includes(id)), @@ -128,12 +128,12 @@ export async function createRunners( } async function terminateFailedInstances( - runnerOperations: Ec2RunnerResourceOperations, + ec2Operations: Ec2RunnerResourceOperations, instanceIds: string[], ): Promise { for (const instanceId of instanceIds) { try { - await runnerOperations.terminate(instanceId); + await ec2Operations.terminate(instanceId); } catch (error) { logger.error('Failed to terminate instance', { instanceId, @@ -143,16 +143,15 @@ async function terminateFailedInstances( } } -function createEc2StartRunnerConfigOptions(runnerOperations: Ec2RunnerResourceOperations): StartRunnerConfigOptions { +function createEc2StartRunnerConfigOptions(ec2Operations: Ec2RunnerResourceOperations): StartRunnerConfigOptions { return { getSsmParameterTags: (instanceId) => [{ Key: 'InstanceId', Value: instanceId }], - onJitConfigCreated: async (instanceId, metadata) => - await tagEc2RunnerMetadata(runnerOperations, instanceId, metadata), + onJitConfigCreated: async (instanceId, metadata) => await tagEc2RunnerMetadata(ec2Operations, instanceId, metadata), }; } async function tagEc2RunnerMetadata( - runnerOperations: Ec2RunnerResourceOperations, + ec2Operations: Ec2RunnerResourceOperations, instanceId: string, metadata: GitHubRunnerMetadata, ): Promise { @@ -162,7 +161,7 @@ async function tagEc2RunnerMetadata( ]; try { - await runnerOperations.tag(instanceId, tags); + await ec2Operations.tag(instanceId, tags); } catch (e) { logger.error(`Failed to mark EC2 runner '${instanceId}' with GitHub runner metadata.`, { error: e }); } diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.test.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.test.ts index fe07cfa056..2977e22221 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.test.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.test.ts @@ -1,32 +1,22 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; import type { RunnerInfo, RunnerType } from '../../../../core'; -import { createEc2ScaleDownProvider } from './scale-down'; -import { listEC2Runners, tag, terminateRunner, untag } from '../runners'; +import { createEc2ScaleDownCapability } from './scale-down'; import type { Ec2RunnerResourceOperations } from '../runners'; -vi.mock('../runners', async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - listEC2Runners: vi.fn(), - tag: vi.fn(), - terminateRunner: vi.fn(), - untag: vi.fn(), - }; -}); - -const mockListRunners = vi.mocked(listEC2Runners); -const mockTagRunner = vi.mocked(tag); -const mockTerminateRunner = vi.mocked(terminateRunner); -const mockUntagRunner = vi.mocked(untag); -const runnerOperations: Ec2RunnerResourceOperations = { +const mockListRunners = vi.fn(); +const mockCreateRunner = vi.fn(); +const mockTagRunner = vi.fn(); +const mockTerminateRunner = vi.fn(); +const mockUntagRunner = vi.fn(); +const ec2Operations: Ec2RunnerResourceOperations = { list: mockListRunners, - create: vi.fn(), + create: mockCreateRunner, terminate: mockTerminateRunner, tag: mockTagRunner, untag: mockUntagRunner, }; +const capability = createEc2ScaleDownCapability(ec2Operations); describe('Scale down runners', () => { beforeEach(() => { @@ -55,10 +45,8 @@ describe('Scale down runners', () => { mockListRunners.mockResolvedValueOnce([]).mockResolvedValueOnce([runner]); mockTagRunner.mockResolvedValue(); mockUntagRunner.mockResolvedValue(); - const provider = createEc2ScaleDownProvider(runnerOperations); - - await expect(provider.list('unit-test-environment')).resolves.toEqual([]); - await expect(provider.list('unit-test-environment', true)).resolves.toEqual([runner]); + await expect(capability.list('unit-test-environment')).resolves.toEqual([]); + await expect(capability.list('unit-test-environment', true)).resolves.toEqual([runner]); expect(mockListRunners).toHaveBeenNthCalledWith(1, { environment: 'unit-test-environment', orphan: undefined, @@ -66,8 +54,8 @@ describe('Scale down runners', () => { expect(mockListRunners).toHaveBeenNthCalledWith(2, { environment: 'unit-test-environment', orphan: true }); expect(mockTerminateRunner).not.toHaveBeenCalled(); - await provider.markOrphan(runner.id); - await provider.unmarkOrphan(runner.id); + await capability.markOrphan(runner.id); + await capability.unmarkOrphan(runner.id); expect(mockTagRunner).toHaveBeenCalledWith(runner.id, [{ Key: 'ghr:orphan', Value: 'true' }]); expect(mockUntagRunner).toHaveBeenCalledWith(runner.id, [{ Key: 'ghr:orphan', Value: 'true' }]); @@ -79,12 +67,11 @@ describe('Scale down runners', () => { launchTime: new Date(), }; process.env.RUNNER_BOOT_TIME_IN_MINUTES = '5'; - const provider = createEc2ScaleDownProvider(runnerOperations); - expect(provider.bootTimeExceeded(scaleDownRunner)).toBe(false); + expect(capability.bootTimeExceeded(scaleDownRunner)).toBe(false); expect(mockTerminateRunner).not.toHaveBeenCalled(); mockTerminateRunner.mockResolvedValue(); - await provider.terminate(runner.id); + await capability.terminate(runner.id); expect(mockTerminateRunner).toHaveBeenCalledWith(runner.id); }); diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.ts index 310e6e3e30..8171550f5c 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-down.ts @@ -1,14 +1,14 @@ import type { ScaleDownComputeProvider } from '../../../../core'; import { bootTimeExceeded, type Ec2RunnerResourceOperations } from '../runners'; -export function createEc2ScaleDownProvider( - runnerOperations: Ec2RunnerResourceOperations, +export function createEc2ScaleDownCapability( + ec2Operations: Ec2RunnerResourceOperations, ): Omit { return { - list: (environment, orphan) => runnerOperations.list({ environment, orphan }), + list: (environment, orphan) => ec2Operations.list({ environment, orphan }), bootTimeExceeded, - markOrphan: (id) => runnerOperations.tag(id, [{ Key: 'ghr:orphan', Value: 'true' }]), - unmarkOrphan: (id) => runnerOperations.untag(id, [{ Key: 'ghr:orphan', Value: 'true' }]), - terminate: (id) => runnerOperations.terminate(id), + markOrphan: (id) => ec2Operations.tag(id, [{ Key: 'ghr:orphan', Value: 'true' }]), + unmarkOrphan: (id) => ec2Operations.untag(id, [{ Key: 'ghr:orphan', Value: 'true' }]), + terminate: (id) => ec2Operations.terminate(id), }; } diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts index a11248a31b..7a25762397 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.test.ts @@ -1,43 +1,35 @@ import type { CreateGitHubRunnerConfig, CreateStartRunnerConfig, RunnerType } from '../../../../core'; import type { Octokit } from '@octokit/rest'; -import { DescribeLaunchTemplateVersionsCommand, EC2Client } from '@aws-sdk/client-ec2'; -import { mockClient } from 'aws-sdk-client-mock'; -import 'aws-sdk-client-mock-jest/vitest'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import { parseEc2OverrideConfig } from './dynamic-labels'; -import { EC2_TAG_VALUE_MAX_LENGTH, RUNNER_LABELS_TAG_MAX_COUNT } from './runner-config'; -import { createRunner, listEC2Runners, tag, terminateRunner } from '../runners'; -import type { Ec2RunnerResourceOperations } from '../runners'; +import { EC2_TAG_VALUE_MAX_LENGTH, RUNNER_LABELS_TAG_MAX_COUNT } from './runner-creation'; +import type { Ec2RunnerProvisioningOperations } from '../runners'; import type { RunnerInputParameters } from '../runners.d'; -import { createEc2ScaleUpProvider } from './scale-up'; - -vi.mock('../runners', () => ({ - createRunner: vi.fn(), - listEC2Runners: vi.fn(), - tag: vi.fn(), - terminateRunner: vi.fn(), -})); - -const mockCreateRunner = vi.mocked(createRunner); -const mockListRunners = vi.mocked(listEC2Runners); -const mockTag = vi.mocked(tag); -const mockTerminateRunner = vi.mocked(terminateRunner); -const mockEC2Client = mockClient(EC2Client); +import { createEc2ScaleUpCapability } from './scale-up'; + +const mockCreateRunner = vi.fn(); +const mockListRunners = vi.fn(); +const mockTag = vi.fn(); +const mockTerminateRunner = vi.fn(); +const mockUntag = vi.fn(); const mockCreateStartRunnerConfig = vi.fn(); +const mockGetDefaultBlockDeviceNameFromLaunchTemplate = + vi.fn(); const githubClient = {} as Octokit; const runnerOwner = 'Codertocat'; const repositoryRunnerOwner = 'Codertocat/hello-world'; const cleanEnv = process.env; -const runnerOperations: Ec2RunnerResourceOperations = { +const ec2Operations: Ec2RunnerProvisioningOperations = { list: mockListRunners, create: mockCreateRunner, terminate: mockTerminateRunner, tag: mockTag, - untag: vi.fn(), + untag: mockUntag, + getDefaultBlockDeviceNameFromLaunchTemplate: mockGetDefaultBlockDeviceNameFromLaunchTemplate, }; -const provider = createEc2ScaleUpProvider(runnerOperations, new EC2Client({}), mockCreateStartRunnerConfig); +const capability = createEc2ScaleUpCapability(ec2Operations, mockCreateStartRunnerConfig); interface CreateProviderRunnersOptions { labels?: string[]; @@ -94,14 +86,14 @@ function expectedRunnerParams( } async function createProviderRunners(options: CreateProviderRunnersOptions = {}) { - const runnerLabelResolution = await provider.resolveLabelsForRunners(options.labels ?? []); + const runnerLabelResolution = await capability.resolveLabelsForRunners(options.labels ?? []); const baseRunnerLabels = options.baseRunnerLabels ?? 'label1,label2'; const githubRunnerConfig = runnerConfig({ runnerLabels: [baseRunnerLabels, ...runnerLabelResolution.runnerLabels].filter(Boolean).join(','), ...options.githubRunnerConfig, }); - return await provider.createRunners({ + return await capability.createRunners({ githubRunnerConfig, numberOfRunners: 1, githubInstallationClient: githubClient, @@ -110,10 +102,13 @@ async function createProviderRunners(options: CreateProviderRunnersOptions = {}) } async function expectCurrentRunners(runnerType: RunnerType, owner: string) { - const runnerLabelResolution = await provider.resolveLabelsForRunners([]); + const runnerLabelResolution = await capability.resolveLabelsForRunners([]); await expect( - provider.getCurrentRunners(runnerLabelResolution.state, { runnerType, runnerOwner: owner }), + capability.getCurrentRunners(runnerLabelResolution.state, { + runnerType, + runnerOwner: owner, + }), ).resolves.toBe(1); expect(mockListRunners).toHaveBeenCalledWith({ environment: 'unit-test-environment', @@ -140,16 +135,7 @@ beforeEach(() => { delete process.env.ENABLE_ON_DEMAND_FAILOVER_FOR_ERRORS; delete process.env.USE_DEDICATED_HOST; - mockEC2Client.reset(); - mockEC2Client.on(DescribeLaunchTemplateVersionsCommand).resolves({ - LaunchTemplateVersions: [ - { - LaunchTemplateData: { - BlockDeviceMappings: [{ DeviceName: '/dev/sda1', Ebs: {} }], - }, - }, - ], - }); + mockGetDefaultBlockDeviceNameFromLaunchTemplate.mockResolvedValue('/dev/sda1'); mockCreateRunner.mockResolvedValue(createRunnerResult(['i-12345'])); mockListRunners.mockResolvedValue([ @@ -285,26 +271,12 @@ describe('scaleUp with GHES', () => { }); it('loads the launch template block device name for dynamic EBS labels without DeviceName', async () => { - mockEC2Client.on(DescribeLaunchTemplateVersionsCommand).resolves({ - LaunchTemplateVersions: [ - { - LaunchTemplateData: { - BlockDeviceMappings: [ - { DeviceName: '/dev/sdb', VirtualName: 'ephemeral0' }, - { DeviceName: '/dev/sdf', Ebs: {} }, - ], - }, - }, - ], - }); + mockGetDefaultBlockDeviceNameFromLaunchTemplate.mockResolvedValueOnce('/dev/sdf'); await createProviderRunners({ baseRunnerLabels: 'base-label', labels: ['ghr-ec2-ebs-volume-size:100', 'ghr-ec2-ebs-volume-type:gp3'], }); - expect(mockEC2Client).toHaveReceivedCommandWith(DescribeLaunchTemplateVersionsCommand, { - LaunchTemplateName: 'lt-1', - Versions: ['$Default'], - }); + expect(mockGetDefaultBlockDeviceNameFromLaunchTemplate).toHaveBeenCalledWith('lt-1'); expect(mockCreateRunner).toHaveBeenCalledWith( expect.objectContaining({ ec2OverrideConfig: expect.objectContaining({ @@ -319,7 +291,7 @@ describe('scaleUp with GHES', () => { baseRunnerLabels: 'base-label', labels: ['ghr-ec2-block-device-name:/dev/sdg', 'ghr-ec2-ebs-volume-size:100', 'ghr-ec2-ebs-volume-type:gp3'], }); - expect(mockEC2Client).not.toHaveReceivedCommand(DescribeLaunchTemplateVersionsCommand); + expect(mockGetDefaultBlockDeviceNameFromLaunchTemplate).not.toHaveBeenCalled(); expect(mockCreateRunner).toHaveBeenCalledWith( expect.objectContaining({ ec2OverrideConfig: expect.objectContaining({ diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.ts index d1bc416ca0..0cbda11a8f 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/scale-up.ts @@ -1,24 +1,12 @@ import { createChildLogger } from '@aws-github-runner/aws-powertools-util'; -import type { EC2Client } from '@aws-sdk/client-ec2'; -import type { - CreateRunnerResult, - CreateScaleUpRunnersInput, - CreateStartRunnerConfig, - CurrentRunnersInput, - RunnerLabelResolution, - ScaleUpComputeProvider, -} from '../../../../core'; +import type { CreateStartRunnerConfig, RunnerLabelResolution, ScaleUpComputeProvider } from '../../../../core'; import yn from 'yn'; -import type { Ec2RunnerResourceOperations } from '../runners'; +import type { Ec2RunnerProvisioningOperations } from '../runners'; import type { Ec2OverrideConfig } from '../runners.d'; -import { - getDefaultBlockDeviceNameFromLaunchTemplate, - parseEc2OverrideConfig, - shouldLoadLaunchTemplateBlockDeviceName, -} from './dynamic-labels'; -import { createRunners, loadEc2ProviderConfig } from './runner-config'; -import type { CreateEC2RunnerConfig } from './runner-config'; +import { parseEc2OverrideConfig, shouldLoadLaunchTemplateBlockDeviceName } from './dynamic-labels'; +import { createRunners, loadEc2ProviderConfig } from './runner-creation'; +import type { CreateEC2RunnerConfig } from './runner-creation'; const logger = createChildLogger('ec2-scale-up'); @@ -33,52 +21,45 @@ function loadEc2ScaleUpProviderConfig(): CreateEC2RunnerConfig { }; } -export function createEc2ScaleUpProvider( - runnerOperations: Ec2RunnerResourceOperations, - ec2Client: EC2Client, - createStartRunnerConfig: CreateStartRunnerConfig, -): Omit, 'type'> { - return { - async resolveLabelsForRunners(messageLabels: string[]): Promise> { - const trimmedLabels = messageLabels.map((label) => label.trim()); - const dynamicEC2Labels = trimmedLabels.filter((label) => label.startsWith('ghr-ec2-')); - const nonEc2DynamicLabels = trimmedLabels.filter( - (label) => label.startsWith('ghr-') && !label.startsWith('ghr-ec2-'), - ); - const runnerLabels = [...nonEc2DynamicLabels, ...dynamicEC2Labels]; - let ec2OverrideConfig: Ec2OverrideConfig | undefined; - - if (dynamicEC2Labels.length > 0) { - const defaultBlockDeviceName = shouldLoadLaunchTemplateBlockDeviceName(dynamicEC2Labels) - ? await getDefaultBlockDeviceNameFromLaunchTemplate(ec2Client, process.env.LAUNCH_TEMPLATE_NAME) - : undefined; +async function resolveEc2ScaleUpRunnerLabels( + ec2Operations: Ec2RunnerProvisioningOperations, + messageLabels: string[], +): Promise> { + const trimmedLabels = messageLabels.map((label) => label.trim()); + const dynamicEC2Labels = trimmedLabels.filter((label) => label.startsWith('ghr-ec2-')); + const nonEc2DynamicLabels = trimmedLabels.filter( + (label) => label.startsWith('ghr-') && !label.startsWith('ghr-ec2-'), + ); + const runnerLabels = [...nonEc2DynamicLabels, ...dynamicEC2Labels]; + let ec2OverrideConfig: Ec2OverrideConfig | undefined; - ec2OverrideConfig = parseEc2OverrideConfig(dynamicEC2Labels, defaultBlockDeviceName); - if (ec2OverrideConfig) { - logger.debug('EC2 override config parsed from labels', { ec2OverrideConfig }); - } - } + if (dynamicEC2Labels.length > 0) { + const defaultBlockDeviceName = shouldLoadLaunchTemplateBlockDeviceName(dynamicEC2Labels) + ? await ec2Operations.getDefaultBlockDeviceNameFromLaunchTemplate(process.env.LAUNCH_TEMPLATE_NAME) + : undefined; - return { runnerLabels, state: { ec2OverrideConfig } }; - }, + ec2OverrideConfig = parseEc2OverrideConfig(dynamicEC2Labels, defaultBlockDeviceName); + if (ec2OverrideConfig) { + logger.debug('EC2 override config parsed from labels', { ec2OverrideConfig }); + } + } - async getCurrentRunners( - _state: Ec2ScaleUpState, - { runnerType, runnerOwner }: CurrentRunnersInput, - ): Promise { - return (await runnerOperations.list({ environment: process.env.ENVIRONMENT, runnerType, runnerOwner })).length; - }, + return { runnerLabels, state: { ec2OverrideConfig } }; +} - async createRunners({ - githubRunnerConfig, - numberOfRunners, - githubInstallationClient, - state, - }: CreateScaleUpRunnersInput): Promise { +export function createEc2ScaleUpCapability( + ec2Operations: Ec2RunnerProvisioningOperations, + createStartRunnerConfig: CreateStartRunnerConfig, +): Omit, 'type'> { + return { + resolveLabelsForRunners: (labels) => resolveEc2ScaleUpRunnerLabels(ec2Operations, labels), + getCurrentRunners: async (_state, { runnerType, runnerOwner }) => + (await ec2Operations.list({ environment: process.env.ENVIRONMENT, runnerType, runnerOwner })).length, + createRunners: async ({ githubRunnerConfig, numberOfRunners, githubInstallationClient, state }) => { const config = loadEc2ScaleUpProviderConfig(); return await createRunners( - runnerOperations, + ec2Operations, githubRunnerConfig, { ...config, diff --git a/lambdas/libs/compute-providers/aws/ec2/src/launch-template.ts b/lambdas/libs/compute-providers/aws/ec2/src/launch-template.ts new file mode 100644 index 0000000000..0dcab42497 --- /dev/null +++ b/lambdas/libs/compute-providers/aws/ec2/src/launch-template.ts @@ -0,0 +1,26 @@ +import { DescribeLaunchTemplateVersionsCommand, type EC2Client } from '@aws-sdk/client-ec2'; + +export async function getDefaultBlockDeviceNameFromLaunchTemplate( + ec2Client: EC2Client, + launchTemplateName: string, + signal: AbortSignal | undefined, +): Promise { + const launchTemplateVersions = await ec2Client.send( + new DescribeLaunchTemplateVersionsCommand({ + LaunchTemplateName: launchTemplateName, + Versions: ['$Default'], + }), + { abortSignal: signal }, + ); + const blockDeviceMappings = + launchTemplateVersions.LaunchTemplateVersions?.[0]?.LaunchTemplateData?.BlockDeviceMappings; + const blockDeviceName = + blockDeviceMappings?.find((blockDeviceMapping) => blockDeviceMapping.DeviceName && blockDeviceMapping.Ebs) + ?.DeviceName ?? blockDeviceMappings?.find((blockDeviceMapping) => blockDeviceMapping.DeviceName)?.DeviceName; + + if (!blockDeviceName) { + throw new Error(`Failed to determine block device name from launch template '${launchTemplateName}'.`); + } + + return blockDeviceName; +} diff --git a/lambdas/libs/compute-providers/aws/ec2/src/runners.test.ts b/lambdas/libs/compute-providers/aws/ec2/src/runners.test.ts index d5a0d13d0a..4cc39e9f91 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/runners.test.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/runners.test.ts @@ -9,6 +9,7 @@ import { DeleteTagsCommand, DescribeInstancesCommand, type DescribeInstancesResult, + DescribeLaunchTemplateVersionsCommand, EC2Client, FleetOnDemandAllocationStrategy, RunInstancesCommand, @@ -22,12 +23,13 @@ import 'aws-sdk-client-mock-jest/vitest'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import type { RunnerInfo, RunnerSource, RunnerType } from '../../../core'; -import { createEc2RunnerClient, createRunner, listEC2Runners, tag, terminateRunner, untag } from './runners'; +import { createEc2RunnerClient } from './runners'; import type { Ec2OverrideConfig, RunnerInputParameters } from './runners.d'; process.env.AWS_REGION = 'eu-east-1'; const mockEC2Client = mockClient(EC2Client); const mockSSMClient = mockClient(SSMClient); +const ec2Operations = createEc2RunnerClient(new EC2Client({})).forRequest({ signal: undefined }); const LAUNCH_TEMPLATE = 'lt-1'; const ORG_NAME = 'SomeAwesomeCoder'; @@ -86,7 +88,7 @@ describe('list instances', () => { it('returns a list of instances (Non JIT)', async () => { mockEC2Client.on(DescribeInstancesCommand).resolves(mockRunningInstances); - const resp = await listEC2Runners(); + const resp = await ec2Operations.list(); expect(resp.length).toBe(1); expect(resp).toContainEqual({ id: 'i-1234', @@ -100,7 +102,7 @@ describe('list instances', () => { it('returns a list of instances (JIT)', async () => { mockEC2Client.on(DescribeInstancesCommand).resolves(mockRunningInstancesJit); - const resp = await listEC2Runners(); + const resp = await ec2Operations.list(); expect(resp.length).toBe(1); expect(resp).toContainEqual({ id: 'i-1234', @@ -121,7 +123,7 @@ describe('list instances', () => { }); mockEC2Client.on(DescribeInstancesCommand).resolves(instances); - const resp = await listEC2Runners(); + const resp = await ec2Operations.list(); expect(resp.length).toBe(1); expect(resp).toContainEqual({ id: instances.Reservations![0].Instances![0].InstanceId!, @@ -135,13 +137,13 @@ describe('list instances', () => { it('calls EC2 describe instances', async () => { mockEC2Client.on(DescribeInstancesCommand).resolves(mockRunningInstances); - await listEC2Runners(); + await ec2Operations.list(); expect(mockEC2Client).toHaveReceivedCommand(DescribeInstancesCommand); }); it('filters instances on repo name', async () => { mockEC2Client.on(DescribeInstancesCommand).resolves(mockRunningInstances); - await listEC2Runners({ + await ec2Operations.list({ runnerType: 'Repo', runnerOwner: REPO_NAME, environment: undefined, @@ -158,7 +160,7 @@ describe('list instances', () => { it('filters instances on org name', async () => { mockEC2Client.on(DescribeInstancesCommand).resolves(mockRunningInstances); - await listEC2Runners({ + await ec2Operations.list({ runnerType: 'Org', runnerOwner: ORG_NAME, environment: undefined, @@ -175,7 +177,7 @@ describe('list instances', () => { it('filters instances on environment', async () => { mockEC2Client.on(DescribeInstancesCommand).resolves(mockRunningInstances); - await listEC2Runners({ environment: ENVIRONMENT }); + await ec2Operations.list({ environment: ENVIRONMENT }); expect(mockEC2Client).toHaveReceivedCommandWith(DescribeInstancesCommand, { Filters: [ { Name: 'instance-state-name', Values: ['running', 'pending'] }, @@ -191,7 +193,7 @@ describe('list instances', () => { Value: 'true', }); mockEC2Client.on(DescribeInstancesCommand).resolves(mockRunningInstances); - await listEC2Runners({ environment: ENVIRONMENT, orphan: true }); + await ec2Operations.list({ environment: ENVIRONMENT, orphan: true }); expect(mockEC2Client).toHaveReceivedCommandWith(DescribeInstancesCommand, { Filters: [ { Name: 'instance-state-name', Values: ['running', 'pending'] }, @@ -207,7 +209,7 @@ describe('list instances', () => { Reservations: undefined, }; mockEC2Client.on(DescribeInstancesCommand).resolves(noInstances); - const resp = await listEC2Runners(); + const resp = await ec2Operations.list(); expect(resp.length).toBe(0); }); @@ -226,13 +228,13 @@ describe('list instances', () => { ], }; mockEC2Client.on(DescribeInstancesCommand).resolves(noInstances); - const resp = await listEC2Runners(); + const resp = await ec2Operations.list(); expect(resp.length).toBe(1); }); it('Filter instances for state running.', async () => { mockEC2Client.on(DescribeInstancesCommand).resolves(mockRunningInstances); - await listEC2Runners({ statuses: ['running'] }); + await ec2Operations.list({ statuses: ['running'] }); expect(mockEC2Client).toHaveReceivedCommandWith(DescribeInstancesCommand, { Filters: [ { Name: 'instance-state-name', Values: ['running'] }, @@ -243,7 +245,7 @@ describe('list instances', () => { it('Filter instances with status undefined, fall back to defaults.', async () => { mockEC2Client.on(DescribeInstancesCommand).resolves(mockRunningInstances); - await listEC2Runners({ statuses: undefined }); + await ec2Operations.list({ statuses: undefined }); expect(mockEC2Client).toHaveReceivedCommandWith(DescribeInstancesCommand, { Filters: [ { Name: 'instance-state-name', Values: ['running', 'pending'] }, @@ -264,7 +266,7 @@ describe('terminate runner', () => { owner: 'owner-2', type: 'Repo', }; - await terminateRunner(runner.id); + await ec2Operations.terminate(runner.id); expect(mockEC2Client).toHaveReceivedCommandWith(TerminateInstancesCommand, { InstanceIds: [runner.id], @@ -283,7 +285,7 @@ describe('tag runner', () => { owner: 'owner-2', type: 'Repo', }; - await tag(runner.id, [{ Key: 'ghr:orphan', Value: 'true' }]); + await ec2Operations.tag(runner.id, [{ Key: 'ghr:orphan', Value: 'true' }]); expect(mockEC2Client).toHaveReceivedCommandWith(CreateTagsCommand, { Resources: [runner.id], @@ -303,12 +305,12 @@ describe('untag runner', () => { owner: 'owner-2', type: 'Repo', }; - await tag(runner.id, [{ Key: 'ghr:orphan', Value: 'true' }]); + await ec2Operations.tag(runner.id, [{ Key: 'ghr:orphan', Value: 'true' }]); expect(mockEC2Client).toHaveReceivedCommandWith(CreateTagsCommand, { Resources: [runner.id], Tags: [{ Key: 'ghr:orphan', Value: 'true' }], }); - await untag(runner.id, [{ Key: 'ghr:orphan', Value: 'true' }]); + await ec2Operations.untag(runner.id, [{ Key: 'ghr:orphan', Value: 'true' }]); expect(mockEC2Client).toHaveReceivedCommandWith(DeleteTagsCommand, { Resources: [runner.id], Tags: [{ Key: 'ghr:orphan', Value: 'true' }], @@ -316,6 +318,38 @@ describe('untag runner', () => { }); }); +describe('runner client', () => { + it('loads the default EBS block device name through the bound EC2 client and request signal', async () => { + const send = vi.fn().mockResolvedValue({ + LaunchTemplateVersions: [ + { + LaunchTemplateData: { + BlockDeviceMappings: [ + { DeviceName: '/dev/sdb', VirtualName: 'ephemeral0' }, + { DeviceName: '/dev/sdf', Ebs: {} }, + ], + }, + }, + ], + }); + const ec2Client = { send } as unknown as EC2Client; + const abortController = new AbortController(); + const operations = createEc2RunnerClient(ec2Client).forRequest({ signal: abortController.signal }); + + await expect(operations.getDefaultBlockDeviceNameFromLaunchTemplate('lt-1')).resolves.toBe('/dev/sdf'); + expect(send).toHaveBeenCalledOnce(); + const [command, options] = send.mock.calls[0] as [ + DescribeLaunchTemplateVersionsCommand, + { abortSignal: AbortSignal }, + ]; + expect(command.input).toEqual({ + LaunchTemplateName: 'lt-1', + Versions: ['$Default'], + }); + expect(options).toEqual({ abortSignal: abortController.signal }); + }); +}); + describe('create runner', () => { const defaultRunnerConfig: RunnerConfig = { allocationStrategy: SpotAllocationStrategy.CAPACITY_OPTIMIZED, @@ -343,7 +377,7 @@ describe('create runner', () => { }); it.each(RUNNER_TYPES)('calls create fleet of 1 instance with the default config for %p', async (type: RunnerType) => { - await createRunner(createRunnerConfig({ ...defaultRunnerConfig, type: type })); + await ec2Operations.create(createRunnerConfig({ ...defaultRunnerConfig, type: type })); expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, { ...expectedCreateFleetRequest({ @@ -358,7 +392,7 @@ describe('create runner', () => { mockEC2Client.on(CreateFleetCommand).resolves({ Instances: instances }); - await createRunner({ + await ec2Operations.create({ ...createRunnerConfig(defaultRunnerConfig), numberOfRunners: 2, }); @@ -376,7 +410,7 @@ describe('create runner', () => { mockEC2Client.on(CreateFleetCommand).resolves({ Instances: instances }); - await createRunner({ + await ec2Operations.create({ ...createRunnerConfig({ ...defaultRunnerConfig, source: 'pool-lambda' }), numberOfRunners: 3, }); @@ -391,7 +425,7 @@ describe('create runner', () => { }); it('calls create fleet of 1 instance with the on-demand capacity', async () => { - await createRunner( + await ec2Operations.create( createRunnerConfig({ ...defaultRunnerConfig, capacityType: 'on-demand', allocationStrategy: 'lowest-price' }), ); expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, { @@ -404,7 +438,7 @@ describe('create runner', () => { }); it('calls create fleet with on-demand capacity and prioritized allocation strategy', async () => { - await createRunner( + await ec2Operations.create( createRunnerConfig({ ...defaultRunnerConfig, capacityType: 'on-demand', @@ -422,7 +456,7 @@ describe('create runner', () => { it('calls create fleet with custom instance type priorities', async () => { const priorities = { 'm5.large': 10, 'c5.large': 5 }; - await createRunner( + await ec2Operations.create( createRunnerConfig({ ...defaultRunnerConfig, capacityType: 'on-demand', @@ -442,7 +476,7 @@ describe('create runner', () => { it('calls create fleet with spot capacity-optimized-prioritized and instance type priorities', async () => { const priorities = { 'm5.large': 10, 'c5.large': 5 }; - await createRunner( + await ec2Operations.create( createRunnerConfig({ ...defaultRunnerConfig, capacityType: 'spot', @@ -461,7 +495,7 @@ describe('create runner', () => { }); it('calls run instances with the on-demand capacity', async () => { - await createRunner(createRunnerConfig({ ...defaultRunnerConfig, maxSpotPrice: '0.1' })); + await ec2Operations.create(createRunnerConfig({ ...defaultRunnerConfig, maxSpotPrice: '0.1' })); expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, { ...expectedCreateFleetRequest({ ...defaultExpectedFleetRequestValues, @@ -472,7 +506,7 @@ describe('create runner', () => { it('does not create ssm parameters when no instance is created', async () => { mockEC2Client.on(CreateFleetCommand).resolves({ Instances: [] }); - await expect(createRunner(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ + await expect(ec2Operations.create(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ instances: [], retryableErrorCount: 0, nonRetryableErrorCount: 1, @@ -487,7 +521,7 @@ describe('create runner', () => { }, }; mockSSMClient.on(GetParameterCommand).resolves(paramValue); - await createRunner( + await ec2Operations.create( createRunnerConfig({ ...defaultRunnerConfig, amiIdSsmParameterName: 'my-ami-id-param', @@ -506,12 +540,12 @@ describe('create runner', () => { it('keeps cancellation request-scoped and rejects before calling AWS', async () => { const abortController = new AbortController(); const abortReason = new Error('service stopping'); - const runnerOperations = createEc2RunnerClient(new EC2Client({})).forRequest({ + const ec2Operations = createEc2RunnerClient(new EC2Client({})).forRequest({ signal: abortController.signal, }); abortController.abort(abortReason); - await expect(runnerOperations.create(createRunnerConfig(defaultRunnerConfig))).rejects.toThrow('service stopping'); + await expect(ec2Operations.create(createRunnerConfig(defaultRunnerConfig))).rejects.toThrow('service stopping'); expect(mockEC2Client).not.toHaveReceivedCommand(CreateFleetCommand); expect(mockSSMClient).not.toHaveReceivedCommand(GetParameterCommand); }); @@ -532,7 +566,7 @@ describe('create runner', () => { it('calls create fleet of 1 instance with runner tracing enabled', async () => { tracer.getRootXrayTraceId = vi.fn().mockReturnValue('123'); - await createRunner(createRunnerConfig({ ...defaultRunnerConfig, tracingEnabled: true })); + await ec2Operations.create(createRunnerConfig({ ...defaultRunnerConfig, tracingEnabled: true })); expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, { ...expectedCreateFleetRequest({ @@ -543,7 +577,7 @@ 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' })); + await ec2Operations.create(createRunnerConfig({ ...defaultRunnerConfig, source: 'scale-up-lambda' })); expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, { ...expectedCreateFleetRequest({ @@ -554,7 +588,7 @@ describe('create runner', () => { }); it('calls create fleet with source set to pool-lambda when source is specified', async () => { - await createRunner(createRunnerConfig({ ...defaultRunnerConfig, source: 'pool-lambda' })); + await ec2Operations.create(createRunnerConfig({ ...defaultRunnerConfig, source: 'pool-lambda' })); expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, { ...expectedCreateFleetRequest({ @@ -565,7 +599,7 @@ describe('create runner', () => { }); it('overrides SubnetId when specified in ec2OverrideConfig', async () => { - await createRunner({ + await ec2Operations.create({ ...createRunnerConfig(defaultRunnerConfig), ec2OverrideConfig: { SubnetId: 'subnet-override', @@ -604,7 +638,7 @@ describe('create runner', () => { }); it('overrides InstanceType when specified in ec2OverrideConfig', async () => { - await createRunner({ + await ec2Operations.create({ ...createRunnerConfig(defaultRunnerConfig), ec2OverrideConfig: { InstanceType: 't3.xlarge', @@ -643,7 +677,7 @@ describe('create runner', () => { }); it('overrides ImageId when specified in ec2OverrideConfig', async () => { - await createRunner({ + await ec2Operations.create({ ...createRunnerConfig(defaultRunnerConfig), ec2OverrideConfig: { ImageId: 'ami-override-123', @@ -694,7 +728,7 @@ describe('create runner', () => { }); it('overrides all three fields (SubnetId, InstanceType, ImageId) when specified in ec2OverrideConfig', async () => { - await createRunner({ + await ec2Operations.create({ ...createRunnerConfig(defaultRunnerConfig), ec2OverrideConfig: { SubnetId: 'subnet-custom', @@ -732,7 +766,7 @@ describe('create runner', () => { }); it('spreads additional ec2OverrideConfig properties to Overrides', async () => { - await createRunner({ + await ec2Operations.create({ ...createRunnerConfig(defaultRunnerConfig), ec2OverrideConfig: { SubnetId: 'subnet-override', @@ -802,7 +836,7 @@ describe('create runner with errors', () => { it('returns one retryable error.', async () => { createFleetMockWithErrors(['UnfulfillableCapacity']); - await expect(createRunner(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ + await expect(ec2Operations.create(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ instances: [], retryableErrorCount: 1, nonRetryableErrorCount: 0, @@ -817,7 +851,7 @@ describe('create runner with errors', () => { it('returns a retryable error for a transient fleet result error without explicit configuration.', async () => { createFleetMockWithErrors(['InternalError']); - await expect(createRunner(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ + await expect(ec2Operations.create(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ instances: [], retryableErrorCount: 1, nonRetryableErrorCount: 0, @@ -827,7 +861,9 @@ describe('create runner with errors', () => { it('retries every missing instance when Fleet reports any retryable error.', async () => { createFleetMockWithErrors(['UnfulfillableCapacity', 'MaxSpotInstanceCountExceeded', 'NotMappedError']); - await expect(createRunner({ ...createRunnerConfig(defaultRunnerConfig), numberOfRunners: 3 })).resolves.toEqual({ + await expect( + ec2Operations.create({ ...createRunnerConfig(defaultRunnerConfig), numberOfRunners: 3 }), + ).resolves.toEqual({ instances: [], retryableErrorCount: 3, nonRetryableErrorCount: 0, @@ -843,7 +879,7 @@ describe('create runner with errors', () => { createFleetMockWithErrors(Array(12).fill('InsufficientFreeAddressesInSubnet')); await expect( - createRunner({ + ec2Operations.create({ ...createRunnerConfig({ ...defaultRunnerConfig, scaleErrors: [...defaultRunnerConfig.scaleErrors, 'InsufficientFreeAddressesInSubnet'], @@ -860,7 +896,7 @@ describe('create runner with errors', () => { it('returns a non-retryable error count for an unmapped error', async () => { createFleetMockWithErrors(['NonMappedError']); - await expect(createRunner(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ + await expect(ec2Operations.create(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ instances: [], retryableErrorCount: 0, nonRetryableErrorCount: 1, @@ -875,7 +911,7 @@ describe('create runner with errors', () => { it('returns a created instance without a failure count', async () => { createFleetMockWithErrors(['NonMappedError'], ['i-123']); - await expect(createRunner(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ + await expect(ec2Operations.create(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ instances: ['i-123'], retryableErrorCount: 0, nonRetryableErrorCount: 0, @@ -889,7 +925,7 @@ describe('create runner with errors', () => { it('returns a non-retryable error count when the create fleet request fails with an unknown exception.', async () => { mockEC2Client.on(CreateFleetCommand).rejects(new Error('Some error')); - await expect(createRunner(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ + await expect(ec2Operations.create(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ instances: [], retryableErrorCount: 0, nonRetryableErrorCount: 1, @@ -905,7 +941,7 @@ describe('create runner with errors', () => { const error = Object.assign(new Error('Not authorized'), { name: 'UnauthorizedOperation' }); mockEC2Client.on(CreateFleetCommand).rejects(error); - await expect(createRunner(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ + await expect(ec2Operations.create(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ instances: [], retryableErrorCount: 0, nonRetryableErrorCount: 1, @@ -917,7 +953,7 @@ describe('create runner with errors', () => { async (errorName) => { mockEC2Client.on(CreateFleetCommand).rejects(Object.assign(new Error(errorName), { name: errorName })); - await expect(createRunner(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ + await expect(ec2Operations.create(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ instances: [], retryableErrorCount: 0, nonRetryableErrorCount: 1, @@ -933,7 +969,7 @@ describe('create runner with errors', () => { }); mockEC2Client.on(CreateFleetCommand).rejects(error); - await expect(createRunner(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ + await expect(ec2Operations.create(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ instances: [], retryableErrorCount: 1, nonRetryableErrorCount: 0, @@ -946,7 +982,7 @@ describe('create runner with errors', () => { .rejects(Object.assign(new Error('Parameter does not exist'), { name: 'ParameterNotFound' })); await expect( - createRunner( + ec2Operations.create( createRunnerConfig({ ...defaultRunnerConfig, amiIdSsmParameterName: 'missing-ami-id-param', @@ -967,7 +1003,7 @@ describe('create runner with errors', () => { ); await expect( - createRunner( + ec2Operations.create( createRunnerConfig({ ...defaultRunnerConfig, amiIdSsmParameterName: 'my-ami-id-param', @@ -982,7 +1018,7 @@ describe('create runner with errors', () => { mockSSMClient.on(GetParameterCommand).rejects(new Error('Some error')); await expect( - createRunner( + ec2Operations.create( createRunnerConfig({ ...defaultRunnerConfig, amiIdSsmParameterName: 'my-ami-id-param', @@ -995,7 +1031,7 @@ describe('create runner with errors', () => { it('returns a non-scale error count with undefined Instances and Errors.', async () => { mockEC2Client.on(CreateFleetCommand).resolvesOnce({ Instances: undefined, Errors: undefined }); - await expect(createRunner(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ + await expect(ec2Operations.create(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ instances: [], retryableErrorCount: 0, nonRetryableErrorCount: 1, @@ -1011,7 +1047,7 @@ describe('create runner with errors', () => { }, ], }); - await expect(createRunner(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ + await expect(ec2Operations.create(createRunnerConfig(defaultRunnerConfig))).resolves.toEqual({ instances: [], retryableErrorCount: 0, nonRetryableErrorCount: 1, @@ -1049,7 +1085,7 @@ describe('create runner with errors fail over to OnDemand', () => { const instancesIds = ['i-123']; createFleetMockWithWithOnDemandFallback(['InsufficientInstanceCapacity'], instancesIds); - const instancesResult = await createRunner(createRunnerConfig(defaultRunnerConfig)); + const instancesResult = await ec2Operations.create(createRunnerConfig(defaultRunnerConfig)); expect(instancesResult).toEqual({ instances: instancesIds, retryableErrorCount: 0, @@ -1080,7 +1116,7 @@ describe('create runner with errors fail over to OnDemand', () => { it('test InsufficientInstanceCapacity no fallback.', async () => { await expect( - createRunner( + ec2Operations.create( createRunnerConfig({ ...defaultRunnerConfig, onDemandFailoverOnError: [], @@ -1093,7 +1129,7 @@ describe('create runner with errors fail over to OnDemand', () => { const instancesIds = ['i-123', 'i-456']; createFleetMockWithWithOnDemandFallback(['InsufficientInstanceCapacity'], instancesIds); - const instancesResult = await createRunner({ + const instancesResult = await ec2Operations.create({ ...createRunnerConfig(defaultRunnerConfig), numberOfRunners: 2, }); @@ -1131,7 +1167,7 @@ describe('create runner with errors fail over to OnDemand', () => { createFleetMockWithWithOnDemandFallback(['UnfulfillableCapacity'], instancesIds); await expect( - createRunner({ + ec2Operations.create({ ...createRunnerConfig(defaultRunnerConfig), numberOfRunners: 2, }), @@ -1363,7 +1399,7 @@ describe('create runner with useDedicatedHost', () => { }); it('uses RunInstances instead of CreateFleet when useDedicatedHost is true', async () => { - const result = await createRunner(createRunnerConfig(dedicatedHostRunnerConfig)); + const result = await ec2Operations.create(createRunnerConfig(dedicatedHostRunnerConfig)); expect(result).toEqual({ instances: ['i-dedicated-1'], @@ -1377,7 +1413,7 @@ describe('create runner with useDedicatedHost', () => { it('uses CreateFleet when useDedicatedHost is false', async () => { mockEC2Client.on(CreateFleetCommand).resolves({ Instances: [{ InstanceIds: ['i-fleet-1'] }] }); - const result = await createRunner( + const result = await ec2Operations.create( createRunnerConfig({ ...dedicatedHostRunnerConfig, useDedicatedHost: false, @@ -1392,7 +1428,7 @@ describe('create runner with useDedicatedHost', () => { it('uses CreateFleet when useDedicatedHost is undefined', async () => { mockEC2Client.on(CreateFleetCommand).resolves({ Instances: [{ InstanceIds: ['i-fleet-1'] }] }); - const result = await createRunner( + const result = await ec2Operations.create( createRunnerConfig({ ...dedicatedHostRunnerConfig, useDedicatedHost: undefined, @@ -1405,7 +1441,7 @@ describe('create runner with useDedicatedHost', () => { }); it('passes correct parameters to RunInstances', async () => { - await createRunner(createRunnerConfig(dedicatedHostRunnerConfig)); + await ec2Operations.create(createRunnerConfig(dedicatedHostRunnerConfig)); expect(mockEC2Client).toHaveReceivedCommandWith(RunInstancesCommand, { LaunchTemplate: { @@ -1444,7 +1480,7 @@ describe('create runner with useDedicatedHost', () => { Instances: [{ InstanceId: 'i-dedicated-1' }, { InstanceId: 'i-dedicated-2' }], }); - const result = await createRunner({ + const result = await ec2Operations.create({ ...createRunnerConfig(dedicatedHostRunnerConfig), numberOfRunners: 2, source: 'scale-up-lambda', @@ -1489,7 +1525,7 @@ describe('create runner with useDedicatedHost', () => { it('returns a non-retryable failure when spot is used with dedicated host', async () => { await expect( - createRunner( + ec2Operations.create( createRunnerConfig({ ...dedicatedHostRunnerConfig, capacityType: 'spot', @@ -1502,7 +1538,7 @@ describe('create runner with useDedicatedHost', () => { it('returns a non-retryable failure when RunInstances returns no instances', async () => { mockEC2Client.on(RunInstancesCommand).resolves({ Instances: [] }); - await expect(createRunner(createRunnerConfig(dedicatedHostRunnerConfig))).resolves.toEqual({ + await expect(ec2Operations.create(createRunnerConfig(dedicatedHostRunnerConfig))).resolves.toEqual({ instances: [], retryableErrorCount: 0, nonRetryableErrorCount: 1, @@ -1512,7 +1548,7 @@ describe('create runner with useDedicatedHost', () => { it('returns a non-retryable failure when RunInstances fails with an unknown exception', async () => { mockEC2Client.on(RunInstancesCommand).rejects(new Error('EC2 error')); - await expect(createRunner(createRunnerConfig(dedicatedHostRunnerConfig))).resolves.toEqual({ + await expect(ec2Operations.create(createRunnerConfig(dedicatedHostRunnerConfig))).resolves.toEqual({ instances: [], retryableErrorCount: 0, nonRetryableErrorCount: 1, @@ -1523,7 +1559,7 @@ describe('create runner with useDedicatedHost', () => { const error = Object.assign(new Error('Invalid subnet'), { name: 'InvalidSubnetID.NotFound' }); mockEC2Client.on(RunInstancesCommand).rejects(error); - await expect(createRunner(createRunnerConfig(dedicatedHostRunnerConfig))).resolves.toEqual({ + await expect(ec2Operations.create(createRunnerConfig(dedicatedHostRunnerConfig))).resolves.toEqual({ instances: [], retryableErrorCount: 0, nonRetryableErrorCount: 1, @@ -1534,7 +1570,7 @@ describe('create runner with useDedicatedHost', () => { const error = Object.assign(new Error('Connection reset'), { code: 'ECONNRESET' }); mockEC2Client.on(RunInstancesCommand).rejects(error); - await expect(createRunner(createRunnerConfig(dedicatedHostRunnerConfig))).resolves.toEqual({ + await expect(ec2Operations.create(createRunnerConfig(dedicatedHostRunnerConfig))).resolves.toEqual({ instances: [], retryableErrorCount: 1, nonRetryableErrorCount: 0, @@ -1546,7 +1582,7 @@ describe('create runner with useDedicatedHost', () => { mockEC2Client.on(RunInstancesCommand).rejects(error); await expect( - createRunner({ + ec2Operations.create({ ...createRunnerConfig({ ...dedicatedHostRunnerConfig, scaleErrors: ['InsufficientInstanceCapacity'], @@ -1562,7 +1598,7 @@ describe('create runner with useDedicatedHost', () => { }); await expect( - createRunner({ + ec2Operations.create({ ...createRunnerConfig(dedicatedHostRunnerConfig), numberOfRunners: 2, }), @@ -1581,7 +1617,7 @@ describe('create runner with useDedicatedHost', () => { }; mockSSMClient.on(GetParameterCommand).resolves(paramValue); - await createRunner( + await ec2Operations.create( createRunnerConfig({ ...dedicatedHostRunnerConfig, amiIdSsmParameterName: 'my-ami-id-param', @@ -1666,7 +1702,7 @@ describe('create runner with useDedicatedHost', () => { WeightedCapacity: 2, }; - await createRunner( + await ec2Operations.create( createRunnerConfig({ ...dedicatedHostRunnerConfig, amiIdSsmParameterName: 'my-ami-id-param', diff --git a/lambdas/libs/compute-providers/aws/ec2/src/runners.ts b/lambdas/libs/compute-providers/aws/ec2/src/runners.ts index f0ef27c9fa..f1df6a19cd 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/runners.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/runners.ts @@ -9,7 +9,7 @@ import { RunInstancesCommand, type RunInstancesCommandInput, RunInstancesCommandOutput, - EC2Client, + type EC2Client, FleetLaunchTemplateOverridesRequest, FleetOnDemandAllocationStrategy, SpotAllocationStrategy, @@ -17,12 +17,12 @@ import { TerminateInstancesCommand, _InstanceType, } from '@aws-sdk/client-ec2'; -import { createChildLogger } from '@aws-github-runner/aws-powertools-util'; -import { getTracedAWSV3Client, tracer } from '@aws-github-runner/aws-powertools-util'; +import { createChildLogger, tracer } from '@aws-github-runner/aws-powertools-util'; import { getParameter } from '@aws-github-runner/aws-ssm-util'; import moment from 'moment'; import type { CreateRunnerResult, RunnerInfo } from '../../../core'; +import { getDefaultBlockDeviceNameFromLaunchTemplate } from './launch-template'; import type { Ec2ListRunnerFilters, Ec2OverrideConfig, RunnerInputParameters } from './runners.d'; const logger = createChildLogger('runners'); @@ -44,8 +44,12 @@ export interface Ec2RunnerResourceOperations { untag(instanceId: string, tags: Tag[]): Promise; } +export interface Ec2RunnerProvisioningOperations extends Ec2RunnerResourceOperations { + getDefaultBlockDeviceNameFromLaunchTemplate(launchTemplateName: string): Promise; +} + export interface Ec2RunnerClient { - forRequest(context: Ec2RunnerRequestContext): Ec2RunnerResourceOperations; + forRequest(context: Ec2RunnerRequestContext): Ec2RunnerProvisioningOperations; } async function runWithRequestSignal( @@ -59,33 +63,24 @@ async function runWithRequestSignal( export function createEc2RunnerClient(ec2Client: EC2Client): Ec2RunnerClient { return { forRequest: ({ signal }) => ({ - list: (filters) => runWithRequestSignal(signal, () => listRunners(ec2Client, filters, signal)), + list: (filters) => runWithRequestSignal(signal, () => listEc2Runners(ec2Client, filters, signal)), create: (runnerParameters) => runWithRequestSignal(signal, () => createEc2Runner(ec2Client, runnerParameters, signal)), terminate: (instanceId) => runWithRequestSignal(signal, () => terminateEc2Runner(ec2Client, instanceId, signal)), tag: (instanceId, tags) => runWithRequestSignal(signal, () => tagEc2Runner(ec2Client, instanceId, tags, signal)), untag: (instanceId, tags) => runWithRequestSignal(signal, () => untagEc2Runner(ec2Client, instanceId, tags, signal)), + getDefaultBlockDeviceNameFromLaunchTemplate: (launchTemplateName) => + runWithRequestSignal(signal, () => + getDefaultBlockDeviceNameFromLaunchTemplate(ec2Client, launchTemplateName, signal), + ), }), }; } -let defaultRunnerOperations: Ec2RunnerResourceOperations | undefined; - -function getDefaultRunnerOperations(): Ec2RunnerResourceOperations { - defaultRunnerOperations ??= createEc2RunnerClient( - getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION })), - ).forRequest({ signal: undefined }); - return defaultRunnerOperations; -} - type FleetError = NonNullable[number]; -export async function listEC2Runners(filters: Ec2ListRunnerFilters | undefined = undefined): Promise { - return await getDefaultRunnerOperations().list(filters); -} - -async function listRunners( +async function listEc2Runners( ec2Client: EC2Client, filters: Ec2ListRunnerFilters | undefined, signal: AbortSignal | undefined, @@ -167,10 +162,6 @@ function getRunnerInfo(runningInstances: DescribeInstancesResult) { return runners; } -export async function terminateRunner(instanceId: string): Promise { - await getDefaultRunnerOperations().terminate(instanceId); -} - async function terminateEc2Runner( ec2Client: EC2Client, instanceId: string, @@ -181,10 +172,6 @@ async function terminateEc2Runner( logger.debug(`Runner ${instanceId} has been terminated.`); } -export async function tag(instanceId: string, tags: Tag[]): Promise { - await getDefaultRunnerOperations().tag(instanceId, tags); -} - async function tagEc2Runner( ec2Client: EC2Client, instanceId: string, @@ -195,10 +182,6 @@ async function tagEc2Runner( await ec2Client.send(new CreateTagsCommand({ Resources: [instanceId], Tags: tags }), { abortSignal: signal }); } -export async function untag(instanceId: string, tags: Tag[]): Promise { - await getDefaultRunnerOperations().untag(instanceId, tags); -} - async function untagEc2Runner( ec2Client: EC2Client, instanceId: string, @@ -380,10 +363,6 @@ function buildRunInstancesOverrides( return overrides; } -export async function createRunner(runnerParameters: RunnerInputParameters): Promise { - return await getDefaultRunnerOperations().create(runnerParameters); -} - async function createEc2Runner( ec2Client: EC2Client, runnerParameters: RunnerInputParameters, diff --git a/lambdas/libs/compute-providers/package.json b/lambdas/libs/compute-providers/package.json index 554b4abdbf..c1806818cc 100644 --- a/lambdas/libs/compute-providers/package.json +++ b/lambdas/libs/compute-providers/package.json @@ -11,7 +11,7 @@ "./aws/ec2/webhook": "./aws/ec2/webhook.ts", "./aws/ec2/control-plane": "./aws/ec2/control-plane.ts", "./aws/ec2/runners": "./aws/ec2/src/runners.ts", - "./aws/ec2/control-plane/runner-config": "./aws/ec2/src/control-plane/runner-config.ts" + "./aws/ec2/control-plane/runner-creation": "./aws/ec2/src/control-plane/runner-creation.ts" }, "type": "module", "license": "MIT", diff --git a/lambdas/libs/compute-providers/templates/provider/control-plane.ts b/lambdas/libs/compute-providers/templates/provider/control-plane.ts index 418ac0dbf9..312bad0168 100644 --- a/lambdas/libs/compute-providers/templates/provider/control-plane.ts +++ b/lambdas/libs/compute-providers/templates/provider/control-plane.ts @@ -16,7 +16,7 @@ function notImplemented(operation: string): never { throw new Error(`Template compute provider must implement ${operation}`); } -export function createTemplatePoolProvider( +export function createTemplatePoolCapability( createStartRunnerConfig: CreateStartRunnerConfig, ): Omit { return { @@ -31,7 +31,7 @@ export function createTemplatePoolProvider( }; } -export function createTemplateScaleUpProvider( +export function createTemplateScaleUpCapability( createStartRunnerConfig: CreateStartRunnerConfig, ): Omit { return { @@ -54,7 +54,7 @@ export function createTemplateScaleUpProvider( }; } -export function createTemplateScaleDownProvider(): Omit { +export function createTemplateScaleDownCapability(): Omit { return { list: async (environment, orphan) => { void environment; @@ -77,9 +77,9 @@ export function createTemplateControlPlanePlugin( return { type: 'template', capabilities: { - pool: () => createTemplatePoolProvider(createStartRunnerConfig), - scaleUp: () => createTemplateScaleUpProvider(createStartRunnerConfig), - scaleDown: createTemplateScaleDownProvider, + pool: () => createTemplatePoolCapability(createStartRunnerConfig), + scaleUp: () => createTemplateScaleUpCapability(createStartRunnerConfig), + scaleDown: createTemplateScaleDownCapability, }, }; }