diff --git a/lambdas/functions/control-plane/src/github/rate-limit.test.ts b/lambdas/functions/control-plane/src/github/rate-limit.test.ts index 591d4a14af..9bee34c94b 100644 --- a/lambdas/functions/control-plane/src/github/rate-limit.test.ts +++ b/lambdas/functions/control-plane/src/github/rate-limit.test.ts @@ -3,6 +3,7 @@ import { createSingleMetric } from '@aws-github-runner/aws-powertools-util'; import { MetricUnit } from '@aws-lambda-powertools/metrics'; import { metricGitHubAppRateLimit } from './rate-limit'; import { describe, it, expect, beforeEach, vi } from 'vitest'; +import { getParameter } from '@aws-github-runner/aws-ssm-util'; process.env.PARAMETER_GITHUB_APP_ID_NAME = 'test'; vi.mock('@aws-github-runner/aws-ssm-util', async () => { @@ -78,4 +79,27 @@ describe('metricGitHubAppRateLimit', () => { expect(createSingleMetric).not.toHaveBeenCalled(); }); + + it('should cache GitHub App ID and only call getParameter once', async () => { + // Reset modules to clear the appIdPromise cache + vi.resetModules(); + const { metricGitHubAppRateLimit: freshMetricFunction } = await import('./rate-limit'); + + process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT = 'true'; + const headers: ResponseHeaders = { + 'x-ratelimit-remaining': '10', + 'x-ratelimit-limit': '60', + }; + + const mockGetParameter = vi.mocked(getParameter); + mockGetParameter.mockClear(); + + await freshMetricFunction(headers); + await freshMetricFunction(headers); + await freshMetricFunction(headers); + + // getParameter should only be called once due to caching + expect(mockGetParameter).toHaveBeenCalledTimes(1); + expect(mockGetParameter).toHaveBeenCalledWith(process.env.PARAMETER_GITHUB_APP_ID_NAME); + }); }); diff --git a/lambdas/functions/control-plane/src/github/rate-limit.ts b/lambdas/functions/control-plane/src/github/rate-limit.ts index 6d2fd5fbe7..8ebb8e3f84 100644 --- a/lambdas/functions/control-plane/src/github/rate-limit.ts +++ b/lambdas/functions/control-plane/src/github/rate-limit.ts @@ -4,6 +4,16 @@ import { MetricUnit } from '@aws-lambda-powertools/metrics'; import yn from 'yn'; import { getParameter } from '@aws-github-runner/aws-ssm-util'; +// Cache the app ID to avoid repeated SSM calls across Lambda invocations +let appIdPromise: Promise | null = null; + +async function getAppId(): Promise { + if (!appIdPromise) { + appIdPromise = getParameter(process.env.PARAMETER_GITHUB_APP_ID_NAME); + } + return appIdPromise; +} + export async function metricGitHubAppRateLimit(headers: ResponseHeaders): Promise { try { const remaining = parseInt(headers['x-ratelimit-remaining'] as string); @@ -13,7 +23,7 @@ export async function metricGitHubAppRateLimit(headers: ResponseHeaders): Promis const updateMetric = yn(process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT); if (updateMetric) { - const appId = await getParameter(process.env.PARAMETER_GITHUB_APP_ID_NAME); + const appId = await getAppId(); const metric = createSingleMetric('GitHubAppRateLimitRemaining', MetricUnit.Count, remaining, { AppId: appId, });