Modify autoscaler to work with self hosted workers (#9)

--------- Co-authored-by: Elie CHARRA <[email protected]> Co-authored-by: jubran nassar <[email protected]>
spacelift-io · Mar 18, 2024 · 5840d19 · 5840d19
2 parents d60d896 + 006088f
commit 5840d19
Show file tree

Hide file tree

Showing 8 changed files with 133 additions and 45 deletions.
diff --git a/iac/api_secret.tf b/iac/api_secret.tf
diff --git a/iac/aws.tf b/iac/aws.tf
@@ -1,3 +1,2 @@
 data "aws_region" "current" {}
 data "aws_partition" "current" {}
-data "aws_region" "current" {}
diff --git a/iac/download.sh b/iac/download.sh
@@ -3,10 +3,11 @@ set -ex
 
 # Download the data.
 code_version=$1
-local_path=$2
+code_architecture=$2
 
-curl -L -o $local_path https://github.com/spacelift-io/ec2-workerpool-autoscaler/releases/download/v${code_version}/ec2-workerpool-autoscaler_${code_version}_linux_amd64.zip
+curl -L -o lambda.zip "https://github.com/spacelift-io/ec2-workerpool-autoscaler/releases/download/${code_version}/ec2-workerpool-autoscaler_linux_${code_architecture}.zip"
 
-source_code_hash=$(openssl dgst -binary -sha256 $local_path | openssl base64 -A)
-
-echo "{\"source_code_hash\": \"${source_code_hash}\"}"
+mkdir -p lambda
+cd lambda
+unzip -o ../lambda.zip
+rm ../lambda.zip
diff --git a/iac/lambda.tf b/iac/lambda.tf
@@ -1,29 +1,57 @@
-data "external" "package" {
-  program = [
-    "${path.module}/download.sh",
-    local.code_version,
-    local.package_path,
-  ]
+locals {
+  base_name      = var.base_name == null ? "sp5ft-${var.worker_pool_id}" : var.base_name
+  function_name  = "${local.base_name}-ec2-autoscaler"
+  use_s3_package = var.autoscaler_s3_package != null
+}
+
+resource "aws_ssm_parameter" "spacelift_api_key_secret" {
+  name  = "/${local.function_name}/spacelift-api-secret-${var.worker_pool_id}"
+  type  = "SecureString"
+  value = var.spacelift_api_key_secret
+}
+
+resource "null_resource" "download" {
+  triggers = {
+    # Always re-download the archive file
+    now = timestamp()
+  }
+  provisioner "local-exec" {
+    command = "${path.module}/download.sh ${var.autoscaler_version} ${var.autoscaler_architecture}"
+  }
+}
+
+data "archive_file" "binary" {
+  type        = "zip"
+  source_file = "lambda/bootstrap"
+  output_path = "ec2-workerpool-autoscaler_${var.autoscaler_version}.zip"
+  depends_on  = [null_resource.download]
 }
 
 resource "aws_lambda_function" "autoscaler" {
-  filename         = local.package_path
-  source_code_hash = data.external.package.result.source_code_hash
-  function_name    = "ec2-autoscaler-${var.worker_pool_id}"
-  role             = aws_iam_role.lambda.arn
-  handler          = "ec2-workerpool-autoscaler_v${var.autoscaler_version}"
+  filename         = !local.use_s3_package ? data.archive_file.binary.output_path : null
+  source_code_hash = !local.use_s3_package ? data.archive_file.binary.output_base64sha256 : null
 
-  reserved_concurrent_executions = 1
-  runtime                        = "go1.x"
+  s3_bucket = local.use_s3_package ? var.autoscaler_s3_package.bucket : null
+  s3_key    = local.use_s3_package ? var.autoscaler_s3_package.key : null
+  s3_object_version = local.use_s3_package ? var.autoscaler_s3_package.object_version : null
+
+  function_name = local.function_name
+  role          = aws_iam_role.autoscaler.arn
+  handler       = "bootstrap"
+  runtime       = "provided.al2"
+  architectures = [var.autoscaler_architecture == "amd64" ? "x86_64" : var.autoscaler_architecture]
+  timeout       = var.autoscaling_timeout
 
   environment {
     variables = {
       AUTOSCALING_GROUP_ARN         = var.autoscaling_group_arn
       AUTOSCALING_REGION            = data.aws_region.current.name
       SPACELIFT_API_KEY_ID          = var.spacelift_api_key_id
       SPACELIFT_API_KEY_SECRET_NAME = aws_ssm_parameter.spacelift_api_key_secret.name
-      SPACELIFT_API_KEY_ENDPOINT    = var.spacelift_url
+      SPACELIFT_API_KEY_ENDPOINT    = var.spacelift_api_key_endpoint
       SPACELIFT_WORKER_POOL_ID      = var.worker_pool_id
+      AUTOSCALING_MAX_CREATE        = var.autoscaling_max_create
+      AUTOSCALING_MAX_KILL          = var.autoscaling_max_terminate
     }
   }
 
@@ -33,20 +61,25 @@ resource "aws_lambda_function" "autoscaler" {
 }
 
 resource "aws_cloudwatch_event_rule" "scheduling" {
-  name                = "spacelift-${var.worker_pool_id}-scheduling"
+  name                = local.function_name
   description         = "Spacelift autoscaler scheduling for worker pool ${var.worker_pool_id}"
-  schedule_expression = "rate(${var.autoscaling_frequency} minutes)"
+  schedule_expression = var.schedule_expression
 }
 
 resource "aws_cloudwatch_event_target" "scheduling" {
   rule = aws_cloudwatch_event_rule.scheduling.name
   arn  = aws_lambda_function.autoscaler.arn
 }
 
-resource "aws_lambda_permission" "allow_cloudwatch_to_call_check_foo" {
+resource "aws_lambda_permission" "allow_cloudwatch_to_call_lambda" {
   statement_id  = "AllowExecutionFromCloudWatch"
   action        = "lambda:InvokeFunction"
   function_name = aws_lambda_function.autoscaler.function_name
   principal     = "events.amazonaws.com"
   source_arn    = aws_cloudwatch_event_rule.scheduling.arn
 }
+
+resource "aws_cloudwatch_log_group" "log_group" {
+  name              = "/aws/lambda/${local.function_name}"
+  retention_in_days = 7
+}
diff --git a/iac/logging.tf b/iac/logging.tf
diff --git a/iac/policy.tf b/iac/policy.tf
@@ -1,4 +1,4 @@
-data "aws_iam_policy_document" "lambda_policy" {
+data "aws_iam_policy_document" "autoscaler" {
   # Allow the Lambda to write CloudWatch Logs.
   statement {
     effect = "Allow"
@@ -7,7 +7,7 @@ data "aws_iam_policy_document" "lambda_policy" {
       "logs:PutLogEvents",
     ]
 
-    resources = ["${aws_cloudwatch_log_group.log_group.arn}/*"]
+    resources = ["${aws_cloudwatch_log_group.log_group.arn}:*"]
   }
 
   # Allow the Lambda to put X-Ray traces.
@@ -28,14 +28,9 @@ data "aws_iam_policy_document" "lambda_policy" {
     actions = [
       "autoscaling:DetachInstances",
       "autoscaling:SetDesiredCapacity",
+      "autoscaling:DescribeAutoScalingGroups",
     ]
 
-    resources = [var.autoscaling_group_arn]
-  }
-
-  statement {
-    effect    = "Allow"
-    actions   = ["autoscaling:DescribeAutoScalingGroups"]
     resources = ["*"]
   }
 
@@ -58,12 +53,24 @@ data "aws_iam_policy_document" "lambda_policy" {
   }
 }
 
-resource "aws_iam_role" "lambda" {
-  name               = "ec2-autoscaler-${var.worker_pool_id}"
-  assume_role_policy = data.aws_iam_policy_document.assume_lambda_role.json
+
+resource "aws_iam_role" "autoscaler" {
+  name = local.function_name
+  assume_role_policy = jsonencode({
+    Version = "2012-10-17"
+    Statement = [
+      {
+        "Effect" : "Allow",
+        "Principal" : {
+          "Service" : "lambda.amazonaws.com"
+        },
+        "Action" : "sts:AssumeRole"
+      },
+    ]
+  })
 
   inline_policy {
     name   = "ec2-autoscaler-${var.worker_pool_id}"
-    policy = data.aws_iam_policy_document.lambda_policy.json
+    policy = data.aws_iam_policy_document.autoscaler.json
   }
 }
diff --git a/iac/providers.tf b/iac/providers.tf
@@ -6,3 +6,7 @@ terraform {
     }
   }
 }
+
+provider "aws" {
+  region = var.region
+}
diff --git a/iac/variables.tf b/iac/variables.tf
@@ -6,6 +6,7 @@ variable "autoscaling_group_arn" {
 variable "autoscaler_version" {
   type        = string
   description = "Version of the autoscaler to deploy"
+  default     = "v0.3.0"
 }
 
 variable "autoscaling_frequency" {
@@ -25,12 +26,64 @@ variable "spacelift_api_key_secret" {
   description = "Secret corresponding to the Spacelift API key to use"
 }
 
-variable "spacelift_url" {
+variable "spacelift_api_key_endpoint" {
   type        = string
   description = "Full URL of the Spacelift API endpoint to use, eg. https://demo.app.spacelift.io"
+  default     = null
 }
 
 variable "worker_pool_id" {
   type        = string
   description = "ID of the Spacelift worker pool to autoscale"
 }
+
+variable "autoscaler_architecture" {
+  type        = string
+  description = "Instruction set architecture of the autoscaler to use"
+  default     = "amd64"
+}
+
+variable "autoscaling_timeout" {
+  type        = number
+  description = "Timeout (in seconds) for a single autoscaling run. The more instances you have, the higher this should be."
+  default     = 30
+}
+
+variable "autoscaling_max_create" {
+  description = "The maximum number of instances the utility is allowed to create in a single run"
+  type        = number
+  default     = 1
+}
+
+variable "autoscaling_max_terminate" {
+  description = "The maximum number of instances the utility is allowed to terminate in a single run"
+  type        = number
+  default     = 1
+}
+
+variable "schedule_expression" {
+  type        = string
+  description = "Autoscaler scheduling expression"
+  default     = "rate(1 minute)"
+}
+
+variable "base_name" {
+  type        = string
+  description = "Base name for resources. If unset, it defaults to `sp5ft-$${var.worker_pool_id}`."
+  nullable    = true
+  default     = null
+}
+
+variable "region" {
+  type        = string
+  description = "AWS Region where the provider will operate"
+}
+
+variable "autoscaler_s3_package" {
+  type = object({
+    bucket = string
+    key    = string
+    object_version = optional(string)
+  })
+  description = "Configuration to retrieve autoscaler lambda package from s3 bucket"
+}