From ae94d1f5492d79d2f6e33e37fa508a6b380166b4 Mon Sep 17 00:00:00 2001
From: Miguel Harmant <miguel.harmant@dominodatalab.com>
Date: Tue, 12 Nov 2024 11:19:50 -0500
Subject: [PATCH] Adds capacity reservation module

---
 modules/capacity-reservation/main.tf      | 71 +++++++++++++++++++++++
 modules/capacity-reservation/outputs.tf   |  0
 modules/capacity-reservation/variables.tf | 21 +++++++
 3 files changed, 92 insertions(+)
 create mode 100644 modules/capacity-reservation/main.tf
 create mode 100644 modules/capacity-reservation/outputs.tf
 create mode 100644 modules/capacity-reservation/variables.tf

diff --git a/modules/capacity-reservation/main.tf b/modules/capacity-reservation/main.tf
new file mode 100644
index 00000000..0ec6b351
--- /dev/null
+++ b/modules/capacity-reservation/main.tf
@@ -0,0 +1,71 @@
+locals {
+  capacity_reservations = merge(flatten([
+    for ng_name, ng in var.node_groups : [
+      for az_id in ng.availability_zone_ids : {
+        for type in ng.instance_types : "${type}-${az_id}" => {
+          capacity      = ng.capacity
+          instance_type = type
+          az            = data.aws_availability_zone.zone_ids[az_id].name
+        }
+      }
+    ]
+  ])...)
+  zone_ids = toset(flatten([for ng in var.node_groups :
+    ng.availability_zone_ids
+  ]))
+
+  instance_types = toset([for cr in local.capacity_reservations : cr.instance_type])
+}
+
+data "aws_availability_zone" "zone_ids" {
+  for_each = local.zone_ids
+  zone_id  = each.key
+}
+
+
+data "aws_ec2_instance_type_offerings" "this" {
+  for_each = local.instance_types
+
+  filter {
+    name   = "instance-type"
+    values = [each.key]
+  }
+
+  location_type = "availability-zone"
+}
+
+
+resource "aws_ec2_capacity_reservation" "this" {
+  for_each          = local.capacity_reservations
+  instance_type     = each.value.instance_type
+  instance_platform = "Linux/UNIX"
+  availability_zone = each.value.az
+  instance_count    = each.value.capacity
+
+  lifecycle {
+    precondition {
+      condition     = contains(data.aws_ec2_instance_type_offerings.this[each.value.instance_type].locations, each.value.az)
+      error_message = <<-EOM
+        Instance type ${each.value.instance_type} is NOT available in availability_zone ${each.value.az}.
+        available = ${jsonencode(data.aws_ec2_instance_type_offerings.this[each.value.instance_type].locations)}
+      EOM
+    }
+  }
+}
+
+resource "terraform_data" "describe_capacity_reservation" {
+  for_each = aws_ec2_capacity_reservation.this
+
+  provisioner "local-exec" {
+    command     = <<EOT
+      aws ec2 describe-capacity-reservations --capacity-reservation-ids ${each.value.id} --region ${var.region}
+    EOT
+    interpreter = ["/bin/bash", "-c"]
+  }
+
+  depends_on = [aws_ec2_capacity_reservation.this]
+}
+
+output "capacity_reservation" {
+  value = aws_ec2_capacity_reservation.this
+}
diff --git a/modules/capacity-reservation/outputs.tf b/modules/capacity-reservation/outputs.tf
new file mode 100644
index 00000000..e69de29b
diff --git a/modules/capacity-reservation/variables.tf b/modules/capacity-reservation/variables.tf
new file mode 100644
index 00000000..b3e94d10
--- /dev/null
+++ b/modules/capacity-reservation/variables.tf
@@ -0,0 +1,21 @@
+
+variable "node_groups" {
+  description = "Additional EKS managed node groups definition."
+  type = map(object({
+    instance_types        = list(string)
+    spot                  = optional(bool, false)
+    capacity              = number
+    availability_zone_ids = list(string)
+  }))
+  default = {}
+}
+
+variable "region" {
+  type        = string
+  description = "AWS region for the deployment"
+  nullable    = false
+  validation {
+    condition     = can(regex("(us(-gov)?|ap|ca|cn|eu|sa|me|af)-(central|(north|south)?(east|west)?)-[0-9]", var.region))
+    error_message = "The provided region must follow the format of AWS region names, e.g., us-west-2, us-gov-west-1."
+  }
+}