fix(scale-down): reselect GitHub Apps and attribute cleanup quota

lambdas/functions/control-plane/src/scale-runners/scale-down.test.ts

@@ -6,6 +6,7 @@ import { defaultComputeProvider } from '@aws-github-runner/compute-providers/pro
 
 import { controlPlaneProviderRegistry } from '../control-plane-providers';
 import * as ghAuth from '../github/auth';
+import * as rateLimit from '../github/rate-limit';
 import { githubCache } from './cache';
 import { newestFirstStrategy, oldestFirstStrategy, scaleDown } from './scale-down';
 import type { RunnerInfo, RunnerType, ScaleDownComputeProvider } from './types';
@@ -18,6 +19,7 @@ vi.mock('../github/auth', () => ({
 }));
 
 const mockOctokit = {
+  hook: { after: vi.fn(), error: vi.fn() },
   apps: {
     getOrgInstallation: vi.fn(),
     getRepoInstallation: vi.fn(),
@@ -262,6 +264,70 @@ describe('Scale down runners', () => {
     mockCreateClient.mockResolvedValue(mockOctokit as unknown as Octokit);
   });
 
+  it('selects Apps before owner-cache lookup and attributes response quota to the selected App', async () => {
+    process.env.SCALE_DOWN_CONFIG = '[]';
+    const runners = ['first', 'second', 'third'].map((id) => createRunnerTestData(id, 'Org', 60, true, false, false));
+    mockGitHubRunners(runners);
+    mockListRunners.mockResolvedValueOnce([]).mockResolvedValueOnce(runners).mockResolvedValue([]);
+    const authentication = { type: 'app' as const, token: 'token', appId: 1 };
+    mockedAppAuth
+      .mockResolvedValueOnce({ ...authentication, appIndex: 0 })
+      .mockResolvedValue({ ...authentication, appIndex: 1 });
+    vi.mocked(ghAuth.getStoredInstallationId).mockResolvedValue(123);
+    const metric = vi.spyOn(rateLimit, 'metricGitHubAppRateLimit').mockResolvedValue();
+    try {
+      await scaleDown();
+      expect(mockedInstallationAuth).toHaveBeenCalledTimes(2);
+      expect(mockedInstallationAuth).toHaveBeenCalledWith(123, '', 0);
+      expect(mockedInstallationAuth).toHaveBeenCalledWith(123, '', 1);
+      expect(githubCache.clients.has(`0:Org:${runners[0].owner}`)).toBe(true);
+      expect(githubCache.clients.has(`1:Org:${runners[0].owner}`)).toBe(true);
+      const headers = { 'x-ratelimit-remaining': '17', 'x-ratelimit-limit': '5000' };
+      await mockOctokit.hook.after.mock.calls[1][1]({ headers });
+      expect(metric).toHaveBeenCalledWith(headers, 1);
+    } finally {
+      metric.mockRestore();
+      vi.mocked(ghAuth.getStoredInstallationId).mockResolvedValue(undefined);
+    }
+  });
+
+  it('attributes successful and failed installation lookup and runner requests to the selected App', async () => {
+    const runner = createRunnerTestData('quota-hooks', 'Org', 60, true, false, false);
+    mockGitHubRunners([runner]);
+    mockListRunners.mockResolvedValueOnce([]).mockResolvedValueOnce([runner]).mockResolvedValue([]);
+    mockedAppAuth.mockResolvedValue({ type: 'app', token: 'token', appId: 1, appIndex: 2 });
+    vi.mocked(ghAuth.getStoredInstallationId).mockResolvedValue(undefined);
+    const metric = vi.spyOn(rateLimit, 'metricGitHubAppRateLimit').mockResolvedValue();
+    try {
+      mockOctokit.apps.getOrgInstallation.mockImplementationOnce(async () => {
+        expect(mockOctokit.hook.after).toHaveBeenCalledTimes(1);
+        expect(mockOctokit.hook.error).toHaveBeenCalledTimes(1);
+        return { data: { id: 123 } };
+      });
+      await scaleDown();
+      expect(mockOctokit.hook.after).toHaveBeenCalledTimes(2);
+      const headers = { 'x-ratelimit-remaining': '0', 'x-ratelimit-limit': '5000' };
+      for (const [, hook] of mockOctokit.hook.after.mock.calls) {
+        await hook({ headers });
+        expect(metric).toHaveBeenLastCalledWith(headers, 2);
+      }
+      for (const [, hook] of mockOctokit.hook.error.mock.calls) {
+        const error = new RequestError('rate limited', 403, {
+          request: { method: 'GET', url: 'https://api.github.com/test', headers: {} },
+          response: { status: 403, url: 'https://api.github.com/test', headers, data: {} },
+        });
+        await expect(hook(error)).rejects.toBe(error);
+        expect(metric).toHaveBeenLastCalledWith(headers, 2);
+        const networkError = new Error('network unavailable');
+        metric.mockClear();
+        await expect(hook(networkError)).rejects.toBe(networkError);
+        expect(metric).not.toHaveBeenCalled();
+      }
+    } finally {
+      metric.mockRestore();
+    }
+  });
+
   const endpoints = ['https://api.github.com', 'https://github.enterprise.something', 'https://companyname.ghe.com'];
 
   describe.each(endpoints)('for %s', (endpoint) => {

lambdas/functions/control-plane/src/scale-runners/scale-down.ts

@@ -24,24 +24,33 @@ type OrgRunnerList = Endpoints['GET /orgs/{org}/actions/runners']['response']['d
 type RepoRunnerList = Endpoints['GET /repos/{owner}/{repo}/actions/runners']['response']['data']['runners'];
 type RunnerState = OrgRunnerList[number] | RepoRunnerList[number];
 
-async function getOrCreateOctokit(runner: RunnerInfo): Promise<Octokit> {
-  const key = runner.owner;
-  const cachedOctokit = githubCache.clients.get(key);
-
-  if (cachedOctokit) {
-    logger.debug(`[createGitHubClientForRunner] Cache hit for ${key}`);
-    return cachedOctokit;
-  }
+function trackAppQuota(client: Octokit, appIdx: number | undefined): void {
+  client.hook.after('request', async (response) => {
+    await metricGitHubAppRateLimit(response.headers, appIdx);
+  });
+  client.hook.error('request', async (error) => {
+    if (error instanceof RequestError && error.response?.headers) {
+      await metricGitHubAppRateLimit(error.response.headers, appIdx);
+    }
+    throw error;
+  });
+}
 
-  logger.debug(`[createGitHubClientForRunner] Cache miss for ${key}`);
+async function getOrCreateOctokit(runner: RunnerInfo): Promise<Octokit> {
   const { ghesApiUrl } = getGitHubEnterpriseApiUrl();
   const ghAuthPre = await createGithubAppAuth(undefined, ghesApiUrl);
   const appIdx = ghAuthPre.appIndex;
+  // Re-evaluate quota before consulting the cache; an owner can use another
+  // installation when the previously selected app becomes exhausted.
+  const key = `${appIdx ?? 0}:${runner.type}:${runner.owner}`;
+  const cachedOctokit = githubCache.clients.get(key);
+  if (cachedOctokit) return cachedOctokit;
 
   // Use the pre-configured installation ID when available (avoids an API call).
   let installationId = await getStoredInstallationId(appIdx);
   if (installationId === undefined) {
     const githubClientPre = await createOctokitClient(ghAuthPre.token, ghesApiUrl, appIdx);
+    trackAppQuota(githubClientPre, appIdx);
     installationId =
       runner.type === 'Org'
         ? (
@@ -58,6 +67,7 @@ async function getOrCreateOctokit(runner: RunnerInfo): Promise<Octokit> {
   }
   const ghAuth = await createGithubInstallationAuth(installationId, ghesApiUrl, appIdx);
   const octokit = await createOctokitClient(ghAuth.token, ghesApiUrl, appIdx);
+  trackAppQuota(octokit, appIdx);
   githubCache.clients.set(key, octokit);
 
   return octokit;
@@ -80,7 +90,6 @@ async function getGitHubSelfHostedRunnerState(
             owner: runner.owner.split('/')[0],
             repo: runner.owner.split('/')[1],
           });
-    metricGitHubAppRateLimit(state.headers);
 
     return state.data;
   } catch (error) {