Improve spot HA by utilising ASG capacity rebalance

README.md

@@ -49,6 +49,7 @@ module "fck-nat" {
 | Name | Type |
 |------|------|
 | [aws_autoscaling_group.main](https://registry.terraform.io/providers/hashicorp/aws/latest/docs/resources/autoscaling_group) | resource |
+| [aws_autoscaling_lifecycle_hook.spot_termination_wait](https://registry.terraform.io/providers/hashicorp/aws/latest/docs/resources/autoscaling_lifecycle_hook) | resource |
 | [aws_iam_instance_profile.main](https://registry.terraform.io/providers/hashicorp/aws/latest/docs/resources/iam_instance_profile) | resource |
 | [aws_iam_role.main](https://registry.terraform.io/providers/hashicorp/aws/latest/docs/resources/iam_role) | resource |
 | [aws_instance.main](https://registry.terraform.io/providers/hashicorp/aws/latest/docs/resources/instance) | resource |
@@ -76,6 +77,7 @@ module "fck-nat" {
 | <a name="input_ebs_root_volume_size"></a> [ebs\_root\_volume\_size](#input\_ebs\_root\_volume\_size) | Size of the EBS root volume in GB | `number` | `8` | no |
 | <a name="input_eip_allocation_ids"></a> [eip\_allocation\_ids](#input\_eip\_allocation\_ids) | EIP allocation IDs to use for the NAT instance. Automatically assign a public IP if none is provided. Note: Currently only supports at most one EIP allocation. | `list(string)` | `[]` | no |
 | <a name="input_encryption"></a> [encryption](#input\_encryption) | Whether or not to encrypt the EBS volume | `bool` | `true` | no |
+| <a name="input_ha_additional_instance_types"></a> [ha\_additional\_instance\_types](#input\_ha\_additional\_instance\_types) | Additional instance types used by autoscaling rebalancing when the primary instance is unavailable | `list(string)` | <pre>[<br>  "t4g.small"<br>]</pre> | no |
 | <a name="input_ha_mode"></a> [ha\_mode](#input\_ha\_mode) | Whether or not high-availability mode should be enabled via autoscaling group | `bool` | `true` | no |
 | <a name="input_instance_type"></a> [instance\_type](#input\_instance\_type) | Instance type to use for the NAT instance | `string` | `"t4g.micro"` | no |
 | <a name="input_kms_key_id"></a> [kms\_key\_id](#input\_kms\_key\_id) | Will use the provided KMS key ID to encrypt the EBS volume. Uses the default KMS key if none provided | `string` | `null` | no |
@@ -116,4 +118,4 @@ module "fck-nat" {
 | <a name="output_security_group_id"></a> [security\_group\_id](#output\_security\_group\_id) | Deprecated. The ID of the security group used by fck-nat ENIs |
 | <a name="output_security_group_ids"></a> [security\_group\_ids](#output\_security\_group\_ids) | List of security group IDs used by fck-nat ENIs |
 | <a name="output_subnet_id"></a> [subnet\_id](#output\_subnet\_id) | Subnet ID to which the fck-nat instance is deployed into |
-| <a name="output_vpc_id"></a> [vpc\_id](#output\_vpc\_id) | VPC ID to which the fck-nat instance is deployed into |
+| <a name="output_vpc_id"></a> [vpc\_id](#output\_vpc\_id) | VPC ID to which the fck-nat instance is deployed into |

asg.tf

@@ -8,9 +8,31 @@ resource "aws_autoscaling_group" "main" {
   health_check_type   = "EC2"
   vpc_zone_identifier = [var.subnet_id]
 
-  launch_template {
-    id      = aws_launch_template.main.id
-    version = "$Latest"
+  capacity_rebalance = var.use_spot_instances
+
+  mixed_instances_policy {
+    instances_distribution {
+      on_demand_percentage_above_base_capacity = var.use_spot_instances ? 0 : 100
+      spot_allocation_strategy                 = "price-capacity-optimized"
+    }
+    launch_template {
+      launch_template_specification {
+        launch_template_id = aws_launch_template.main.id
+        version            = aws_launch_template.main.latest_version
+      }
+
+      override {
+        instance_type = var.instance_type
+      }
+
+      dynamic "override" {
+        for_each = toset(var.ha_additional_instance_types)
+
+        content {
+          instance_type = override.value
+        }
+      }
+    }
   }
 
   dynamic "tag" {
@@ -60,3 +82,12 @@ resource "aws_autoscaling_group" "main" {
     delete = "15m"
   }
 }
+
+resource "aws_autoscaling_lifecycle_hook" "spot_termination_wait" {
+  count = var.ha_mode && var.use_spot_instances ? 1 : 0
+
+  name                   = "TerminationWait"
+  autoscaling_group_name = aws_autoscaling_group.main[0].name
+  heartbeat_timeout      = 300
+  lifecycle_transition   = "autoscaling:EC2_INSTANCE_TERMINATING"
+}

ec2.tf

@@ -61,7 +61,7 @@ resource "aws_launch_template" "main" {
   }
 
   dynamic "instance_market_options" {
-    for_each = var.use_spot_instances ? ["x"] : []
+    for_each = var.use_spot_instances && !var.ha_mode ? ["x"] : []
 
     content {
       market_type = "spot"
@@ -100,7 +100,7 @@ resource "aws_instance" "main" {
 
   launch_template {
     id      = aws_launch_template.main.id
-    version = "$Latest"
+    version = aws_launch_template.main.latest_version
   }
 
   tags = var.tags

variables.tf

@@ -55,6 +55,12 @@ variable "ha_mode" {
   default     = true
 }
 
+variable "ha_additional_instance_types" {
+  description = "Additional instance types used by autoscaling rebalancing when the primary instance is unavailable"
+  type        = list(string)
+  default     = ["t4g.small"]
+}
+
 variable "instance_type" {
   description = "Instance type to use for the NAT instance"
   type        = string