diff --git a/lambdas/libs/compute-providers/aws/ec2/control-plane.ts b/lambdas/libs/compute-providers/aws/ec2/control-plane.ts index 6b4fb7c0bf..07b5ce1f1d 100644 --- a/lambdas/libs/compute-providers/aws/ec2/control-plane.ts +++ b/lambdas/libs/compute-providers/aws/ec2/control-plane.ts @@ -1,20 +1,26 @@ 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 { 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, ): ComputeProviderPlugin { + const ec2Client = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION })); + const ec2Operations = createEc2RunnerClient(ec2Client).forRequest({ signal: undefined }); + return { type: 'ec2', capabilities: { - pool: () => createEc2PoolProvider(createStartRunnerConfig), - scaleUp: () => createEc2ScaleUpProvider(createStartRunnerConfig), - scaleDown: createEc2ScaleDownProvider, + 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 8de7882e24..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, - EC2Client, FleetBlockDeviceMappingRequest, FleetEbsBlockDeviceRequest, InstanceGeneration, @@ -30,7 +28,6 @@ import { VCpuCountRangeRequest, VolumeType, } from '@aws-sdk/client-ec2'; -import { getTracedAWSV3Client } from '@aws-github-runner/aws-powertools-util'; import { Ec2OverrideConfig } from '../runners.d'; @@ -365,24 +362,3 @@ 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 })); - 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 5667cf9bc0..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,26 +1,34 @@ import type { Octokit } from '@octokit/rest'; import type { CreateGitHubRunnerConfig, CreateStartRunnerConfig, RunnerInfo } from '../../../../core'; -import { bootTimeExceeded, listEC2Runners } from '../runners'; -import { calculateEc2PoolSize, createEc2PoolProvider } from './pool'; -import { createRunners, type Ec2ProviderConfig, loadEc2ProviderConfig } from './runner-config'; +import { bootTimeExceeded, type Ec2RunnerResourceOperations } from '../runners'; +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(), - listEC2Runners: vi.fn(), })); -vi.mock('./runner-config', () => ({ +vi.mock('./runner-creation', () => ({ createRunners: vi.fn(), loadEc2ProviderConfig: vi.fn(), })); const mockBootTimeExceeded = vi.mocked(bootTimeExceeded); -const mockListRunners = vi.mocked(listEC2Runners); const mockCreateRunners = vi.mocked(createRunners); const mockLoadProviderConfig = vi.mocked(loadEc2ProviderConfig); -describe('calculateEc2PoolSize', () => { +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('createEc2PoolCapability.countAvailableRunners', () => { beforeEach(() => { vi.clearAllMocks(); }); @@ -29,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(); }); @@ -43,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(); }); @@ -51,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(); }); @@ -59,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, @@ -105,42 +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' }]; - mockListRunners.mockResolvedValue(runners); - const provider = createEc2PoolProvider(createStartRunnerConfig); - - await expect( - provider.listRunners({ - environment: 'test-environment', - runnerOwner: 'owner', - runnerType: 'Org', - }), - ).resolves.toBe(runners); - expect(mockListRunners).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(createStartRunnerConfig); await expect( - provider.createRunners({ + capability.createRunners({ githubRunnerConfig, numberOfRunners: 1, githubInstallationClient, }), ).resolves.toEqual(['i-created']); expect(mockCreateRunners).toHaveBeenCalledWith( + 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 d1d243ca51..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,67 +1,11 @@ import { createChildLogger } from '@aws-github-runner/aws-powertools-util'; -import type { - CreateStartRunnerConfig, - CreatePoolRunnersInput, - ListPoolRunnersInput, - PoolComputeProvider, - RunnerInfo, - RunnerStatus, -} from '../../../../core'; -import { createRunners, loadEc2ProviderConfig } from './runner-config'; -import { bootTimeExceeded, listEC2Runners } from '../runners'; +import type { CreateStartRunnerConfig, PoolComputeProvider, RunnerInfo, RunnerStatus } from '../../../../core'; +import { bootTimeExceeded, type Ec2RunnerResourceOperations } from '../runners'; +import { createRunners, loadEc2ProviderConfig } from './runner-creation'; const logger = createChildLogger('pool'); -async function listEc2PoolRunners({ - environment, - runnerOwner, - runnerType, -}: ListPoolRunnersInput): Promise { - return await listEC2Runners({ - environment, - runnerOwner, - runnerType, - statuses: ['running'], - }); -} - -async function createEc2PoolRunners( - { githubRunnerConfig, numberOfRunners, githubInstallationClient }: CreatePoolRunnersInput, - createStartRunnerConfig: CreateStartRunnerConfig, -): Promise { - const config = loadEc2ProviderConfig(); - - const { instances } = await createRunners( - 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( - createStartRunnerConfig: CreateStartRunnerConfig, -): Omit, 'type'> { - return { - listRunners: listEc2PoolRunners, - countAvailableRunners: calculateEc2PoolSize, - createRunners: (input) => createEc2PoolRunners(input, createStartRunnerConfig), - }; -} - -export function calculateEc2PoolSize( +function countAvailableEc2PoolRunners( ec2runners: RunnerInfo[], runnerStatus: Map, includeBusyRunners = false, @@ -86,3 +30,42 @@ export function calculateEc2PoolSize( } return numberOfRunnersInPool; } + +export function createEc2PoolCapability( + ec2Operations: Ec2RunnerResourceOperations, + createStartRunnerConfig: CreateStartRunnerConfig, +): Omit, 'type'> { + return { + listRunners: ({ environment, runnerOwner, runnerType }) => + ec2Operations.list({ + environment, + runnerOwner, + runnerType, + statuses: ['running'], + }), + countAvailableRunners: countAvailableEc2PoolRunners, + createRunners: async ({ githubRunnerConfig, numberOfRunners, githubInstallationClient }) => { + const config = loadEc2ProviderConfig(); + + const { instances } = await createRunners( + ec2Operations, + 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-creation.ts similarity index 87% 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 34f95aaf97..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 @@ -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 { Ec2RunnerResourceOperations } 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( + ec2Operations: Ec2RunnerResourceOperations, githubRunnerConfig: CreateGitHubRunnerConfig, ec2RunnerConfig: CreateEC2RunnerConfig, numberOfRunners: number, @@ -70,7 +71,7 @@ export async function createRunners( ): Promise { let result: CreateRunnerResult; try { - result = await createRunner({ + result = await ec2Operations.create({ runnerType: githubRunnerConfig.runnerType, runnerOwner: githubRunnerConfig.runnerOwner, numberOfRunners, @@ -93,7 +94,7 @@ export async function createRunners( githubRunnerConfig, result.instances, ghClient, - createEc2StartRunnerConfigOptions(), + createEc2StartRunnerConfigOptions(ec2Operations), ); } 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(ec2Operations, failedInstances); return { instances: result.instances.filter((id) => !failedInstances.includes(id)), @@ -126,10 +127,13 @@ export async function createRunners( return result; } -async function terminateFailedInstances(instanceIds: string[]): Promise { +async function terminateFailedInstances( + ec2Operations: Ec2RunnerResourceOperations, + instanceIds: string[], +): Promise { for (const instanceId of instanceIds) { try { - await terminateRunner(instanceId); + await ec2Operations.terminate(instanceId); } catch (error) { logger.error('Failed to terminate instance', { instanceId, @@ -139,21 +143,25 @@ async function terminateFailedInstances(instanceIds: string[]): Promise { } } -function createEc2StartRunnerConfigOptions(): StartRunnerConfigOptions { +function createEc2StartRunnerConfigOptions(ec2Operations: Ec2RunnerResourceOperations): StartRunnerConfigOptions { return { getSsmParameterTags: (instanceId) => [{ Key: 'InstanceId', Value: instanceId }], - onJitConfigCreated: async (instanceId, metadata) => await tagEc2RunnerMetadata(instanceId, metadata), + onJitConfigCreated: async (instanceId, metadata) => await tagEc2RunnerMetadata(ec2Operations, instanceId, metadata), }; } -async function tagEc2RunnerMetadata(instanceId: string, metadata: GitHubRunnerMetadata): Promise { +async function tagEc2RunnerMetadata( + ec2Operations: Ec2RunnerResourceOperations, + instanceId: string, + metadata: GitHubRunnerMetadata, +): Promise { const tags = [ { Key: 'ghr:github_runner_id', Value: metadata.githubRunnerId }, ...generateRunnerLabelsTags(metadata.runnerLabels), ]; try { - await 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 f122f3b3a0..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,24 +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 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: mockCreateRunner, + terminate: mockTerminateRunner, + tag: mockTagRunner, + untag: mockUntagRunner, +}; +const capability = createEc2ScaleDownCapability(ec2Operations); describe('Scale down runners', () => { beforeEach(() => { @@ -47,10 +45,8 @@ describe('Scale down runners', () => { mockListRunners.mockResolvedValueOnce([]).mockResolvedValueOnce([runner]); mockTagRunner.mockResolvedValue(); mockUntagRunner.mockResolvedValue(); - const provider = createEc2ScaleDownProvider(); - - 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, @@ -58,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' }]); @@ -71,12 +67,11 @@ describe('Scale down runners', () => { launchTime: new Date(), }; process.env.RUNNER_BOOT_TIME_IN_MINUTES = '5'; - const provider = createEc2ScaleDownProvider(); - 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 3084eebe21..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,24 +1,14 @@ -import type { RunnerInfo, ScaleDownComputeProvider } from '../../../../core'; -import { bootTimeExceeded, listEC2Runners, tag, terminateRunner, untag } from '../runners'; +import type { ScaleDownComputeProvider } from '../../../../core'; +import { bootTimeExceeded, type Ec2RunnerResourceOperations } from '../runners'; -async function listEc2ScaleDownRunners(environment: string, orphan?: boolean): Promise { - return await listEC2Runners({ environment, orphan }); -} - -async function markEc2RunnerOrphan(id: string): Promise { - await tag(id, [{ Key: 'ghr:orphan', Value: 'true' }]); -} - -async function unmarkEc2RunnerOrphan(id: string): Promise { - await untag(id, [{ Key: 'ghr:orphan', Value: 'true' }]); -} - -export function createEc2ScaleDownProvider(): Omit { +export function createEc2ScaleDownCapability( + ec2Operations: Ec2RunnerResourceOperations, +): Omit { return { - list: listEc2ScaleDownRunners, + list: (environment, orphan) => ec2Operations.list({ environment, orphan }), bootTimeExceeded, - markOrphan: markEc2RunnerOrphan, - unmarkOrphan: unmarkEc2RunnerOrphan, - terminate: terminateRunner, + 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 1a967fc4cc..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,35 +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 { 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 provider = createEc2ScaleUpProvider(mockCreateStartRunnerConfig); +const ec2Operations: Ec2RunnerProvisioningOperations = { + list: mockListRunners, + create: mockCreateRunner, + terminate: mockTerminateRunner, + tag: mockTag, + untag: mockUntag, + getDefaultBlockDeviceNameFromLaunchTemplate: mockGetDefaultBlockDeviceNameFromLaunchTemplate, +}; +const capability = createEc2ScaleUpCapability(ec2Operations, mockCreateStartRunnerConfig); interface CreateProviderRunnersOptions { labels?: string[]; @@ -86,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, @@ -102,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', @@ -132,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([ @@ -277,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({ @@ -311,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 6bebaa31ba..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,23 +1,12 @@ import { createChildLogger } from '@aws-github-runner/aws-powertools-util'; -import type { - CreateRunnerResult, - CreateScaleUpRunnersInput, - CreateStartRunnerConfig, - CurrentRunnersInput, - RunnerLabelResolution, - ScaleUpComputeProvider, -} from '../../../../core'; +import type { CreateStartRunnerConfig, RunnerLabelResolution, ScaleUpComputeProvider } from '../../../../core'; import yn from 'yn'; -import { listEC2Runners } 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'); @@ -32,7 +21,10 @@ function loadEc2ScaleUpProviderConfig(): CreateEC2RunnerConfig { }; } -async function resolveEc2LabelsForRunners(messageLabels: string[]): Promise> { +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( @@ -43,7 +35,7 @@ async function resolveEc2LabelsForRunners(messageLabels: string[]): Promise 0) { const defaultBlockDeviceName = shouldLoadLaunchTemplateBlockDeviceName(dynamicEC2Labels) - ? await getDefaultBlockDeviceNameFromLaunchTemplate(process.env.LAUNCH_TEMPLATE_NAME) + ? await ec2Operations.getDefaultBlockDeviceNameFromLaunchTemplate(process.env.LAUNCH_TEMPLATE_NAME) : undefined; ec2OverrideConfig = parseEc2OverrideConfig(dynamicEC2Labels, defaultBlockDeviceName); @@ -55,38 +47,29 @@ async function resolveEc2LabelsForRunners(messageLabels: string[]): Promise { - return (await listEC2Runners({ environment: process.env.ENVIRONMENT, runnerType, runnerOwner })).length; -} - -async function createEc2ScaleUpRunners( - { githubRunnerConfig, numberOfRunners, githubInstallationClient, state }: CreateScaleUpRunnersInput, - createStartRunnerConfig: CreateStartRunnerConfig, -): Promise { - const config = loadEc2ScaleUpProviderConfig(); - - return await createRunners( - githubRunnerConfig, - { - ...config, - ec2OverrideConfig: state.ec2OverrideConfig, - }, - numberOfRunners, - githubInstallationClient, - createStartRunnerConfig, - 'scale-up-lambda', - ); -} - -export function createEc2ScaleUpProvider( +export function createEc2ScaleUpCapability( + ec2Operations: Ec2RunnerProvisioningOperations, createStartRunnerConfig: CreateStartRunnerConfig, ): Omit, 'type'> { return { - resolveLabelsForRunners: resolveEc2LabelsForRunners, - getCurrentRunners: getCurrentEc2Runners, - createRunners: (input) => createEc2ScaleUpRunners(input, createStartRunnerConfig), + 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( + ec2Operations, + githubRunnerConfig, + { + ...config, + ec2OverrideConfig: state.ec2OverrideConfig, + }, + numberOfRunners, + githubInstallationClient, + createStartRunnerConfig, + 'scale-up-lambda', + ); + }, }; } 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 affd063c39..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 { 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', @@ -502,10 +536,37 @@ 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 ec2Operations = createEc2RunnerClient(new EC2Client({})).forRequest({ + signal: abortController.signal, + }); + abortController.abort(abortReason); + + await expect(ec2Operations.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'); - await createRunner(createRunnerConfig({ ...defaultRunnerConfig, tracingEnabled: true })); + await ec2Operations.create(createRunnerConfig({ ...defaultRunnerConfig, tracingEnabled: true })); expect(mockEC2Client).toHaveReceivedCommandWith(CreateFleetCommand, { ...expectedCreateFleetRequest({ @@ -516,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({ @@ -527,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({ @@ -538,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', @@ -577,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', @@ -616,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', @@ -667,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', @@ -705,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', @@ -775,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, @@ -790,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, @@ -800,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, @@ -816,7 +879,7 @@ describe('create runner with errors', () => { createFleetMockWithErrors(Array(12).fill('InsufficientFreeAddressesInSubnet')); await expect( - createRunner({ + ec2Operations.create({ ...createRunnerConfig({ ...defaultRunnerConfig, scaleErrors: [...defaultRunnerConfig.scaleErrors, 'InsufficientFreeAddressesInSubnet'], @@ -833,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, @@ -848,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, @@ -862,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, @@ -878,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, @@ -890,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, @@ -906,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, @@ -919,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', @@ -940,7 +1003,7 @@ describe('create runner with errors', () => { ); await expect( - createRunner( + ec2Operations.create( createRunnerConfig({ ...defaultRunnerConfig, amiIdSsmParameterName: 'my-ami-id-param', @@ -955,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', @@ -968,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, @@ -984,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, @@ -1022,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, @@ -1053,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: [], @@ -1066,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, }); @@ -1104,7 +1167,7 @@ describe('create runner with errors fail over to OnDemand', () => { createFleetMockWithWithOnDemandFallback(['UnfulfillableCapacity'], instancesIds); await expect( - createRunner({ + ec2Operations.create({ ...createRunnerConfig(defaultRunnerConfig), numberOfRunners: 2, }), @@ -1336,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'], @@ -1350,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, @@ -1365,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, @@ -1378,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: { @@ -1417,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', @@ -1462,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', @@ -1475,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, @@ -1485,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, @@ -1496,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, @@ -1507,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, @@ -1519,7 +1582,7 @@ describe('create runner with useDedicatedHost', () => { mockEC2Client.on(RunInstancesCommand).rejects(error); await expect( - createRunner({ + ec2Operations.create({ ...createRunnerConfig({ ...dedicatedHostRunnerConfig, scaleErrors: ['InsufficientInstanceCapacity'], @@ -1535,7 +1598,7 @@ describe('create runner with useDedicatedHost', () => { }); await expect( - createRunner({ + ec2Operations.create({ ...createRunnerConfig(dedicatedHostRunnerConfig), numberOfRunners: 2, }), @@ -1554,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', @@ -1639,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 4e187437ec..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'); @@ -32,13 +32,63 @@ interface Ec2Filter { Values: string[]; } +export interface Ec2RunnerRequestContext { + readonly signal: AbortSignal | undefined; +} + +export interface Ec2RunnerResourceOperations { + 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 Ec2RunnerProvisioningOperations extends Ec2RunnerResourceOperations { + getDefaultBlockDeviceNameFromLaunchTemplate(launchTemplateName: string): Promise; +} + +export interface Ec2RunnerClient { + forRequest(context: Ec2RunnerRequestContext): Ec2RunnerProvisioningOperations; +} + +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, () => 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), + ), + }), + }; +} + type FleetError = NonNullable[number]; -export async function listEC2Runners(filters: Ec2ListRunnerFilters | undefined = undefined): Promise { +async function listEc2Runners( + 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 +118,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; @@ -108,23 +162,34 @@ function getRunnerInfo(runningInstances: DescribeInstancesResult) { return runners; } -export async function terminateRunner(instanceId: string): Promise { +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 { +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 { +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 = [ @@ -298,7 +363,11 @@ function buildRunInstancesOverrides( return overrides; } -export async function createRunner(runnerParameters: RunnerInputParameters): Promise { +async function createEc2Runner( + ec2Client: EC2Client, + runnerParameters: RunnerInputParameters, + signal: AbortSignal | undefined, +): Promise { logger.debug('Runner configuration.', { runner: { configuration: { @@ -307,11 +376,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 +394,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 +419,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 +448,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( + ec2Client, + { + ...runnerParameters, + numberOfRunners: numberOfInstances, + onDemandFailoverOnError: ['InsufficientInstanceCapacity'], + ec2instanceCriteria: { + ...runnerParameters.ec2instanceCriteria, + targetCapacityType: 'on-demand', + instanceAllocationStrategy: failoverAllocationStrategy, + }, }, - }); + signal, + ); instances.push(...onDemandResult.instances); return { instances, @@ -474,6 +550,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 +578,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 +651,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 +663,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 +710,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); 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, }, }; }