feat: Add metric to track GitHub app rate limit

Author: npalm

docs/configuration.md

@@ -191,6 +191,15 @@ This feature has been disabled by default.
 
 The watcher will act on all spot termination notificatins and log all onses relevant to the runner module. Therefor we suggest to only deploy the watcher once. You can either deploy the watcher by enabling in one of your deployments or deploy the watcher as a stand alone module.
 
+## Metrics
+
+The module supports metrics (experimental feature) to monitor the system. The metrics are disabled by default. To enable the metrics set `metrics.enable = true`. Once put by true all module managed metrics are used, you can configure the one bye one via the `metrics` object. The metrics are created in the namespace `GitHub Runners`.
+
+### Supported metrics
+
+- **GitHubAppRateLimitRemaining**: Remaining rate limit for the GitHub App.
+- **JobRetry**: Number of job retries, only relevant when job retry is enabled.
+- **SpotInterruptionWarning**: Number of spot interruption warnings received by the termination watcher, only relevant when the termination watcher is enabled.
 
 ## Debugging
 

examples/default/main.tf

@@ -114,20 +114,24 @@ module "runners" {
 
   instance_termination_watcher = {
     enable = true
-    enable_metric = {
-      spot_warning = true
-    }
   }
 
-  # enable job_retry feature. Be careful with this feature, it can lead to API rate limits.
-  # job_retry = {
-  #   enable           = true
-  #   max_attempts     = 1
-  #   delay_in_seconds = 180
+  # enable metric creation  (experimental)
+  # metrics = {
+  #   enable = true
+  #   metric = {
+  #     enable_spot_termination_warning = true
+  #     enable_job_retry                = false
+  #     enable_github_app_rate_limit    = true
+  #   }
   # }
 
-  # enable metric creation by the control plane (experimental)
-  # enable_metrics_control_plane = true
+  # enable job_retry feature. Be careful with this feature, it can lead to API rate limits.
+  job_retry = {
+    enable           = true
+    max_attempts     = 1
+    delay_in_seconds = 180
+  }
 
   # enable CMK instead of aws managed key for encryptions
   # kms_key_arn = aws_kms_key.github.arn

examples/multi-runner/main.tf

@@ -103,8 +103,15 @@ module "runners" {
   # Enable to track the spot instance termination warning
   # instance_termination_watcher = {
   #   enable         = true
-  #   enable_metric = {
-  #     spot_warning = true
+  # }
+
+  # Enable metrics
+  # metrics = {
+  #   enable = true
+  #   metric = {
+  #     enable_github_app_rate_limit    = true
+  #     enable_job_retry                = false
+  #     enable_spot_termination_warning = true
   #   }
   # }
 }

lambdas/functions/control-plane/package.json

@@ -38,6 +38,7 @@
     "ts-node-dev": "^2.0.0"
   },
   "dependencies": {
+    "@aws-lambda-powertools/parameters": "^2.7.0",
     "@aws-sdk/client-ec2": "^3.637.0",
     "@aws-sdk/client-sqs": "^3.637.0",
     "@aws-sdk/types": "^3.609.0",

lambdas/functions/control-plane/src/gh-auth/rate-limit.test.ts

@@ -0,0 +1,70 @@
+import { ResponseHeaders } from '@octokit/types';
+import { createSingleMetric } from '@terraform-aws-github-runner/aws-powertools-util';
+import { MetricUnit } from '@aws-lambda-powertools/metrics';
+import { metricGitHubAppRateLimit } from './rate-limit';
+
+process.env.PARAMETER_GITHUB_APP_ID_NAME = 'test';
+jest.mock('@terraform-aws-github-runner/aws-ssm-util', () => ({
+  ...jest.requireActual('@terraform-aws-github-runner/aws-ssm-util'),
+  // get parameter name from process.env.PARAMETER_GITHUB_APP_ID_NAME rerunt 1234
+  getParameter: jest.fn((name: string) => {
+    if (name === process.env.PARAMETER_GITHUB_APP_ID_NAME) {
+      return '1234';
+    } else {
+      return '';
+    }
+  }),
+}));
+
+jest.mock('@terraform-aws-github-runner/aws-powertools-util', () => ({
+  ...jest.requireActual('@terraform-aws-github-runner/aws-powertools-util'),
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  createSingleMetric: jest.fn((name: string, unit: string, value: number, dimensions?: Record<string, string>) => {
+    return {
+      addMetadata: jest.fn(),
+    };
+  }),
+}));
+
+describe('metricGitHubAppRateLimit', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it('should update rate limit metric', async () => {
+    // set process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT to true
+    process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT = 'true';
+    const headers: ResponseHeaders = {
+      'x-ratelimit-remaining': '10',
+      'x-ratelimit-limit': '60',
+    };
+
+    await metricGitHubAppRateLimit(headers);
+
+    expect(createSingleMetric).toHaveBeenCalledWith('GitHubAppRateLimitRemaining', MetricUnit.Count, 10, {
+      AppId: '1234',
+    });
+  });
+
+  it('should not update rate limit metric', async () => {
+    // set process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT to false
+    process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT = 'false';
+    const headers: ResponseHeaders = {
+      'x-ratelimit-remaining': '10',
+      'x-ratelimit-limit': '60',
+    };
+
+    await metricGitHubAppRateLimit(headers);
+
+    expect(createSingleMetric).not.toHaveBeenCalled();
+  });
+
+  it('should not update rate limit metric if headers are undefined', async () => {
+    // set process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT to true
+    process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT = 'true';
+
+    await metricGitHubAppRateLimit(undefined as unknown as ResponseHeaders);
+
+    expect(createSingleMetric).not.toHaveBeenCalled();
+  });
+});

lambdas/functions/control-plane/src/gh-auth/rate-limit.ts

@@ -0,0 +1,25 @@
+import { ResponseHeaders } from '@octokit/types';
+import { createSingleMetric, logger } from '@terraform-aws-github-runner/aws-powertools-util';
+import { MetricUnit } from '@aws-lambda-powertools/metrics';
+import yn from 'yn';
+import { getParameter } from '@terraform-aws-github-runner/aws-ssm-util';
+
+export async function metricGitHubAppRateLimit(headers: ResponseHeaders): Promise<void> {
+  try {
+    const remaining = parseInt(headers['x-ratelimit-remaining'] as string);
+    const limit = parseInt(headers['x-ratelimit-limit'] as string);
+
+    logger.debug(`Rate limit remaining: ${remaining}, limit: ${limit}`);
+
+    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 metric = createSingleMetric('GitHubAppRateLimitRemaining', MetricUnit.Count, remaining, {
+        AppId: appId,
+      });
+      metric.addMetadata('AppId', appId);
+    }
+  } catch (e) {
+    logger.debug(`Error updating rate limit metric`, { error: e });
+  }
+}

lambdas/functions/control-plane/src/modules.d.ts

@@ -1,6 +1,7 @@
 declare namespace NodeJS {
   export interface ProcessEnv {
     AWS_REGION: string;
+    ENABLE_METRIC_GITHUB_APP_RATE_LIMIT: string;
     ENABLE_ON_DEMAND_FAILOVER_FOR_ERRORS: string;
     ENVIRONMENT: string;
     GHES_URL: string;

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

@@ -179,7 +179,7 @@ describe(`Test job retry check`, () => {
     process.env.ENABLE_ORGANIZATION_RUNNERS = 'true';
     process.env.ENVIRONMENT = 'test';
     process.env.RUNNER_NAME_PREFIX = 'test';
-    process.env.ENABLE_METRICS = 'true';
+    process.env.ENABLE_METRIC_JOB_RETRY = 'true';
     process.env.JOB_QUEUE_SCALE_UP_URL =
       'https://sqs.eu-west-1.amazonaws.com/123456789/webhook_events_workflow_job_queue';
 

lambdas/functions/control-plane/src/scale-runners/job-retry.ts

@@ -46,7 +46,7 @@ export async function checkAndRetryJob(payload: ActionRequestMessageRetry): Prom
   const runnerOwner = enableOrgLevel ? payload.repositoryOwner : `${payload.repositoryOwner}/${payload.repositoryName}`;
   const runnerNamePrefix = process.env.RUNNER_NAME_PREFIX ?? '';
   const jobQueueUrl = process.env.JOB_QUEUE_SCALE_UP_URL ?? '';
-  const enableMetrics = yn(process.env.ENABLE_METRICS, { default: false });
+  const enableMetrics = yn(process.env.ENABLE_METRIC_JOB_RETRY, { default: false });
   const environment = process.env.ENVIRONMENT;
 
   addPersistentContextToChildLogger({

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

@@ -7,6 +7,7 @@ import { bootTimeExceeded, listEC2Runners, tag, terminateRunner } from './../aws
 import { RunnerInfo, RunnerList } from './../aws/runners.d';
 import { GhRunners, githubCache } from './cache';
 import { ScalingDownConfig, getEvictionStrategy, getIdleRunnerCount } from './scale-down-config';
+import { metricGitHubAppRateLimit } from '../gh-auth/rate-limit';
 
 const logger = createChildLogger('scale-down');
 
@@ -63,6 +64,8 @@ async function getGitHubRunnerBusyState(client: Octokit, ec2runner: RunnerInfo,
 
   logger.info(`Runner '${ec2runner.instanceId}' - GitHub Runner ID '${runnerId}' - Busy: ${state.data.busy}`);
 
+  metricGitHubAppRateLimit(state.headers);
+
   return state.data.busy;
 }
 

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

@@ -1,4 +1,4 @@
-import { GetParameterCommand, PutParameterCommand, SSMClient } from '@aws-sdk/client-ssm';
+import { PutParameterCommand, SSMClient } from '@aws-sdk/client-ssm';
 import { Octokit } from '@octokit/rest';
 import { mockClient } from 'aws-sdk-client-mock';
 import 'aws-sdk-client-mock-jest';
@@ -11,6 +11,7 @@ import { createRunner, listEC2Runners } from './../aws/runners';
 import { RunnerInputParameters } from './../aws/runners.d';
 import ScaleError from './ScaleError';
 import * as scaleUpModule from './scale-up';
+import { getParameter } from '@terraform-aws-github-runner/aws-ssm-util';
 
 const mockOctokit = {
   paginate: jest.fn(),
@@ -30,13 +31,20 @@ const mockOctokit = {
 const mockCreateRunner = mocked(createRunner);
 const mockListRunners = mocked(listEC2Runners);
 const mockSSMClient = mockClient(SSMClient);
+const mockSSM = mocked(getParameter);
 
 jest.mock('@octokit/rest', () => ({
   Octokit: jest.fn().mockImplementation(() => mockOctokit),
 }));
 
 jest.mock('./../aws/runners');
 jest.mock('./../gh-auth/gh-auth');
+
+jest.mock('@terraform-aws-github-runner/aws-ssm-util', () => ({
+  ...jest.requireActual('@terraform-aws-github-runner/aws-ssm-util'),
+  getParameter: jest.fn(),
+}));
+
 export type RunnerType = 'ephemeral' | 'non-ephemeral';
 
 // for ephemeral and non-ephemeral runners
@@ -172,6 +180,10 @@ beforeEach(() => {
   });
 
   mockCreateClient.mockResolvedValue(new mocktokit());
+
+  mockSSM.mockImplementation(async () => {
+    return '1';
+  });
 });
 
 describe('scaleUp with GHES', () => {
@@ -213,7 +225,6 @@ describe('scaleUp with GHES', () => {
 
       expectedRunnerParams = { ...EXPECTED_RUNNER_PARAMS };
       mockSSMClient.reset();
-      mockSSMClient.on(GetParameterCommand).resolves({ Parameter: { Value: '1' } });
     });
 
     it('gets the current org level runners', async () => {
@@ -270,7 +281,9 @@ describe('scaleUp with GHES', () => {
 
     it('Throws an error if runner group doesnt exist for ephemeral runners', async () => {
       process.env.RUNNER_GROUP_NAME = 'test-runner-group';
-      mockSSMClient.on(GetParameterCommand).rejects();
+      mockSSM.mockImplementation(async () => {
+        throw new Error('ParameterNotFound');
+      });
       await expect(scaleUpModule.scaleUp('aws:sqs', TEST_DATA)).rejects.toBeInstanceOf(Error);
       expect(mockOctokit.paginate).toHaveBeenCalledTimes(1);
     });
@@ -284,7 +297,9 @@ describe('scaleUp with GHES', () => {
     });
 
     it('create SSM parameter for runner group id if it doesnt exist', async () => {
-      mockSSMClient.on(GetParameterCommand).rejects();
+      mockSSM.mockImplementation(async () => {
+        throw new Error('ParameterNotFound');
+      });
       await scaleUpModule.scaleUp('aws:sqs', TEST_DATA);
       expect(mockOctokit.paginate).toHaveBeenCalledTimes(1);
       expect(mockSSMClient).toHaveReceivedCommandTimes(PutParameterCommand, 2);
@@ -295,16 +310,14 @@ describe('scaleUp with GHES', () => {
       });
     });
 
-    it('Doesnt create SSM parameter for runner group id if it exists', async () => {
-      mockSSMClient.on(GetParameterCommand).resolves({ Parameter: { Value: '1' } });
+    it('Does not create SSM parameter for runner group id if it exists', async () => {
       await scaleUpModule.scaleUp('aws:sqs', TEST_DATA);
       expect(mockOctokit.paginate).toHaveBeenCalledTimes(0);
       expect(mockSSMClient).toHaveReceivedCommandTimes(PutParameterCommand, 1);
     });
 
     it('create start runner config for ephemeral runners ', async () => {
       process.env.RUNNERS_MAXIMUM_COUNT = '2';
-      mockSSMClient.on(GetParameterCommand).resolves({ Parameter: { Value: '1' } });
       await scaleUpModule.scaleUp('aws:sqs', TEST_DATA);
       expect(mockOctokit.actions.generateRunnerJitconfigForOrg).toBeCalledWith({
         org: TEST_DATA.repositoryOwner,
@@ -356,7 +369,6 @@ describe('scaleUp with GHES', () => {
         mockListRunners.mockImplementation(async () => {
           return [];
         });
-        mockSSMClient.on(GetParameterCommand).resolves({ Parameter: { Value: '1' } });
         const startTime = performance.now();
         const instances = [
           'i-1234',

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

@@ -8,6 +8,7 @@ import { createRunner, listEC2Runners } from './../aws/runners';
 import { RunnerInputParameters } from './../aws/runners.d';
 import ScaleError from './ScaleError';
 import { publishRetryMessage } from './job-retry';
+import { metricGitHubAppRateLimit } from '../gh-auth/rate-limit';
 
 const logger = createChildLogger('scale-up');
 
@@ -94,6 +95,9 @@ async function getGithubRunnerRegistrationToken(githubRunnerConfig: CreateGitHub
           owner: githubRunnerConfig.runnerOwner.split('/')[0],
           repo: githubRunnerConfig.runnerOwner.split('/')[1],
         });
+
+  const appId = parseInt(await getParameter(process.env.PARAMETER_GITHUB_APP_ID_NAME));
+  logger.info('App id from SSM', { appId: appId });
   return registrationToken.data.token;
 }
 
@@ -142,6 +146,7 @@ export async function isJobQueued(githubInstallationClient: Octokit, payload: Ac
       owner: payload.repositoryOwner,
       repo: payload.repositoryName,
     });
+    metricGitHubAppRateLimit(jobForWorkflowRun.headers);
     isQueued = jobForWorkflowRun.data.status === 'queued';
   } else {
     throw Error(`Event ${payload.eventType} is not supported`);
@@ -169,7 +174,7 @@ async function getRunnerGroupId(githubRunnerConfig: CreateGitHubRunnerConfig, gh
     }
     if (runnerGroup === undefined) {
       // get runner group id from GitHub
-      runnerGroupId = await GetRunnerGroupByName(ghClient, githubRunnerConfig);
+      runnerGroupId = await getRunnerGroupByName(ghClient, githubRunnerConfig);
       // store runner group id in SSM
       try {
         await putParameter(
@@ -188,7 +193,7 @@ async function getRunnerGroupId(githubRunnerConfig: CreateGitHubRunnerConfig, gh
   return runnerGroupId;
 }
 
-async function GetRunnerGroupByName(ghClient: Octokit, githubRunnerConfig: CreateGitHubRunnerConfig): Promise<number> {
+async function getRunnerGroupByName(ghClient: Octokit, githubRunnerConfig: CreateGitHubRunnerConfig): Promise<number> {
   const runnerGroups: RunnerGroup[] = await ghClient.paginate(`GET /orgs/{org}/actions/runner-groups`, {
     org: githubRunnerConfig.runnerOwner,
     per_page: 100,
@@ -432,6 +437,8 @@ async function createJitConfig(githubRunnerConfig: CreateGitHubRunnerConfig, ins
             labels: ephemeralRunnerConfig.runnerLabels,
           });
 
+    metricGitHubAppRateLimit(runnerConfig.headers);
+
     // store jit config in ssm parameter store
     logger.debug('Runner JIT config for ephemeral runner generated.', {
       instance: instance,

lambdas/libs/aws-ssm-util/src/index.test.ts

@@ -136,9 +136,6 @@ describe('Test getParameter and putParameter', () => {
     mockSSMClient.on(GetParameterCommand).resolves(output);
 
     // Act
-    const result = await getParameter(parameterName);
-
-    // Assert
-    expect(result).toBe(undefined);
+    await expect(getParameter(parameterName)).rejects.toThrow(`Parameter ${parameterName} not found`);
   });
 });