From c975c771e9358fd73a7e6b0f97324ac37266ebc8 Mon Sep 17 00:00:00 2001
From: Falk Scheerschmidt <falk@staffbase.com>
Date: Mon, 16 Oct 2023 15:31:49 +0200
Subject: [PATCH] Support rolling_upgrade_policy for vm module

---
 modules/vm/main.tf      | 13 +++++++++++++
 modules/vm/variables.tf | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/modules/vm/main.tf b/modules/vm/main.tf
index d754ac03..04672a31 100644
--- a/modules/vm/main.tf
+++ b/modules/vm/main.tf
@@ -59,6 +59,19 @@ resource "azurerm_linux_virtual_machine_scale_set" "tfe_vmss" {
     force_deletion_enabled = var.vm_vmss_scale_in_force_deletion_enabled
   }
 
+  dynamic "rolling_upgrade_policy" {
+    for_each = var.vm_upgrade_mode == "Rolling" ? [1] : []
+
+    content {
+      cross_zone_upgrades_enabled             = var.vm_cross_zone_upgrades_enabled
+      max_batch_instance_percent              = var.vm_max_batch_instance_percent
+      max_unhealthy_instance_percent          = var.vm_max_unhealthy_instance_percent
+      max_unhealthy_upgraded_instance_percent = var.vm_max_unhealthy_upgraded_instance_percent
+      pause_time_between_batches              = var.vm_pause_time_between_batches
+      prioritize_unhealthy_instances_enabled  = var.vm_prioritize_unhealthy_instances_enabled
+    }
+  }
+
   # Source image reference will be used if vm_image_id is 'ubuntu' or 'rhel'
   dynamic "source_image_reference" {
     for_each = var.vm_image_id == "ubuntu" || var.vm_image_id == "rhel" ? [1] : []
diff --git a/modules/vm/variables.tf b/modules/vm/variables.tf
index ffce17ab..75d52add 100644
--- a/modules/vm/variables.tf
+++ b/modules/vm/variables.tf
@@ -147,6 +147,42 @@ variable "vm_vmss_scale_in_force_deletion_enabled" {
   description = "Should the virtual machines chosen for removal be force deleted when the virtual machine scale set is being scaled-in?"
 }
 
+variable "vm_cross_zone_upgrades_enabled" {
+  default     = false
+  type        = bool
+  description = "Should the Virtual Machine Scale Set be upgraded across zones in the same region when a Virtual Machine Image is updated?"
+}
+
+variable "vm_max_batch_instance_percent" {
+  default     = 20
+  type        = number
+  description = "The maximum percent of total virtual machine instances that will be upgraded simultaneously by the rolling upgrade in one batch. As this is a maximum, unhealthy instances in previous or future batches can cause the percentage of instances in a batch to decrease to ensure higher reliability."
+}
+
+variable "vm_max_unhealthy_instance_percent" {
+  default     = 20
+  type        = number
+  description = "The maximum percentage of the total virtual machine instances in the scale set that can be simultaneously unhealthy, either as a result of being upgraded, or by being found in an unhealthy state by the virtual machine health checks before the rolling upgrade aborts. This constraint will be checked prior to starting any batch."
+}
+
+variable "vm_max_unhealthy_upgraded_instance_percent" {
+  default     = 20
+  type        = number
+  description = "The maximum percentage of upgraded virtual machine instances that can be found to be in an unhealthy state. This check will happen after each batch is upgraded. If this percentage is ever exceeded, the rolling update aborts."
+}
+
+variable "vm_pause_time_between_batches" {
+  default     = "PT0S"
+  type        = string
+  description = "The wait time between completing the update for all virtual machines in one batch and starting the next batch. The time duration should be specified in ISO 8601 format."
+}
+
+variable "vm_prioritize_unhealthy_instances_enabled" {
+  default     = false
+  type        = bool
+  description = "Upgrade all unhealthy instances in a scale set before any healthy instances. Possible values are true or false."
+}
+
 variable "ca_certificate_secret" {
   type = object({
     key_vault_id = string