feat(termination-watcher): deregister runners from GitHub on EC2 termination

README.md

@@ -146,7 +146,7 @@ Join our discord community via [this invite link](https://discord.gg/bxgXW8jJGh)
 | <a name="input_instance_max_spot_price"></a> [instance\_max\_spot\_price](#input\_instance\_max\_spot\_price) | Max price price for spot instances per hour. This variable will be passed to the create fleet as max spot price for the fleet. | `string` | `null` | no |
 | <a name="input_instance_profile_path"></a> [instance\_profile\_path](#input\_instance\_profile\_path) | The path that will be added to the instance\_profile, if not set the environment name will be used. | `string` | `null` | no |
 | <a name="input_instance_target_capacity_type"></a> [instance\_target\_capacity\_type](#input\_instance\_target\_capacity\_type) | Default lifecycle used for runner instances, can be either `spot` or `on-demand`. | `string` | `"spot"` | no |
-| <a name="input_instance_termination_watcher"></a> [instance\_termination\_watcher](#input\_instance\_termination\_watcher) | Configuration for the instance termination watcher. This feature is Beta, changes will not trigger a major release as long in beta.<br/><br/>`enable`: Enable or disable the spot termination watcher.<br/>'features': Enable or disable features of the termination watcher.<br/>`memory_size`: Memory size limit in MB of the lambda.<br/>`s3_key`: S3 key for syncer lambda function. Required if using S3 bucket to specify lambdas.<br/>`s3_object_version`: S3 object version for syncer lambda function. Useful if S3 versioning is enabled on source bucket.<br/>`timeout`: Time out of the lambda in seconds.<br/>`zip`: File location of the lambda zip file. | <pre>object({<br/>    enable = optional(bool, false)<br/>    features = optional(object({<br/>      enable_spot_termination_handler              = optional(bool, true)<br/>      enable_spot_termination_notification_watcher = optional(bool, true)<br/>    }), {})<br/>    memory_size       = optional(number, null)<br/>    s3_key            = optional(string, null)<br/>    s3_object_version = optional(string, null)<br/>    timeout           = optional(number, null)<br/>    zip               = optional(string, null)<br/>  })</pre> | `{}` | no |
+| <a name="input_instance_termination_watcher"></a> [instance\_termination\_watcher](#input\_instance\_termination\_watcher) | Configuration for the instance termination watcher. This feature is Beta, changes will not trigger a major release as long in beta.<br/><br/>`enable`: Enable or disable the spot termination watcher.<br/>'features': Enable or disable features of the termination watcher.<br/>`enable_runner_deregistration`: Enable or disable deregistering the runner from GitHub when its EC2 instance is terminated.<br/>`memory_size`: Memory size limit in MB of the lambda.<br/>`s3_key`: S3 key for syncer lambda function. Required if using S3 bucket to specify lambdas.<br/>`s3_object_version`: S3 object version for syncer lambda function. Useful if S3 versioning is enabled on source bucket.<br/>`timeout`: Time out of the lambda in seconds.<br/>`zip`: File location of the lambda zip file. | <pre>object({<br/>    enable = optional(bool, false)<br/>    features = optional(object({<br/>      enable_spot_termination_handler              = optional(bool, true)<br/>      enable_spot_termination_notification_watcher = optional(bool, true)<br/>    }), {})<br/>    enable_runner_deregistration = optional(bool, true)<br/>    memory_size                  = optional(number, null)<br/>    s3_key                       = optional(string, null)<br/>    s3_object_version            = optional(string, null)<br/>    timeout                      = optional(number, null)<br/>    zip                          = optional(string, null)<br/>  })</pre> | `{}` | no |
 | <a name="input_instance_types"></a> [instance\_types](#input\_instance\_types) | List of instance types for the action runner. Defaults are based on runner\_os (al2023 for linux, macOS Sequoia for osx, Windows Server Core for win). | `list(string)` | <pre>[<br/>  "m5.large",<br/>  "c5.large"<br/>]</pre> | no |
 | <a name="input_job_queue_retention_in_seconds"></a> [job\_queue\_retention\_in\_seconds](#input\_job\_queue\_retention\_in\_seconds) | The number of seconds the job is held in the queue before it is purged. | `number` | `86400` | no |
 | <a name="input_job_retry"></a> [job\_retry](#input\_job\_retry) | Experimental! Can be removed / changed without trigger a major release.Configure job retries. The configuration enables job retries (for ephemeral runners). After creating the instances a message will be published to a job retry queue. The job retry check lambda is checking after a delay if the job is queued. If not the message will be published again on the scale-up (build queue). Using this feature can impact the rate limit of the GitHub app.<br/><br/>`enable`: Enable or disable the job retry feature.<br/>`delay_in_seconds`: The delay in seconds before the job retry check lambda will check the job status.<br/>`delay_backoff`: The backoff factor for the delay.<br/>`lambda_memory_size`: Memory size limit in MB for the job retry check lambda.<br/>`lambda_timeout`: Time out of the job retry check lambda in seconds.<br/>`max_attempts`: The maximum number of attempts to retry the job. | <pre>object({<br/>    enable             = optional(bool, false)<br/>    delay_in_seconds   = optional(number, 300)<br/>    delay_backoff      = optional(number, 2)<br/>    lambda_memory_size = optional(number, 256)<br/>    lambda_timeout     = optional(number, 30)<br/>    max_attempts       = optional(number, 1)<br/>  })</pre> | `{}` | no |

lambdas/functions/termination-watcher/package.json

@@ -24,8 +24,15 @@
   },
   "dependencies": {
     "@aws-github-runner/aws-powertools-util": "*",
+    "@aws-github-runner/aws-ssm-util": "*",
     "@aws-sdk/client-ec2": "^3.1009.0",
-    "@middy/core": "^6.4.5"
+    "@aws-sdk/client-sqs": "^3.1009.0",
+    "@middy/core": "^6.4.5",
+    "@octokit/auth-app": "8.2.0",
+    "@octokit/core": "7.0.6",
+    "@octokit/plugin-throttling": "11.0.3",
+    "@octokit/request": "^9.2.2",
+    "@octokit/rest": "22.0.1"
   },
   "nx": {
     "includedScripts": [

lambdas/functions/termination-watcher/src/ConfigResolver.test.ts

@@ -37,6 +37,8 @@ describe('Test ConfigResolver', () => {
       delete process.env.ENABLE_METRICS_SPOT_WARNING;
       delete process.env.PREFIX;
       delete process.env.TAG_FILTERS;
+      delete process.env.ENABLE_RUNNER_DEREGISTRATION;
+      delete process.env.GHES_URL;
     });
 
     it(description, async () => {
@@ -55,4 +57,29 @@ describe('Test ConfigResolver', () => {
       expect(config.tagFilters).toEqual(output.tagFilters);
     });
   });
+
+  describe('runner deregistration config', () => {
+    beforeEach(() => {
+      delete process.env.ENABLE_RUNNER_DEREGISTRATION;
+      delete process.env.GHES_URL;
+    });
+
+    it('should default to disabled', () => {
+      const config = new Config();
+      expect(config.enableRunnerDeregistration).toBe(false);
+      expect(config.ghesApiUrl).toBe('');
+    });
+
+    it('should enable deregistration when env var is true', () => {
+      process.env.ENABLE_RUNNER_DEREGISTRATION = 'true';
+      const config = new Config();
+      expect(config.enableRunnerDeregistration).toBe(true);
+    });
+
+    it('should set GHES URL when provided', () => {
+      process.env.GHES_URL = 'https://github.internal.co/api/v3';
+      const config = new Config();
+      expect(config.ghesApiUrl).toBe('https://github.internal.co/api/v3');
+    });
+  });
 });

lambdas/functions/termination-watcher/src/ConfigResolver.ts

@@ -5,6 +5,8 @@ export class Config {
   createSpotTerminationMetric: boolean;
   tagFilters: Record<string, string>;
   prefix: string;
+  enableRunnerDeregistration: boolean;
+  ghesApiUrl: string;
 
   constructor() {
     const logger = createChildLogger('config-resolver');
@@ -14,6 +16,8 @@ export class Config {
     this.createSpotWarningMetric = process.env.ENABLE_METRICS_SPOT_WARNING === 'true';
     this.createSpotTerminationMetric = process.env.ENABLE_METRICS_SPOT_TERMINATION === 'true';
     this.prefix = process.env.PREFIX ?? '';
+    this.enableRunnerDeregistration = process.env.ENABLE_RUNNER_DEREGISTRATION === 'true';
+    this.ghesApiUrl = process.env.GHES_URL ?? '';
     this.tagFilters = { 'ghr:environment': this.prefix };
 
     const rawTagFilters = process.env.TAG_FILTERS;

lambdas/functions/termination-watcher/src/deregister.test.ts

@@ -0,0 +1,295 @@
+import { Instance } from '@aws-sdk/client-ec2';
+import { describe, it, expect, beforeEach, vi } from 'vitest';
+import { deregisterRunner, createThrottleOptions } from './deregister';
+import { Config } from './ConfigResolver';
+import type { EndpointDefaults } from '@octokit/types';
+
+const mockGetParameter = vi.fn();
+vi.mock('@aws-github-runner/aws-ssm-util', () => ({
+  getParameter: (...args: unknown[]) => mockGetParameter(...args),
+}));
+
+const mockCreateAppAuth = vi.fn();
+vi.mock('@octokit/auth-app', () => ({
+  createAppAuth: (...args: unknown[]) => mockCreateAppAuth(...args),
+}));
+
+const mockPaginate = {
+  iterator: vi.fn(),
+};
+
+const mockActions = {
+  listSelfHostedRunnersForOrg: vi.fn(),
+  listSelfHostedRunnersForRepo: vi.fn(),
+  deleteSelfHostedRunnerFromOrg: vi.fn(),
+  deleteSelfHostedRunnerFromRepo: vi.fn(),
+};
+
+const mockApps = {
+  getOrgInstallation: vi.fn(),
+  getRepoInstallation: vi.fn(),
+};
+
+function MockOctokit() {
+  return {
+    actions: mockActions,
+    apps: mockApps,
+    paginate: mockPaginate,
+  };
+}
+MockOctokit.plugin = vi.fn().mockReturnValue(MockOctokit);
+
+vi.mock('@octokit/rest', () => ({
+  Octokit: MockOctokit,
+}));
+
+vi.mock('@octokit/plugin-throttling', () => ({
+  throttling: vi.fn(),
+}));
+
+vi.mock('@octokit/request', () => ({
+  request: {
+    defaults: vi.fn().mockReturnValue(vi.fn()),
+  },
+}));
+
+const baseConfig: Config = {
+  createSpotWarningMetric: false,
+  createSpotTerminationMetric: true,
+  tagFilters: { 'ghr:environment': 'test' },
+  prefix: 'runners',
+  enableRunnerDeregistration: true,
+  ghesApiUrl: '',
+};
+
+const orgInstance: Instance = {
+  InstanceId: 'i-12345678901234567',
+  InstanceType: 't2.micro',
+  Tags: [
+    { Key: 'Name', Value: 'test-instance' },
+    { Key: 'ghr:environment', Value: 'test' },
+    { Key: 'ghr:Owner', Value: 'test-org' },
+    { Key: 'ghr:Type', Value: 'Org' },
+  ],
+  State: { Name: 'running' },
+  LaunchTime: new Date('2021-01-01'),
+};
+
+const repoInstance: Instance = {
+  InstanceId: 'i-repo12345678901234',
+  InstanceType: 't2.micro',
+  Tags: [
+    { Key: 'Name', Value: 'test-repo-instance' },
+    { Key: 'ghr:environment', Value: 'test' },
+    { Key: 'ghr:Owner', Value: 'test-org/test-repo' },
+    { Key: 'ghr:Type', Value: 'Repo' },
+  ],
+  State: { Name: 'running' },
+  LaunchTime: new Date('2021-01-01'),
+};
+
+function setupAuthMocks() {
+  const appPrivateKey = Buffer.from('fake-private-key').toString('base64');
+  mockGetParameter.mockImplementation((name: string) => {
+    if (name === 'github-app-id') return Promise.resolve('12345');
+    if (name === 'github-app-key') return Promise.resolve(appPrivateKey);
+    return Promise.reject(new Error(`Unknown parameter: ${name}`));
+  });
+
+  // App auth returns app token
+  const mockAuth = vi.fn();
+  mockAuth.mockImplementation((opts: { type: string }) => {
+    if (opts.type === 'app') {
+      return Promise.resolve({ token: 'app-token' });
+    }
+    return Promise.resolve({ token: 'installation-token' });
+  });
+  mockCreateAppAuth.mockReturnValue(mockAuth);
+}
+
+describe('deregisterRunner', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    process.env.PARAMETER_GITHUB_APP_ID_NAME = 'github-app-id';
+    process.env.PARAMETER_GITHUB_APP_KEY_BASE64_NAME = 'github-app-key';
+    setupAuthMocks();
+  });
+
+  it('should skip deregistration when disabled', async () => {
+    await deregisterRunner(orgInstance, { ...baseConfig, enableRunnerDeregistration: false });
+    expect(mockGetParameter).not.toHaveBeenCalled();
+  });
+
+  it('should skip deregistration when instance ID is missing', async () => {
+    const instance: Instance = { ...orgInstance, InstanceId: undefined };
+    await deregisterRunner(instance, baseConfig);
+    expect(mockGetParameter).not.toHaveBeenCalled();
+  });
+
+  it('should skip deregistration when ghr:Owner tag is missing', async () => {
+    const instance: Instance = {
+      ...orgInstance,
+      Tags: [{ Key: 'Name', Value: 'test' }],
+    };
+    await deregisterRunner(instance, baseConfig);
+    // Auth should not be called since we bail early
+    expect(mockCreateAppAuth).not.toHaveBeenCalled();
+  });
+
+  it('should deregister an org runner successfully', async () => {
+    mockApps.getOrgInstallation.mockResolvedValue({ data: { id: 999 } });
+
+    async function* fakeIterator() {
+      yield { data: [{ id: 42, name: `runner-i-12345678901234567` }] };
+    }
+    mockPaginate.iterator.mockReturnValue(fakeIterator());
+
+    mockActions.deleteSelfHostedRunnerFromOrg.mockResolvedValue({});
+
+    await deregisterRunner(orgInstance, baseConfig);
+
+    expect(mockApps.getOrgInstallation).toHaveBeenCalledWith({ org: 'test-org' });
+    expect(mockActions.deleteSelfHostedRunnerFromOrg).toHaveBeenCalledWith({
+      org: 'test-org',
+      runner_id: 42,
+    });
+  });
+
+  it('should deregister a repo runner successfully', async () => {
+    mockApps.getRepoInstallation.mockResolvedValue({ data: { id: 888 } });
+
+    async function* fakeIterator() {
+      yield { data: [{ id: 55, name: `runner-i-repo12345678901234` }] };
+    }
+    mockPaginate.iterator.mockReturnValue(fakeIterator());
+
+    mockActions.deleteSelfHostedRunnerFromRepo.mockResolvedValue({});
+
+    await deregisterRunner(repoInstance, baseConfig);
+
+    expect(mockApps.getRepoInstallation).toHaveBeenCalledWith({ owner: 'test-org', repo: 'test-repo' });
+    expect(mockActions.deleteSelfHostedRunnerFromRepo).toHaveBeenCalledWith({
+      owner: 'test-org',
+      repo: 'test-repo',
+      runner_id: 55,
+    });
+  });
+
+  it('should handle runner not found gracefully', async () => {
+    mockApps.getOrgInstallation.mockResolvedValue({ data: { id: 999 } });
+
+    async function* fakeIterator() {
+      yield { data: [{ id: 42, name: 'runner-other-instance' }] };
+    }
+    mockPaginate.iterator.mockReturnValue(fakeIterator());
+
+    await deregisterRunner(orgInstance, baseConfig);
+
+    expect(mockActions.deleteSelfHostedRunnerFromOrg).not.toHaveBeenCalled();
+  });
+
+  it('should handle GitHub API errors gracefully', async () => {
+    mockApps.getOrgInstallation.mockRejectedValue(new Error('GitHub API error'));
+
+    await deregisterRunner(orgInstance, baseConfig);
+
+    // Should not throw — error is caught internally
+    expect(mockActions.deleteSelfHostedRunnerFromOrg).not.toHaveBeenCalled();
+  });
+
+  it('should default to Org runner type when ghr:Type tag is missing', async () => {
+    const instance: Instance = {
+      ...orgInstance,
+      Tags: [
+        { Key: 'ghr:environment', Value: 'test' },
+        { Key: 'ghr:Owner', Value: 'test-org' },
+      ],
+    };
+
+    mockApps.getOrgInstallation.mockResolvedValue({ data: { id: 999 } });
+
+    async function* fakeIterator() {
+      yield { data: [{ id: 42, name: `runner-i-12345678901234567` }] };
+    }
+    mockPaginate.iterator.mockReturnValue(fakeIterator());
+
+    mockActions.deleteSelfHostedRunnerFromOrg.mockResolvedValue({});
+
+    await deregisterRunner(instance, baseConfig);
+
+    expect(mockApps.getOrgInstallation).toHaveBeenCalledWith({ org: 'test-org' });
+    expect(mockActions.deleteSelfHostedRunnerFromOrg).toHaveBeenCalledWith({
+      org: 'test-org',
+      runner_id: 42,
+    });
+  });
+
+  it('should use GHES API URL when configured', async () => {
+    const ghesConfig = { ...baseConfig, ghesApiUrl: 'https://github.internal.co/api/v3' };
+
+    mockApps.getOrgInstallation.mockResolvedValue({ data: { id: 999 } });
+
+    async function* fakeIterator() {
+      yield { data: [{ id: 42, name: `runner-i-12345678901234567` }] };
+    }
+    mockPaginate.iterator.mockReturnValue(fakeIterator());
+
+    mockActions.deleteSelfHostedRunnerFromOrg.mockResolvedValue({});
+
+    await deregisterRunner(orgInstance, ghesConfig);
+
+    expect(mockActions.deleteSelfHostedRunnerFromOrg).toHaveBeenCalled();
+  });
+
+  it('should paginate through multiple pages to find runner', async () => {
+    mockApps.getOrgInstallation.mockResolvedValue({ data: { id: 999 } });
+
+    async function* fakeIterator() {
+      yield { data: [{ id: 1, name: 'runner-other-1' }] };
+      yield { data: [{ id: 2, name: 'runner-other-2' }] };
+      yield { data: [{ id: 42, name: `runner-i-12345678901234567` }] };
+    }
+    mockPaginate.iterator.mockReturnValue(fakeIterator());
+
+    mockActions.deleteSelfHostedRunnerFromOrg.mockResolvedValue({});
+
+    await deregisterRunner(orgInstance, baseConfig);
+
+    expect(mockActions.deleteSelfHostedRunnerFromOrg).toHaveBeenCalledWith({
+      org: 'test-org',
+      runner_id: 42,
+    });
+  });
+
+  it('should handle repo runner not found gracefully', async () => {
+    mockApps.getRepoInstallation.mockResolvedValue({ data: { id: 888 } });
+
+    async function* fakeIterator() {
+      yield { data: [{ id: 99, name: 'runner-other-instance' }] };
+    }
+    mockPaginate.iterator.mockReturnValue(fakeIterator());
+
+    await deregisterRunner(repoInstance, baseConfig);
+
+    expect(mockActions.deleteSelfHostedRunnerFromRepo).not.toHaveBeenCalled();
+  });
+
+  it('should handle instance with no tags', async () => {
+    const instance: Instance = {
+      InstanceId: 'i-12345678901234567',
+      Tags: undefined,
+    };
+    await deregisterRunner(instance, baseConfig);
+    expect(mockCreateAppAuth).not.toHaveBeenCalled();
+  });
+});
+
+describe('createThrottleOptions', () => {
+  it('should return false for rate limit and log warning', () => {
+    const options = createThrottleOptions();
+    const endpointDefaults = { method: 'GET', url: '/test' } as Required<EndpointDefaults>;
+
+    expect(options.onRateLimit(60, endpointDefaults)).toBe(false);
+    expect(options.onSecondaryRateLimit(60, endpointDefaults)).toBe(false);
+  });
+});