feat(pool): Allow to optionally include busy runners into the pool (github-aws-runners#5078)

## Description
'pool_config' is useful when one wants to ensure that a certain number
of runners is running within given time periods.
However, when 'pool_config' is used in a combination with persistent
runners,
the pool lambda tends to spin up new runner instance when it sees that
some
existing runners are online, but busy.

As a result, with frequent pool lambda checks, the total size of the
runners
pool tends to grow more than desired (without respecting
'runners_maximum_count', which is only considered by the 'scale_up'
lambda).

These changes introduce a new 'pool_include_busy_runners' module
variable
to make it possbile to include all online runners (both idle and busy)
into the pool, so that the pool's lambda only tops up the pool if not
enough runners are online.
<!-- Briefly describe the changes in this PR. -->

## Test Plan

Tested on a module instance with persistent runners, also added a few
unit tests to the pool lambda itself.
<!--
Please describe how you tested these changes. This helps reviewers
understand
  the scope and confidence level of the change. Examples:

  - Ran `terraform plan` against an existing deployment
  - Added/updated unit tests (link to test file or describe coverage)
  - Deployed to a test environment and triggered a workflow run
  - Validated with `terraform validate` and `tflint`
-->



## Related Issues

<!-- Link any related issues, e.g. Fixes github-aws-runners#123, Closes github-aws-runners#456 -->

Author: rvem

README.md

@@ -170,6 +170,7 @@ Join our discord community via [this invite link](https://discord.gg/bxgXW8jJGh)
 | <a name="input_minimum_running_time_in_minutes"></a> [minimum\_running\_time\_in\_minutes](#input\_minimum\_running\_time\_in\_minutes) | The time an ec2 action runner should be running at minimum before terminated, if not busy. | `number` | `null` | no |
 | <a name="input_parameter_store_tags"></a> [parameter\_store\_tags](#input\_parameter\_store\_tags) | Map of tags that will be added to all the SSM Parameter Store parameters created by the Lambda function. | `map(string)` | `{}` | no |
 | <a name="input_pool_config"></a> [pool\_config](#input\_pool\_config) | The configuration for updating the pool. The `pool_size` to adjust to by the events triggered by the `schedule_expression`. For example you can configure a cron expression for weekdays to adjust the pool to 10 and another expression for the weekend to adjust the pool to 1. Use `schedule_expression_timezone` to override the schedule time zone (defaults to UTC). | <pre>list(object({<br/>    schedule_expression          = string<br/>    schedule_expression_timezone = optional(string)<br/>    size                         = number<br/>  }))</pre> | `[]` | no |
+| <a name="input_pool_include_busy_runners"></a> [pool\_include\_busy\_runners](#input\_pool\_include\_busy\_runners) | Include busy runners in the pool calculation. By default busy runners are not included in the pool. | `bool` | `false` | no |
 | <a name="input_pool_lambda_memory_size"></a> [pool\_lambda\_memory\_size](#input\_pool\_lambda\_memory\_size) | Memory size limit for scale-up lambda. | `number` | `512` | no |
 | <a name="input_pool_lambda_reserved_concurrent_executions"></a> [pool\_lambda\_reserved\_concurrent\_executions](#input\_pool\_lambda\_reserved\_concurrent\_executions) | Amount of reserved concurrent executions for the scale-up lambda function. A value of 0 disables lambda from being triggered and -1 removes any concurrency limitations. | `number` | `1` | no |
 | <a name="input_pool_lambda_timeout"></a> [pool\_lambda\_timeout](#input\_pool\_lambda\_timeout) | Time out for the pool lambda in seconds. | `number` | `60` | no |

lambdas/functions/control-plane/src/pool/pool.test.ts

@@ -424,4 +424,28 @@ describe('Test simple pool.', () => {
       );
     });
   });
+
+  describe('With INCLUDE_BUSY_RUNNERS enabled', () => {
+    beforeEach(() => {
+      process.env.INCLUDE_BUSY_RUNNERS = 'true';
+    });
+
+    it('Should not top up when pool size matches runners including busy online runners.', async () => {
+      // Without INCLUDE_BUSY_RUNNERS: 2 in pool (i-1-idle, i-4-idle-older). With it: 3 (adds i-2-busy).
+      await adjust({ poolSize: 3 });
+      expect(createRunners).not.toHaveBeenCalled();
+    });
+
+    it('Should top up by two runners when pool size is 5 and busy runners count toward the pool.', async () => {
+      await adjust({ poolSize: 5 });
+      // 3 in pool (idle, busy, older idle); need 2 more
+      expect(createRunners).toHaveBeenCalledWith(
+        expect.anything(),
+        expect.anything(),
+        2,
+        expect.anything(),
+        'pool-lambda',
+      );
+    });
+  });
 });

lambdas/functions/control-plane/src/pool/pool.ts

@@ -53,6 +53,7 @@ export async function adjust(event: PoolEvent): Promise<void> {
   // -1 disables the maximum check, matching the scale-up lambda's semantics. Defaults to unlimited
   // when unset so the pool keeps its previous behavior on stacks that do not provide the variable.
   const maximumRunners = parseInt(process.env.RUNNERS_MAXIMUM_COUNT || '-1');
+  const includeBusyRunners = yn(process.env.INCLUDE_BUSY_RUNNERS, { default: false });
 
   const { ghesApiUrl, ghesBaseUrl } = getGitHubEnterpriseApiUrl();
 
@@ -75,7 +76,7 @@ export async function adjust(event: PoolEvent): Promise<void> {
     statuses: ['running'],
   });
 
-  const numberOfRunnersInPool = calculatePooSize(ec2runners, runnerStatusses);
+  const numberOfRunnersInPool = calculatePooSize(ec2runners, runnerStatusses, includeBusyRunners);
   let topUp = event.poolSize - numberOfRunnersInPool;
 
   // The pool must never push the total number of runners (busy + idle) past the configured maximum.
@@ -146,12 +147,16 @@ async function getInstallationId(ghesApiUrl: string, org: string): Promise<numbe
   ).data.id;
 }
 
-function calculatePooSize(ec2runners: RunnerList[], runnerStatus: Map<string, RunnerStatus>): number {
+function calculatePooSize(
+  ec2runners: RunnerList[],
+  runnerStatus: Map<string, RunnerStatus>,
+  includeBusyRunners: boolean,
+): number {
   // Runner should be considered idle if it is still booting, or is idle in GitHub
   let numberOfRunnersInPool = 0;
   for (const ec2Instance of ec2runners) {
     if (
-      runnerStatus.get(ec2Instance.instanceId)?.busy === false &&
+      (runnerStatus.get(ec2Instance.instanceId)?.busy === false || includeBusyRunners) &&
       runnerStatus.get(ec2Instance.instanceId)?.status === 'online'
     ) {
       numberOfRunnersInPool++;

main.tf

@@ -280,6 +280,7 @@ module "runners" {
   pool_lambda_timeout                        = var.pool_lambda_timeout
   pool_runner_owner                          = var.pool_runner_owner
   pool_lambda_reserved_concurrent_executions = var.pool_lambda_reserved_concurrent_executions
+  pool_include_busy_runners                  = var.pool_include_busy_runners
 
   ssm_housekeeper = var.runners_ssm_housekeeper
   ebs_optimized   = var.runners_ebs_optimized

modules/runners/README.md

@@ -198,6 +198,7 @@ yarn run dist
 | <a name="input_parameter_store_tags"></a> [parameter\_store\_tags](#input\_parameter\_store\_tags) | Map of tags that will be added to all the SSM Parameter Store parameters created by the Lambda function. | `map(string)` | `{}` | no |
 | <a name="input_placement"></a> [placement](#input\_placement) | The placement options for the instance. See https://registry.terraform.io/providers/hashicorp/aws/latest/docs/resources/launch_template#placement for details. | <pre>object({<br/>    affinity                = optional(string)<br/>    availability_zone       = optional(string)<br/>    group_id                = optional(string)<br/>    group_name              = optional(string)<br/>    host_id                 = optional(string)<br/>    host_resource_group_arn = optional(string)<br/>    spread_domain           = optional(string)<br/>    tenancy                 = optional(string)<br/>    partition_number        = optional(number)<br/>  })</pre> | `null` | no |
 | <a name="input_pool_config"></a> [pool\_config](#input\_pool\_config) | The configuration for updating the pool. The `pool_size` to adjust to by the events triggered by the `schedule_expression`. For example you can configure a cron expression for week days to adjust the pool to 10 and another expression for the weekend to adjust the pool to 1. Use `schedule_expression_timezone ` to override the schedule time zone (defaults to UTC). | <pre>list(object({<br/>    schedule_expression          = string<br/>    schedule_expression_timezone = optional(string)<br/>    size                         = number<br/>  }))</pre> | `[]` | no |
+| <a name="input_pool_include_busy_runners"></a> [pool\_include\_busy\_runners](#input\_pool\_include\_busy\_runners) | Include busy runners in the pool calculation. By default busy runners are not included in the pool. | `bool` | `false` | no |
 | <a name="input_pool_lambda_memory_size"></a> [pool\_lambda\_memory\_size](#input\_pool\_lambda\_memory\_size) | Lambda Memory size limit in MB for pool lambda | `number` | `512` | no |
 | <a name="input_pool_lambda_reserved_concurrent_executions"></a> [pool\_lambda\_reserved\_concurrent\_executions](#input\_pool\_lambda\_reserved\_concurrent\_executions) | Amount of reserved concurrent executions for the scale-up lambda function. A value of 0 disables lambda from being triggered and -1 removes any concurrency limitations. | `number` | `1` | no |
 | <a name="input_pool_lambda_timeout"></a> [pool\_lambda\_timeout](#input\_pool\_lambda\_timeout) | Time out for the pool lambda in seconds. | `number` | `60` | no |

modules/runners/pool.tf

@@ -39,6 +39,7 @@ module "pool" {
       parameter_store_tags           = local.parameter_store_tags
     }
     pool                      = var.pool_config
+    include_busy_runners      = var.pool_include_busy_runners
     role_path                 = local.role_path
     role_permissions_boundary = var.role_permissions_boundary
     runner = {

modules/runners/pool/main.tf

@@ -60,6 +60,7 @@ resource "aws_lambda_function" "pool" {
       SSM_PARAMETER_STORE_TAGS                 = var.config.lambda.parameter_store_tags
       SCALE_ERRORS                             = jsonencode(var.config.runner.scale_errors)
       USE_DEDICATED_HOST                       = var.config.runner.use_dedicated_host
+      INCLUDE_BUSY_RUNNERS                     = var.config.include_busy_runners
     }
   }
 

modules/runners/pool/variables.tf

@@ -60,6 +60,7 @@ variable "config" {
       schedule_expression_timezone = string
       size                         = number
     }))
+    include_busy_runners                 = bool
     role_permissions_boundary            = string
     kms_key_arn                          = string
     ami_kms_key_arn                      = string

modules/runners/variables.tf

@@ -606,6 +606,12 @@ variable "pool_config" {
   default = []
 }
 
+variable "pool_include_busy_runners" {
+  description = "Include busy runners in the pool calculation. By default busy runners are not included in the pool."
+  type        = bool
+  default     = false
+}
+
 variable "disable_runner_autoupdate" {
   description = "Disable the auto update of the github runner agent. Be aware there is a grace period of 30 days, see also the [GitHub article](https://github.blog/changelog/2022-02-01-github-actions-self-hosted-runners-can-now-disable-automatic-updates/)"
   type        = bool

variables.tf

@@ -847,6 +847,12 @@ variable "pool_config" {
   default = []
 }
 
+variable "pool_include_busy_runners" {
+  description = "Include busy runners in the pool calculation. By default busy runners are not included in the pool."
+  type        = bool
+  default     = false
+}
+
 variable "aws_partition" {
   description = "(optiona) partition in the arn namespace to use if not 'aws'"
   type        = string