feat(lambdas): add batch SSM parameter fetching to reduce API calls

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

@@ -7,19 +7,17 @@ import {
   EC2Client,
   Image,
 } from '@aws-sdk/client-ec2';
-import {
-  DescribeParametersCommand,
-  DescribeParametersCommandOutput,
-  GetParameterCommand,
-  SSMClient,
-} from '@aws-sdk/client-ssm';
+import { DescribeParametersCommand, DescribeParametersCommandOutput, SSMClient } from '@aws-sdk/client-ssm';
+import { getParameters } from '@aws-github-runner/aws-ssm-util';
 import { mockClient } from 'aws-sdk-client-mock';
 import 'aws-sdk-client-mock-jest/vitest';
 
 import { AmiCleanupOptions, amiCleanup, defaultAmiCleanupOptions } from './ami';
 import { describe, it, expect, beforeEach, vi } from 'vitest';
 import { fail } from 'assert';
 
+vi.mock('@aws-github-runner/aws-ssm-util');
+
 process.env.AWS_REGION = 'eu-east-1';
 const deleteAmisOlderThenDays = 30;
 const date31DaysAgo = new Date(new Date().setDate(new Date().getDate() - (deleteAmisOlderThenDays + 1)));
@@ -83,22 +81,12 @@ describe("delete AMI's", () => {
     mockSSMClient.reset();
 
     mockSSMClient.on(DescribeParametersCommand).resolves(ssmParameters);
-    mockSSMClient.on(GetParameterCommand, { Name: 'ami-id/ami-ssm0001' }).resolves({
-      Parameter: {
-        Name: 'ami-id/ami-ssm0001',
-        Type: 'String',
-        Value: 'ami-ssm0001',
-        Version: 1,
-      },
-    });
-    mockSSMClient.on(GetParameterCommand, { Name: 'ami-id/ami-ssm0002' }).resolves({
-      Parameter: {
-        Name: 'ami-id/ami-ssm0002',
-        Type: 'String',
-        Value: 'ami-ssm0002',
-        Version: 1,
-      },
-    });
+    vi.mocked(getParameters).mockResolvedValue(
+      new Map([
+        ['ami-id/ami-ssm0001', 'ami-ssm0001'],
+        ['ami-id/ami-ssm0002', 'ami-ssm0002'],
+      ]),
+    );
 
     mockEC2Client.on(DescribeLaunchTemplatesCommand).resolves({
       LaunchTemplates: [
@@ -143,13 +131,7 @@ describe("delete AMI's", () => {
     expect(mockEC2Client).toHaveReceivedCommand(DescribeLaunchTemplatesCommand);
     expect(mockEC2Client).toHaveReceivedCommand(DescribeLaunchTemplateVersionsCommand);
     expect(mockSSMClient).toHaveReceivedCommand(DescribeParametersCommand);
-    expect(mockSSMClient).toHaveReceivedCommandTimes(GetParameterCommand, 2);
-    expect(mockSSMClient).toHaveReceivedCommandWith(GetParameterCommand, {
-      Name: 'ami-id/ami-ssm0001',
-    });
-    expect(mockSSMClient).toHaveReceivedCommandWith(GetParameterCommand, {
-      Name: 'ami-id/ami-ssm0002',
-    });
+    expect(getParameters).toHaveBeenCalledWith(['ami-id/ami-ssm0001', 'ami-id/ami-ssm0002']);
   });
 
   it('should NOT delete instances in use.', async () => {
@@ -485,14 +467,7 @@ describe("delete AMI's", () => {
         ],
       });
 
-      mockSSMClient.on(GetParameterCommand, { Name: '/github-runner/config/ami_id' }).resolves({
-        Parameter: {
-          Name: '/github-runner/config/ami_id',
-          Type: 'String',
-          Value: 'ami-underscore0001',
-          Version: 1,
-        },
-      });
+      vi.mocked(getParameters).mockResolvedValue(new Map([['/github-runner/config/ami_id', 'ami-underscore0001']]));
 
       await amiCleanup({
         minimumDaysOld: 0,
@@ -501,9 +476,7 @@ describe("delete AMI's", () => {
 
       // AMI should not be deleted because it's referenced in SSM
       expect(mockEC2Client).not.toHaveReceivedCommand(DeregisterImageCommand);
-      expect(mockSSMClient).toHaveReceivedCommandWith(GetParameterCommand, {
-        Name: '/github-runner/config/ami_id',
-      });
+      expect(getParameters).toHaveBeenCalledWith(['/github-runner/config/ami_id']);
       expect(mockSSMClient).not.toHaveReceivedCommand(DescribeParametersCommand);
     });
 
@@ -518,14 +491,7 @@ describe("delete AMI's", () => {
         ],
       });
 
-      mockSSMClient.on(GetParameterCommand, { Name: '/github-runner/config/ami-id' }).resolves({
-        Parameter: {
-          Name: '/github-runner/config/ami-id',
-          Type: 'String',
-          Value: 'ami-hyphen0001',
-          Version: 1,
-        },
-      });
+      vi.mocked(getParameters).mockResolvedValue(new Map([['/github-runner/config/ami-id', 'ami-hyphen0001']]));
 
       await amiCleanup({
         minimumDaysOld: 0,
@@ -534,9 +500,7 @@ describe("delete AMI's", () => {
 
       // AMI should not be deleted because it's referenced in SSM
       expect(mockEC2Client).not.toHaveReceivedCommand(DeregisterImageCommand);
-      expect(mockSSMClient).toHaveReceivedCommandWith(GetParameterCommand, {
-        Name: '/github-runner/config/ami-id',
-      });
+      expect(getParameters).toHaveBeenCalledWith(['/github-runner/config/ami-id']);
       expect(mockSSMClient).not.toHaveReceivedCommand(DescribeParametersCommand);
     });
 
@@ -561,14 +525,7 @@ describe("delete AMI's", () => {
         ],
       });
 
-      mockSSMClient.on(GetParameterCommand, { Name: '/some/path/ami-id' }).resolves({
-        Parameter: {
-          Name: '/some/path/ami-id',
-          Type: 'String',
-          Value: 'ami-wildcard0001',
-          Version: 1,
-        },
-      });
+      vi.mocked(getParameters).mockResolvedValue(new Map([['/some/path/ami-id', 'ami-wildcard0001']]));
 
       await amiCleanup({
         minimumDaysOld: 0,
@@ -580,9 +537,7 @@ describe("delete AMI's", () => {
       expect(mockSSMClient).toHaveReceivedCommandWith(DescribeParametersCommand, {
         ParameterFilters: [{ Key: 'Name', Option: 'Contains', Values: ['ami-id'] }],
       });
-      expect(mockSSMClient).toHaveReceivedCommandWith(GetParameterCommand, {
-        Name: '/some/path/ami-id',
-      });
+      expect(getParameters).toHaveBeenCalledWith(['/some/path/ami-id']);
     });
 
     it('handles wildcard SSM parameter patterns (*ami_id)', async () => {
@@ -606,14 +561,9 @@ describe("delete AMI's", () => {
         ],
       });
 
-      mockSSMClient.on(GetParameterCommand, { Name: '/github-runner/config/ami_id' }).resolves({
-        Parameter: {
-          Name: '/github-runner/config/ami_id',
-          Type: 'String',
-          Value: 'ami-wildcard-underscore0001',
-          Version: 1,
-        },
-      });
+      vi.mocked(getParameters).mockResolvedValue(
+        new Map([['/github-runner/config/ami_id', 'ami-wildcard-underscore0001']]),
+      );
 
       await amiCleanup({
         minimumDaysOld: 0,
@@ -625,9 +575,7 @@ describe("delete AMI's", () => {
       expect(mockSSMClient).toHaveReceivedCommandWith(DescribeParametersCommand, {
         ParameterFilters: [{ Key: 'Name', Option: 'Contains', Values: ['ami_id'] }],
       });
-      expect(mockSSMClient).toHaveReceivedCommandWith(GetParameterCommand, {
-        Name: '/github-runner/config/ami_id',
-      });
+      expect(getParameters).toHaveBeenCalledWith(['/github-runner/config/ami_id']);
     });
 
     it('handles mixed explicit names and wildcard patterns', async () => {
@@ -649,14 +597,9 @@ describe("delete AMI's", () => {
         ],
       });
 
-      mockSSMClient.on(GetParameterCommand, { Name: '/explicit/ami_id' }).resolves({
-        Parameter: {
-          Name: '/explicit/ami_id',
-          Type: 'String',
-          Value: 'ami-explicit0001',
-          Version: 1,
-        },
-      });
+      vi.mocked(getParameters)
+        .mockResolvedValueOnce(new Map([['/explicit/ami_id', 'ami-explicit0001']]))
+        .mockResolvedValueOnce(new Map([['/discovered/ami-id', 'ami-wildcard0001']]));
 
       mockSSMClient.on(DescribeParametersCommand).resolves({
         Parameters: [
@@ -668,15 +611,6 @@ describe("delete AMI's", () => {
         ],
       });
 
-      mockSSMClient.on(GetParameterCommand, { Name: '/discovered/ami-id' }).resolves({
-        Parameter: {
-          Name: '/discovered/ami-id',
-          Type: 'String',
-          Value: 'ami-wildcard0001',
-          Version: 1,
-        },
-      });
-
       await amiCleanup({
         minimumDaysOld: 0,
         ssmParameterNames: ['/explicit/ami_id', '*ami-id'],
@@ -688,15 +622,11 @@ describe("delete AMI's", () => {
         ImageId: 'ami-unused0001',
       });
 
-      expect(mockSSMClient).toHaveReceivedCommandWith(GetParameterCommand, {
-        Name: '/explicit/ami_id',
-      });
+      expect(getParameters).toHaveBeenCalledWith(['/explicit/ami_id']);
       expect(mockSSMClient).toHaveReceivedCommandWith(DescribeParametersCommand, {
         ParameterFilters: [{ Key: 'Name', Option: 'Contains', Values: ['ami-id'] }],
       });
-      expect(mockSSMClient).toHaveReceivedCommandWith(GetParameterCommand, {
-        Name: '/discovered/ami-id',
-      });
+      expect(getParameters).toHaveBeenCalledWith(['/discovered/ami-id']);
     });
 
     it('handles SSM parameter fetch failures gracefully', async () => {
@@ -710,7 +640,7 @@ describe("delete AMI's", () => {
         ],
       });
 
-      mockSSMClient.on(GetParameterCommand, { Name: '/nonexistent/ami_id' }).rejects(new Error('ParameterNotFound'));
+      vi.mocked(getParameters).mockRejectedValue(new Error('ParameterNotFound'));
 
       // Should not throw and should delete the AMI since SSM reference failed
       await amiCleanup({
@@ -768,7 +698,7 @@ describe("delete AMI's", () => {
         ImageId: 'ami-no-ssm0001',
       });
       expect(mockSSMClient).not.toHaveReceivedCommand(DescribeParametersCommand);
-      expect(mockSSMClient).not.toHaveReceivedCommand(GetParameterCommand);
+      expect(getParameters).not.toHaveBeenCalled();
     });
   });
 });

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

@@ -8,9 +8,10 @@ import {
   Filter,
   Image,
 } from '@aws-sdk/client-ec2';
-import { GetParameterCommand, SSMClient, DescribeParametersCommand } from '@aws-sdk/client-ssm';
+import { SSMClient, DescribeParametersCommand } from '@aws-sdk/client-ssm';
 import { createChildLogger } from '@aws-github-runner/aws-powertools-util';
 import { getTracedAWSV3Client } from '@aws-github-runner/aws-powertools-util';
+import { getParameters } from '@aws-github-runner/aws-ssm-util';
 
 const logger = createChildLogger('ami');
 
@@ -184,22 +185,34 @@ async function deleteSnapshot(options: AmiCleanupOptions, amiDetails: Image, ec2
 }
 
 /**
- * Resolves the value of an SSM parameter by its name. Doesn't fail on errors,
- * but warns instead, as this process is best-effort.
+ * Resolves the values of multiple SSM parameters by their names.
+ * Delegates batching to the shared `getParameters` utility.
+ * Doesn't fail on errors, but warns instead, as this process is best-effort.
  *
- * @param name - The SSM parameter name to resolve
- * @param ssmClient - Configured SSM client for making API calls
- * @returns The parameter value if successful, undefined if parameter doesn't exist or access fails
+ * @param names - Array of SSM parameter names to resolve
+ * @returns Array of parameter values in the same order as input (undefined for missing/failed parameters)
  */
-async function resolveSsmParameterValue(name: string, ssmClient: SSMClient): Promise<string | undefined> {
+async function resolveSsmParameterValues(names: string[]): Promise<(string | undefined)[]> {
+  if (names.length === 0) {
+    return [];
+  }
+
   try {
-    const { Parameter: parameter } = await ssmClient.send(new GetParameterCommand({ Name: name }));
+    const parameterMap = await getParameters(names);
 
-    return parameter?.Value;
-  } catch (error: unknown) {
-    logger.warn(`Failed to resolve image id from SSM parameter ${name}`, { error });
+    // Log warnings for parameters that couldn't be resolved
+    for (const name of names) {
+      if (!parameterMap.has(name)) {
+        logger.warn(`Failed to resolve image id from SSM parameter ${name}: Parameter not found or access denied`);
+      }
+    }
 
-    return undefined;
+    // Return values in the same order as input names
+    return names.map((name) => parameterMap.get(name));
+  } catch (error: unknown) {
+    logger.warn(`Failed to resolve image ids from SSM parameters ${names.join(', ')}`, { error });
+    // Mark all parameters as undefined on failure
+    return names.map(() => undefined);
   }
 }
 
@@ -273,11 +286,12 @@ async function getAmisReferedInSSM(options: AmiCleanupOptions): Promise<(string
   const explicitNames = options.ssmParameterNames.filter((n) => !n.startsWith('*'));
   const wildcardPatterns = options.ssmParameterNames.filter((n) => n.startsWith('*'));
 
-  const explicitValuesPromise = Promise.all(explicitNames.map((name) => resolveSsmParameterValue(name, ssmClient)));
+  // Batch fetch explicit parameter values in chunks of 10 (AWS API limit)
+  const explicitValuesPromise = resolveSsmParameterValues(explicitNames);
 
   // Handle wildcard patterns by first discovering matching parameters, then
   // fetching their values
-  let wildcardValues: Promise<(string | undefined)[]> = Promise.resolve([]);
+  let wildcardValuesPromise: Promise<(string | undefined)[]> = Promise.resolve([]);
   if (wildcardPatterns.length > 0) {
     // Convert wildcard patterns to SSM ParameterFilters using Contains logic
     // Example: "*ami-id" becomes a filter for parameters containing "ami-id"
@@ -287,24 +301,30 @@ async function getAmisReferedInSSM(options: AmiCleanupOptions): Promise<(string
       Values: [p.replace(/^\*/g, '')],
     }));
 
-    try {
-      // Discover parameters matching the wildcard patterns
-      logger.debug('Describing SSM parameter', { filters });
-      const ssmParameters = await ssmClient.send(new DescribeParametersCommand({ ParameterFilters: filters }));
-
-      // Fetch the actual values of discovered parameters
-      wildcardValues = Promise.all(
-        (ssmParameters.Parameters ?? []).map((param) => resolveSsmParameterValue(param.Name!, ssmClient)),
-      );
-    } catch (e) {
-      logger.warn('Failed to describe SSM parameters using wildcard patterns', { error: e });
-    }
+    wildcardValuesPromise = (async () => {
+      try {
+        // Discover parameters matching the wildcard patterns
+        logger.debug('Describing SSM parameter', { filters });
+        const ssmParameters = await ssmClient.send(new DescribeParametersCommand({ ParameterFilters: filters }));
+
+        // Batch fetch the actual values of discovered parameters
+        const discoveredNames = (ssmParameters.Parameters ?? [])
+          .map((param) => param.Name)
+          .filter((name): name is string => name !== undefined);
+
+        return resolveSsmParameterValues(discoveredNames);
+      } catch (e) {
+        logger.warn('Failed to describe SSM parameters using wildcard patterns', { error: e });
+        return [];
+      }
+    })();
   }
 
   // Combine results from both explicit and wildcard parameter resolution
-  const values = await Promise.all([explicitValuesPromise, wildcardValues]);
+  const [explicitValues, wildcardValues] = await Promise.all([explicitValuesPromise, wildcardValuesPromise]);
+  const values = [...explicitValues, ...wildcardValues];
   logger.debug('Resolved SSM parameter values', { values });
-  return values.flat();
+  return values;
 }
 
 export { amiCleanup, getAmisNotInUse };