Skip to content

Commit 24857c2

Browse files
thomasnemerBrend-SmitsCopilot
authored
feat(lambdas): add batch SSM parameter fetching to reduce API calls (#5017)
This PR intends to reduce SSM AWS API calls by doing the following: Add `getParameters()` function to aws-ssm-util that fetches multiple SSM parameters in a single API call with automatic chunking (max 10 per call per AWS API limits). Apply batch fetching to: - auth.ts: fetch App ID and Private Key in one call (2 calls → 1) - ConfigLoader.ts: fetch multiple matcher config paths in one call - ami.ts: batch resolve SSM parameter values for AMI lookups Also remove redundant appId SSM fetch in scale-up.ts that was only used for logging. --------- Co-authored-by: Brend Smits <brend.smits@philips.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
1 parent 1f9805d commit 24857c2

File tree

14 files changed

+378
-172
lines changed

14 files changed

+378
-172
lines changed

lambdas/functions/ami-housekeeper/src/ami.test.ts

Lines changed: 28 additions & 98 deletions
Original file line numberDiff line numberDiff line change
@@ -7,19 +7,17 @@ import {
77
EC2Client,
88
Image,
99
} from '@aws-sdk/client-ec2';
10-
import {
11-
DescribeParametersCommand,
12-
DescribeParametersCommandOutput,
13-
GetParameterCommand,
14-
SSMClient,
15-
} from '@aws-sdk/client-ssm';
10+
import { DescribeParametersCommand, DescribeParametersCommandOutput, SSMClient } from '@aws-sdk/client-ssm';
11+
import { getParameters } from '@aws-github-runner/aws-ssm-util';
1612
import { mockClient } from 'aws-sdk-client-mock';
1713
import 'aws-sdk-client-mock-jest/vitest';
1814

1915
import { AmiCleanupOptions, amiCleanup, defaultAmiCleanupOptions } from './ami';
2016
import { describe, it, expect, beforeEach, vi } from 'vitest';
2117
import { fail } from 'assert';
2218

19+
vi.mock('@aws-github-runner/aws-ssm-util');
20+
2321
process.env.AWS_REGION = 'eu-east-1';
2422
const deleteAmisOlderThenDays = 30;
2523
const date31DaysAgo = new Date(new Date().setDate(new Date().getDate() - (deleteAmisOlderThenDays + 1)));
@@ -83,22 +81,12 @@ describe("delete AMI's", () => {
8381
mockSSMClient.reset();
8482

8583
mockSSMClient.on(DescribeParametersCommand).resolves(ssmParameters);
86-
mockSSMClient.on(GetParameterCommand, { Name: 'ami-id/ami-ssm0001' }).resolves({
87-
Parameter: {
88-
Name: 'ami-id/ami-ssm0001',
89-
Type: 'String',
90-
Value: 'ami-ssm0001',
91-
Version: 1,
92-
},
93-
});
94-
mockSSMClient.on(GetParameterCommand, { Name: 'ami-id/ami-ssm0002' }).resolves({
95-
Parameter: {
96-
Name: 'ami-id/ami-ssm0002',
97-
Type: 'String',
98-
Value: 'ami-ssm0002',
99-
Version: 1,
100-
},
101-
});
84+
vi.mocked(getParameters).mockResolvedValue(
85+
new Map([
86+
['ami-id/ami-ssm0001', 'ami-ssm0001'],
87+
['ami-id/ami-ssm0002', 'ami-ssm0002'],
88+
]),
89+
);
10290

10391
mockEC2Client.on(DescribeLaunchTemplatesCommand).resolves({
10492
LaunchTemplates: [
@@ -143,13 +131,7 @@ describe("delete AMI's", () => {
143131
expect(mockEC2Client).toHaveReceivedCommand(DescribeLaunchTemplatesCommand);
144132
expect(mockEC2Client).toHaveReceivedCommand(DescribeLaunchTemplateVersionsCommand);
145133
expect(mockSSMClient).toHaveReceivedCommand(DescribeParametersCommand);
146-
expect(mockSSMClient).toHaveReceivedCommandTimes(GetParameterCommand, 2);
147-
expect(mockSSMClient).toHaveReceivedCommandWith(GetParameterCommand, {
148-
Name: 'ami-id/ami-ssm0001',
149-
});
150-
expect(mockSSMClient).toHaveReceivedCommandWith(GetParameterCommand, {
151-
Name: 'ami-id/ami-ssm0002',
152-
});
134+
expect(getParameters).toHaveBeenCalledWith(['ami-id/ami-ssm0001', 'ami-id/ami-ssm0002']);
153135
});
154136

155137
it('should NOT delete instances in use.', async () => {
@@ -485,14 +467,7 @@ describe("delete AMI's", () => {
485467
],
486468
});
487469

488-
mockSSMClient.on(GetParameterCommand, { Name: '/github-runner/config/ami_id' }).resolves({
489-
Parameter: {
490-
Name: '/github-runner/config/ami_id',
491-
Type: 'String',
492-
Value: 'ami-underscore0001',
493-
Version: 1,
494-
},
495-
});
470+
vi.mocked(getParameters).mockResolvedValue(new Map([['/github-runner/config/ami_id', 'ami-underscore0001']]));
496471

497472
await amiCleanup({
498473
minimumDaysOld: 0,
@@ -501,9 +476,7 @@ describe("delete AMI's", () => {
501476

502477
// AMI should not be deleted because it's referenced in SSM
503478
expect(mockEC2Client).not.toHaveReceivedCommand(DeregisterImageCommand);
504-
expect(mockSSMClient).toHaveReceivedCommandWith(GetParameterCommand, {
505-
Name: '/github-runner/config/ami_id',
506-
});
479+
expect(getParameters).toHaveBeenCalledWith(['/github-runner/config/ami_id']);
507480
expect(mockSSMClient).not.toHaveReceivedCommand(DescribeParametersCommand);
508481
});
509482

@@ -518,14 +491,7 @@ describe("delete AMI's", () => {
518491
],
519492
});
520493

521-
mockSSMClient.on(GetParameterCommand, { Name: '/github-runner/config/ami-id' }).resolves({
522-
Parameter: {
523-
Name: '/github-runner/config/ami-id',
524-
Type: 'String',
525-
Value: 'ami-hyphen0001',
526-
Version: 1,
527-
},
528-
});
494+
vi.mocked(getParameters).mockResolvedValue(new Map([['/github-runner/config/ami-id', 'ami-hyphen0001']]));
529495

530496
await amiCleanup({
531497
minimumDaysOld: 0,
@@ -534,9 +500,7 @@ describe("delete AMI's", () => {
534500

535501
// AMI should not be deleted because it's referenced in SSM
536502
expect(mockEC2Client).not.toHaveReceivedCommand(DeregisterImageCommand);
537-
expect(mockSSMClient).toHaveReceivedCommandWith(GetParameterCommand, {
538-
Name: '/github-runner/config/ami-id',
539-
});
503+
expect(getParameters).toHaveBeenCalledWith(['/github-runner/config/ami-id']);
540504
expect(mockSSMClient).not.toHaveReceivedCommand(DescribeParametersCommand);
541505
});
542506

@@ -561,14 +525,7 @@ describe("delete AMI's", () => {
561525
],
562526
});
563527

564-
mockSSMClient.on(GetParameterCommand, { Name: '/some/path/ami-id' }).resolves({
565-
Parameter: {
566-
Name: '/some/path/ami-id',
567-
Type: 'String',
568-
Value: 'ami-wildcard0001',
569-
Version: 1,
570-
},
571-
});
528+
vi.mocked(getParameters).mockResolvedValue(new Map([['/some/path/ami-id', 'ami-wildcard0001']]));
572529

573530
await amiCleanup({
574531
minimumDaysOld: 0,
@@ -580,9 +537,7 @@ describe("delete AMI's", () => {
580537
expect(mockSSMClient).toHaveReceivedCommandWith(DescribeParametersCommand, {
581538
ParameterFilters: [{ Key: 'Name', Option: 'Contains', Values: ['ami-id'] }],
582539
});
583-
expect(mockSSMClient).toHaveReceivedCommandWith(GetParameterCommand, {
584-
Name: '/some/path/ami-id',
585-
});
540+
expect(getParameters).toHaveBeenCalledWith(['/some/path/ami-id']);
586541
});
587542

588543
it('handles wildcard SSM parameter patterns (*ami_id)', async () => {
@@ -606,14 +561,9 @@ describe("delete AMI's", () => {
606561
],
607562
});
608563

609-
mockSSMClient.on(GetParameterCommand, { Name: '/github-runner/config/ami_id' }).resolves({
610-
Parameter: {
611-
Name: '/github-runner/config/ami_id',
612-
Type: 'String',
613-
Value: 'ami-wildcard-underscore0001',
614-
Version: 1,
615-
},
616-
});
564+
vi.mocked(getParameters).mockResolvedValue(
565+
new Map([['/github-runner/config/ami_id', 'ami-wildcard-underscore0001']]),
566+
);
617567

618568
await amiCleanup({
619569
minimumDaysOld: 0,
@@ -625,9 +575,7 @@ describe("delete AMI's", () => {
625575
expect(mockSSMClient).toHaveReceivedCommandWith(DescribeParametersCommand, {
626576
ParameterFilters: [{ Key: 'Name', Option: 'Contains', Values: ['ami_id'] }],
627577
});
628-
expect(mockSSMClient).toHaveReceivedCommandWith(GetParameterCommand, {
629-
Name: '/github-runner/config/ami_id',
630-
});
578+
expect(getParameters).toHaveBeenCalledWith(['/github-runner/config/ami_id']);
631579
});
632580

633581
it('handles mixed explicit names and wildcard patterns', async () => {
@@ -649,14 +597,9 @@ describe("delete AMI's", () => {
649597
],
650598
});
651599

652-
mockSSMClient.on(GetParameterCommand, { Name: '/explicit/ami_id' }).resolves({
653-
Parameter: {
654-
Name: '/explicit/ami_id',
655-
Type: 'String',
656-
Value: 'ami-explicit0001',
657-
Version: 1,
658-
},
659-
});
600+
vi.mocked(getParameters)
601+
.mockResolvedValueOnce(new Map([['/explicit/ami_id', 'ami-explicit0001']]))
602+
.mockResolvedValueOnce(new Map([['/discovered/ami-id', 'ami-wildcard0001']]));
660603

661604
mockSSMClient.on(DescribeParametersCommand).resolves({
662605
Parameters: [
@@ -668,15 +611,6 @@ describe("delete AMI's", () => {
668611
],
669612
});
670613

671-
mockSSMClient.on(GetParameterCommand, { Name: '/discovered/ami-id' }).resolves({
672-
Parameter: {
673-
Name: '/discovered/ami-id',
674-
Type: 'String',
675-
Value: 'ami-wildcard0001',
676-
Version: 1,
677-
},
678-
});
679-
680614
await amiCleanup({
681615
minimumDaysOld: 0,
682616
ssmParameterNames: ['/explicit/ami_id', '*ami-id'],
@@ -688,15 +622,11 @@ describe("delete AMI's", () => {
688622
ImageId: 'ami-unused0001',
689623
});
690624

691-
expect(mockSSMClient).toHaveReceivedCommandWith(GetParameterCommand, {
692-
Name: '/explicit/ami_id',
693-
});
625+
expect(getParameters).toHaveBeenCalledWith(['/explicit/ami_id']);
694626
expect(mockSSMClient).toHaveReceivedCommandWith(DescribeParametersCommand, {
695627
ParameterFilters: [{ Key: 'Name', Option: 'Contains', Values: ['ami-id'] }],
696628
});
697-
expect(mockSSMClient).toHaveReceivedCommandWith(GetParameterCommand, {
698-
Name: '/discovered/ami-id',
699-
});
629+
expect(getParameters).toHaveBeenCalledWith(['/discovered/ami-id']);
700630
});
701631

702632
it('handles SSM parameter fetch failures gracefully', async () => {
@@ -710,7 +640,7 @@ describe("delete AMI's", () => {
710640
],
711641
});
712642

713-
mockSSMClient.on(GetParameterCommand, { Name: '/nonexistent/ami_id' }).rejects(new Error('ParameterNotFound'));
643+
vi.mocked(getParameters).mockRejectedValue(new Error('ParameterNotFound'));
714644

715645
// Should not throw and should delete the AMI since SSM reference failed
716646
await amiCleanup({
@@ -768,7 +698,7 @@ describe("delete AMI's", () => {
768698
ImageId: 'ami-no-ssm0001',
769699
});
770700
expect(mockSSMClient).not.toHaveReceivedCommand(DescribeParametersCommand);
771-
expect(mockSSMClient).not.toHaveReceivedCommand(GetParameterCommand);
701+
expect(getParameters).not.toHaveBeenCalled();
772702
});
773703
});
774704
});

lambdas/functions/ami-housekeeper/src/ami.ts

Lines changed: 48 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,10 @@ import {
88
Filter,
99
Image,
1010
} from '@aws-sdk/client-ec2';
11-
import { GetParameterCommand, SSMClient, DescribeParametersCommand } from '@aws-sdk/client-ssm';
11+
import { SSMClient, DescribeParametersCommand } from '@aws-sdk/client-ssm';
1212
import { createChildLogger } from '@aws-github-runner/aws-powertools-util';
1313
import { getTracedAWSV3Client } from '@aws-github-runner/aws-powertools-util';
14+
import { getParameters } from '@aws-github-runner/aws-ssm-util';
1415

1516
const logger = createChildLogger('ami');
1617

@@ -184,22 +185,34 @@ async function deleteSnapshot(options: AmiCleanupOptions, amiDetails: Image, ec2
184185
}
185186

186187
/**
187-
* Resolves the value of an SSM parameter by its name. Doesn't fail on errors,
188-
* but warns instead, as this process is best-effort.
188+
* Resolves the values of multiple SSM parameters by their names.
189+
* Delegates batching to the shared `getParameters` utility.
190+
* Doesn't fail on errors, but warns instead, as this process is best-effort.
189191
*
190-
* @param name - The SSM parameter name to resolve
191-
* @param ssmClient - Configured SSM client for making API calls
192-
* @returns The parameter value if successful, undefined if parameter doesn't exist or access fails
192+
* @param names - Array of SSM parameter names to resolve
193+
* @returns Array of parameter values in the same order as input (undefined for missing/failed parameters)
193194
*/
194-
async function resolveSsmParameterValue(name: string, ssmClient: SSMClient): Promise<string | undefined> {
195+
async function resolveSsmParameterValues(names: string[]): Promise<(string | undefined)[]> {
196+
if (names.length === 0) {
197+
return [];
198+
}
199+
195200
try {
196-
const { Parameter: parameter } = await ssmClient.send(new GetParameterCommand({ Name: name }));
201+
const parameterMap = await getParameters(names);
197202

198-
return parameter?.Value;
199-
} catch (error: unknown) {
200-
logger.warn(`Failed to resolve image id from SSM parameter ${name}`, { error });
203+
// Log warnings for parameters that couldn't be resolved
204+
for (const name of names) {
205+
if (!parameterMap.has(name)) {
206+
logger.warn(`Failed to resolve image id from SSM parameter ${name}: Parameter not found or access denied`);
207+
}
208+
}
201209

202-
return undefined;
210+
// Return values in the same order as input names
211+
return names.map((name) => parameterMap.get(name));
212+
} catch (error: unknown) {
213+
logger.warn(`Failed to resolve image ids from SSM parameters ${names.join(', ')}`, { error });
214+
// Mark all parameters as undefined on failure
215+
return names.map(() => undefined);
203216
}
204217
}
205218

@@ -273,11 +286,12 @@ async function getAmisReferedInSSM(options: AmiCleanupOptions): Promise<(string
273286
const explicitNames = options.ssmParameterNames.filter((n) => !n.startsWith('*'));
274287
const wildcardPatterns = options.ssmParameterNames.filter((n) => n.startsWith('*'));
275288

276-
const explicitValuesPromise = Promise.all(explicitNames.map((name) => resolveSsmParameterValue(name, ssmClient)));
289+
// Batch fetch explicit parameter values in chunks of 10 (AWS API limit)
290+
const explicitValuesPromise = resolveSsmParameterValues(explicitNames);
277291

278292
// Handle wildcard patterns by first discovering matching parameters, then
279293
// fetching their values
280-
let wildcardValues: Promise<(string | undefined)[]> = Promise.resolve([]);
294+
let wildcardValuesPromise: Promise<(string | undefined)[]> = Promise.resolve([]);
281295
if (wildcardPatterns.length > 0) {
282296
// Convert wildcard patterns to SSM ParameterFilters using Contains logic
283297
// Example: "*ami-id" becomes a filter for parameters containing "ami-id"
@@ -287,24 +301,30 @@ async function getAmisReferedInSSM(options: AmiCleanupOptions): Promise<(string
287301
Values: [p.replace(/^\*/g, '')],
288302
}));
289303

290-
try {
291-
// Discover parameters matching the wildcard patterns
292-
logger.debug('Describing SSM parameter', { filters });
293-
const ssmParameters = await ssmClient.send(new DescribeParametersCommand({ ParameterFilters: filters }));
294-
295-
// Fetch the actual values of discovered parameters
296-
wildcardValues = Promise.all(
297-
(ssmParameters.Parameters ?? []).map((param) => resolveSsmParameterValue(param.Name!, ssmClient)),
298-
);
299-
} catch (e) {
300-
logger.warn('Failed to describe SSM parameters using wildcard patterns', { error: e });
301-
}
304+
wildcardValuesPromise = (async () => {
305+
try {
306+
// Discover parameters matching the wildcard patterns
307+
logger.debug('Describing SSM parameter', { filters });
308+
const ssmParameters = await ssmClient.send(new DescribeParametersCommand({ ParameterFilters: filters }));
309+
310+
// Batch fetch the actual values of discovered parameters
311+
const discoveredNames = (ssmParameters.Parameters ?? [])
312+
.map((param) => param.Name)
313+
.filter((name): name is string => name !== undefined);
314+
315+
return resolveSsmParameterValues(discoveredNames);
316+
} catch (e) {
317+
logger.warn('Failed to describe SSM parameters using wildcard patterns', { error: e });
318+
return [];
319+
}
320+
})();
302321
}
303322

304323
// Combine results from both explicit and wildcard parameter resolution
305-
const values = await Promise.all([explicitValuesPromise, wildcardValues]);
324+
const [explicitValues, wildcardValues] = await Promise.all([explicitValuesPromise, wildcardValuesPromise]);
325+
const values = [...explicitValues, ...wildcardValues];
306326
logger.debug('Resolved SSM parameter values', { values });
307-
return values.flat();
327+
return values;
308328
}
309329

310330
export { amiCleanup, getAmisNotInUse };

0 commit comments

Comments
 (0)