feat: implement strict multi-zone pod distribution for StatefulSets

Author: Abhishek Kumar

docs/CONFIG-VARS.md

@@ -332,6 +332,43 @@ Notes:
     - For example, defining `V4_CFG_VIYA_STOP_SCHEDULE` and not `V4_CFG_VIYA_START_SCHEDULE` will result in a Viya stop job that runs on a schedule and a suspended Viya start job that you will be able to manually trigger.
   - Defining both `V4_CFG_VIYA_START_SCHEDULE` and `V4_CFG_VIYA_STOP_SCHEDULE` will result in a non-suspended Viya start and stop job that runs on the schedule you defined.
 
+## Multi-Zone Pod Distribution
+
+| Name | Description | Type | Default | Required | Notes | Tasks |
+| :--- | ---: | ---: | ---: | ---: | ---: | ---: |
+| V4_CFG_MULTI_ZONE_ENABLED | Enable strict multi-zone pod distribution and anti-affinity for StatefulSets | bool | true | false | Adds restrictive topology spread constraints (maxSkew: 0) and required pod anti-affinity to prevent StatefulSet quorum loss during zone failures | viya |
+| V4_CFG_MULTI_ZONE_RABBITMQ_ENABLED | Enable strict multi-zone distribution for RabbitMQ StatefulSet | bool | true | false | Ensures RabbitMQ pods are strictly distributed across zones with required anti-affinity to maintain quorum during zone failures | viya |
+| V4_CFG_MULTI_ZONE_POSTGRES_ENABLED | Enable strict multi-zone distribution for PostgreSQL StatefulSet | bool | true | false | Ensures PostgreSQL pods are strictly distributed across zones for high availability. Only applies to internal PostgreSQL deployments | viya |
+| V4_CFG_STATEFUL_NODEPOOL_RESTRICTION | Restrict StatefulSets to dedicated stateful nodepools | bool | true | false | Adds node affinity to ensure StatefulSets only run on nodes labeled with workload.sas.com/class=stateful | viya |
+| V4_CFG_STATEFUL_NODEPOOL_LABEL | Label key used to identify stateful workload nodepools | string | workload.sas.com/class | false | Node label used for nodepool restriction. Stateful pods will require this label with value 'stateful' | viya |
+| V4_CFG_MULTI_ZONE_AUTO_DETECT | Automatically detect if cluster has multiple zones | bool | true | false | When enabled, automatically detects cluster topology and applies appropriate constraints. Prevents scheduling issues in single-zone clusters | viya |
+| V4_CFG_SINGLE_ZONE_FALLBACK | Enable relaxed constraints for single-zone clusters | bool | true | false | Applies relaxed topology constraints and preferred (not required) anti-affinity in single-zone deployments to avoid scheduling failures | viya |
+
+**Notes:**
+
+**Multi-Zone Clusters** (2+ zones detected):
+- **Zone Distribution**: `maxSkew: 0` with `DoNotSchedule` - pods must be evenly distributed across zones
+- **Required Pod Anti-Affinity**: Pods cannot be scheduled in the same zone as another pod of the same StatefulSet
+- **Nodepool Restriction**: StatefulSets are restricted to nodes with `workload.sas.com/class=stateful` label
+- **Node Distribution**: `maxSkew: 1` with `DoNotSchedule` - prevents multiple pods on the same node
+
+**Single-Zone Clusters** (1 zone detected):
+- **Node Distribution**: `maxSkew: 1` with `ScheduleAnyway` - spreads pods across nodes when possible
+- **Preferred Pod Anti-Affinity**: Attempts to avoid co-location on same node (weight: 100)
+- **Preferred Node Affinity**: Prefers stateful nodepool but allows scheduling elsewhere if needed
+
+**Automatic Behavior**:
+- Detects cluster zones by querying node labels (`topology.kubernetes.io/zone`)
+- Applies strict constraints only in true multi-zone environments  
+- Falls back to relaxed constraints in single-zone clusters
+- Ensures compatibility across different infrastructure deployments
+
+This configuration ensures:
+- No StatefulSet quorum loss during zone failures in multi-zone clusters
+- No scheduling failures in single-zone deployments
+- Optimal resource distribution based on cluster topology
+- Supports AKS, EKS, and GKE clusters
+
 ## Third-Party Tools
 
 ### Cert-manager

docs/user/MultiZoneDistribution.md

@@ -0,0 +1,69 @@
+# Strict Multi-Zone Pod Distribution - Implementation Guide
+
+## Overview
+This implementation provides strict multi-zone pod distribution for StatefulSets in AKS, EKS, and GKE clusters to prevent quorum loss during zone failures.
+
+## Features Added
+
+### 1. Configuration Variables (roles/vdm/defaults/main.yaml)
+- `V4_CFG_MULTI_ZONE_ENABLED`: Master switch for strict multi-zone distribution (default: true)
+- `V4_CFG_MULTI_ZONE_RABBITMQ_ENABLED`: RabbitMQ-specific distribution (default: true)  
+- `V4_CFG_MULTI_ZONE_POSTGRES_ENABLED`: PostgreSQL-specific distribution (default: true)
+- `V4_CFG_STATEFUL_NODEPOOL_RESTRICTION`: Restrict StatefulSets to dedicated nodepools (default: true)
+- `V4_CFG_STATEFUL_NODEPOOL_LABEL`: Label for stateful nodepool identification (default: workload.sas.com/class)
+
+### 2. Transformers Created
+- `rabbitmq-zone-distribution.yaml`: RabbitMQ strict zone distribution and nodepool restriction
+- `postgres-zone-distribution.yaml`: PostgreSQL strict zone distribution and nodepool restriction
+- `multi-zone-pod-distribution.yaml`: General StatefulSet strict distribution rules
+
+### 3. Restrictive Implementation Details
+
+#### Strict Topology Spread Constraints
+- **Zone Distribution**: `maxSkew: 0` on `topology.kubernetes.io/zone` with `DoNotSchedule`
+  - Ensures perfectly even distribution across zones
+  - Prevents scheduling if it would create imbalance
+- **Node Distribution**: `maxSkew: 1` on `kubernetes.io/hostname` with `DoNotSchedule`
+  - Prevents multiple pods on same node
+
+#### Required Pod Anti-Affinity Rules
+- **Zone-level**: Required anti-affinity to prevent pods in same zone
+- **Node-level**: Preferred anti-affinity to avoid same node (weight: 100)
+
+#### Nodepool Restrictions
+- **Required Node Affinity**: Stateful workloads must run on nodes with:
+  - `workload.sas.com/class: stateful` label
+  - `agentpool` label (ensures managed nodepool)
+
+### 4. Acceptance Criteria Compliance
+
+**More Restrictive Pod Topology Constraints**:
+- `maxSkew: 0` for zone distribution (most restrictive)
+- `DoNotSchedule` for both zone and node constraints
+- Required anti-affinity rules
+
+**Restrict StatefulSets to One Nodepool**:
+- Node affinity requires `workload.sas.com/class=stateful`
+- Ensures StatefulSets run only on designated stateful nodepool
+
+### 5. Benefits
+- **Zero Quorum Loss Risk**: Strict zone distribution prevents cluster failures
+- **Dedicated Resources**: Stateful workloads isolated to specific nodepool
+- **Predictable Scheduling**: Clear constraints for placement decisions
+- **Multi-Cloud Support**: Works with AKS, EKS, and GKE
+
+## Usage
+
+Enable in your ansible-vars.yaml:
+```yaml
+V4_CFG_MULTI_ZONE_ENABLED: true
+V4_CFG_MULTI_ZONE_RABBITMQ_ENABLED: true
+V4_CFG_MULTI_ZONE_POSTGRES_ENABLED: true
+V4_CFG_STATEFUL_NODEPOOL_RESTRICTION: true
+```
+
+## Nodepool Requirements
+Ensure your stateful nodepool is labeled:
+```bash
+kubectl label nodes <stateful-node> workload.sas.com/class=stateful
+```

roles/vdm/defaults/main.yaml

@@ -114,3 +114,12 @@ V4_WORKLOAD_ORCHESTRATOR_ENABLED: true
 
 ## NIST Features
 V4_CFG_NIST_FEATURES_ENABLED: false
+
+## Multi-Zone Pod Distribution
+V4_CFG_MULTI_ZONE_ENABLED: true
+V4_CFG_MULTI_ZONE_RABBITMQ_ENABLED: true
+V4_CFG_MULTI_ZONE_POSTGRES_ENABLED: true
+V4_CFG_STATEFUL_NODEPOOL_RESTRICTION: true
+V4_CFG_STATEFUL_NODEPOOL_LABEL: "workload.sas.com/class"
+V4_CFG_MULTI_ZONE_AUTO_DETECT: true
+V4_CFG_SINGLE_ZONE_FALLBACK: true

roles/vdm/tasks/main.yaml

@@ -213,6 +213,14 @@
     - uninstall
     - update
 
+# Include Multi-Zone Pod Distribution configuration
+- name: Include Multi-Zone Distribution
+  include_tasks: multi_zone_distribution.yaml
+  tags:
+    - install
+    - uninstall
+    - update
+
 # Include Sizing configuration and resources
 - name: Include Sizing
   include_tasks: sizing.yaml

roles/vdm/tasks/multi_zone_distribution.yaml

@@ -0,0 +1,130 @@
+# Copyright © 2020-2025, SAS Institute Inc., Cary, NC, USA. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+---
+# This file contains tasks for configuring multi-zone pod distribution and anti-affinity rules
+# to prevent StatefulSet quorum loss during zone failures in multi-zone clusters.
+# For single-zone deployments, applies relaxed constraints to avoid scheduling issues.
+
+# Detect if cluster has multiple zones
+- name: Multi-Zone - Detect cluster zones
+  kubernetes.core.k8s_info:
+    api_version: v1
+    kind: Node
+    kubeconfig: "{{ KUBECONFIG }}"
+  register: cluster_nodes
+  when:
+    - V4_CFG_MULTI_ZONE_ENABLED | bool
+    - V4_CFG_MULTI_ZONE_AUTO_DETECT | bool
+  tags:
+    - install
+    - uninstall
+    - update
+
+# Set fact for zone count
+- name: Multi-Zone - Calculate zone count
+  set_fact:
+    cluster_zone_count: "{{ cluster_nodes.resources | map(attribute='metadata.labels') | map('dict2items') | flatten | selectattr('key', 'equalto', 'topology.kubernetes.io/zone') | map(attribute='value') | unique | length }}"
+    is_multi_zone: "{{ (cluster_nodes.resources | map(attribute='metadata.labels') | map('dict2items') | flatten | selectattr('key', 'equalto', 'topology.kubernetes.io/zone') | map(attribute='value') | unique | length | int) > 1 }}"
+  when:
+    - V4_CFG_MULTI_ZONE_ENABLED | bool
+    - V4_CFG_MULTI_ZONE_AUTO_DETECT | bool
+    - cluster_nodes.resources is defined
+  tags:
+    - install
+    - uninstall
+    - update
+
+# Add multi-zone distribution overlay for RabbitMQ StatefulSet (multi-zone clusters)
+- name: Multi-Zone - RabbitMQ zone distribution with restrictive constraints
+  overlay_facts:
+    cadence_name: "{{ V4_CFG_CADENCE_NAME }}"
+    cadence_number: "{{ V4_CFG_CADENCE_VERSION }}"
+    existing: "{{ vdm_overlays }}"
+    add:
+      - { transformers: rabbitmq-zone-distribution.yaml, vdm: true }
+  when:
+    - V4_CFG_MULTI_ZONE_ENABLED | bool
+    - V4_CFG_MULTI_ZONE_RABBITMQ_ENABLED | bool
+    - PROVIDER in ["AKS", "azure", "EKS", "aws", "GKE", "gcp"]
+    - not V4_CFG_MULTI_ZONE_AUTO_DETECT | bool or (is_multi_zone | default(true) | bool)
+  tags:
+    - install
+    - uninstall
+    - update
+
+# Add relaxed distribution overlay for RabbitMQ StatefulSet (single-zone clusters)
+- name: Multi-Zone - RabbitMQ single-zone distribution fallback
+  overlay_facts:
+    cadence_name: "{{ V4_CFG_CADENCE_NAME }}"
+    cadence_number: "{{ V4_CFG_CADENCE_VERSION }}"
+    existing: "{{ vdm_overlays }}"
+    add:
+      - { transformers: rabbitmq-single-zone-distribution.yaml, vdm: true }
+  when:
+    - V4_CFG_MULTI_ZONE_ENABLED | bool
+    - V4_CFG_MULTI_ZONE_RABBITMQ_ENABLED | bool
+    - V4_CFG_SINGLE_ZONE_FALLBACK | bool
+    - PROVIDER in ["AKS", "azure", "EKS", "aws", "GKE", "gcp"]
+    - V4_CFG_MULTI_ZONE_AUTO_DETECT | bool and not (is_multi_zone | default(true) | bool)
+  tags:
+    - install
+    - uninstall
+    - update
+
+# Add multi-zone distribution overlay for PostgreSQL StatefulSet (multi-zone clusters)
+- name: Multi-Zone - PostgreSQL zone distribution with restrictive constraints
+  overlay_facts:
+    cadence_name: "{{ V4_CFG_CADENCE_NAME }}"
+    cadence_number: "{{ V4_CFG_CADENCE_VERSION }}"
+    existing: "{{ vdm_overlays }}"
+    add:
+      - { transformers: postgres-zone-distribution.yaml, vdm: true }
+  when:
+    - V4_CFG_MULTI_ZONE_ENABLED | bool
+    - V4_CFG_MULTI_ZONE_POSTGRES_ENABLED | bool
+    - V4_CFG_POSTGRES_SERVERS.default.internal | bool
+    - PROVIDER in ["AKS", "azure", "EKS", "aws", "GKE", "gcp"]
+    - not V4_CFG_MULTI_ZONE_AUTO_DETECT | bool or (is_multi_zone | default(true) | bool)
+  tags:
+    - install
+    - uninstall
+    - update
+
+# Add relaxed distribution overlay for PostgreSQL StatefulSet (single-zone clusters)
+- name: Multi-Zone - PostgreSQL single-zone distribution fallback
+  overlay_facts:
+    cadence_name: "{{ V4_CFG_CADENCE_NAME }}"
+    cadence_number: "{{ V4_CFG_CADENCE_VERSION }}"
+    existing: "{{ vdm_overlays }}"
+    add:
+      - { transformers: postgres-single-zone-distribution.yaml, vdm: true }
+  when:
+    - V4_CFG_MULTI_ZONE_ENABLED | bool
+    - V4_CFG_MULTI_ZONE_POSTGRES_ENABLED | bool
+    - V4_CFG_POSTGRES_SERVERS.default.internal | bool
+    - V4_CFG_SINGLE_ZONE_FALLBACK | bool
+    - PROVIDER in ["AKS", "azure", "EKS", "aws", "GKE", "gcp"]
+    - V4_CFG_MULTI_ZONE_AUTO_DETECT | bool and not (is_multi_zone | default(true) | bool)
+  tags:
+    - install
+    - uninstall
+    - update
+
+# Add general multi-zone distribution overlay for other StatefulSets (multi-zone clusters)
+- name: Multi-Zone - General StatefulSet zone distribution with restrictive constraints
+  overlay_facts:
+    cadence_name: "{{ V4_CFG_CADENCE_NAME }}"
+    cadence_number: "{{ V4_CFG_CADENCE_VERSION }}"
+    existing: "{{ vdm_overlays }}"
+    add:
+      - { transformers: multi-zone-pod-distribution.yaml, vdm: true }
+  when:
+    - V4_CFG_MULTI_ZONE_ENABLED | bool
+    - V4_CFG_STATEFUL_NODEPOOL_RESTRICTION | bool
+    - PROVIDER in ["AKS", "azure", "EKS", "aws", "GKE", "gcp"]
+    - not V4_CFG_MULTI_ZONE_AUTO_DETECT | bool or (is_multi_zone | default(true) | bool)
+  tags:
+    - install
+    - uninstall
+    - update

roles/vdm/templates/transformers/multi-zone-pod-distribution.yaml

@@ -0,0 +1,64 @@
+# Copyright © 2020-2025, SAS Institute Inc., Cary, NC, USA. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Transformer to add restrictive pod topology spread constraints and anti-affinity rules 
+# for StatefulSets to ensure strict zone distribution and quorum protection
+apiVersion: builtin
+kind: PatchTransformer
+metadata:
+  name: multi-zone-pod-distribution
+patch: |-
+  - op: add
+    path: /spec/template/spec/topologySpreadConstraints
+    value:
+      - maxSkew: 0
+        topologyKey: topology.kubernetes.io/zone
+        whenUnsatisfiable: DoNotSchedule
+        labelSelector:
+          matchLabels:
+            app.kubernetes.io/name: PLACEHOLDER_APP_NAME
+      - maxSkew: 1
+        topologyKey: kubernetes.io/hostname
+        whenUnsatisfiable: DoNotSchedule
+        labelSelector:
+          matchLabels:
+            app.kubernetes.io/name: PLACEHOLDER_APP_NAME
+  - op: add
+    path: /spec/template/spec/affinity
+    value:
+      nodeAffinity:
+        requiredDuringSchedulingIgnoredDuringExecution:
+          nodeSelectorTerms:
+            - matchExpressions:
+                - key: workload.sas.com/class
+                  operator: In
+                  values:
+                    - stateful
+                - key: agentpool
+                  operator: Exists
+      podAntiAffinity:
+        requiredDuringSchedulingIgnoredDuringExecution:
+          - labelSelector:
+              matchLabels:
+                app.kubernetes.io/name: PLACEHOLDER_APP_NAME
+            topologyKey: topology.kubernetes.io/zone
+        preferredDuringSchedulingIgnoredDuringExecution:
+          - weight: 100
+            podAffinityTerm:
+              labelSelector:
+                matchLabels:
+                  app.kubernetes.io/name: PLACEHOLDER_APP_NAME
+              topologyKey: kubernetes.io/hostname
+target:
+  kind: StatefulSet
+  name: ".*rabbitmq.*|.*postgres.*|.*cas.*"
+  version: v1
+replacements:
+  - source:
+      kind: StatefulSet
+      fieldPath: metadata.labels['app.kubernetes.io/name']
+    targets:
+      - select:
+          kind: PatchTransformer
+        fieldPaths:
+          - patch

roles/vdm/templates/transformers/postgres-single-zone-distribution.yaml

@@ -0,0 +1,42 @@
+# Copyright © 2020-2025, SAS Institute Inc., Cary, NC, USA. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Transformer for PostgreSQL StatefulSet in single-zone clusters with relaxed constraints
+apiVersion: builtin
+kind: PatchTransformer
+metadata:
+  name: postgres-single-zone-distribution
+patch: |-
+  - op: add
+    path: /spec/template/spec/topologySpreadConstraints
+    value:
+      - maxSkew: 1
+        topologyKey: kubernetes.io/hostname
+        whenUnsatisfiable: ScheduleAnyway
+        labelSelector:
+          matchLabels:
+            postgres-operator.crunchydata.com/cluster: shared-postgres
+  - op: add
+    path: /spec/template/spec/affinity
+    value:
+      nodeAffinity:
+        preferredDuringSchedulingIgnoredDuringExecution:
+          - weight: 100
+            preference:
+              matchExpressions:
+                - key: workload.sas.com/class
+                  operator: In
+                  values:
+                    - stateful
+      podAntiAffinity:
+        preferredDuringSchedulingIgnoredDuringExecution:
+          - weight: 100
+            podAffinityTerm:
+              labelSelector:
+                matchLabels:
+                  postgres-operator.crunchydata.com/cluster: shared-postgres
+              topologyKey: kubernetes.io/hostname
+target:
+  kind: StatefulSet
+  name: ".*postgres.*"
+  version: v1