From 82966207a832bb2bad652cc1025390434a8a6727 Mon Sep 17 00:00:00 2001
From: Alexander Berger <alex-berger@users.noreply.github.com>
Date: Tue, 19 Mar 2024 21:39:18 +0100
Subject: [PATCH] Adding support for Availability Zone Affinity (#3470)

Fixes https://github.com/kubernetes-sigs/aws-load-balancer-controller/issues/3431

Signed-off-by: Alex Berger <alex-berger@gmx.ch>
---
 docs/guide/service/annotations.md             |  4 +++
 pkg/service/model_build_load_balancer.go      | 28 +++++++++++++++----
 pkg/service/model_build_load_balancer_test.go |  6 +++-
 3 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/docs/guide/service/annotations.md b/docs/guide/service/annotations.md
index cdab12199..cd4e1a787 100644
--- a/docs/guide/service/annotations.md
+++ b/docs/guide/service/annotations.md
@@ -258,6 +258,10 @@ for proxy protocol v2 configuration.
         ```
         service.beta.kubernetes.io/aws-load-balancer-attributes: load_balancing.cross_zone.enabled=true
         ```
+        - enable client availability zone affinity
+        ```
+        service.beta.kubernetes.io/aws-load-balancer-attributes: dns_record.client_routing_policy=availability_zone_affinity
+        ```
 
 - <a name="deprecated-attributes"></a>the following annotations are deprecated in v2.3.0 release in favor of [service.beta.kubernetes.io/aws-load-balancer-attributes](#load-balancer-attributes)
 
diff --git a/pkg/service/model_build_load_balancer.go b/pkg/service/model_build_load_balancer.go
index 471f39d69..9e5040a52 100644
--- a/pkg/service/model_build_load_balancer.go
+++ b/pkg/service/model_build_load_balancer.go
@@ -26,12 +26,16 @@ import (
 )
 
 const (
-	lbAttrsAccessLogsS3Enabled           = "access_logs.s3.enabled"
-	lbAttrsAccessLogsS3Bucket            = "access_logs.s3.bucket"
-	lbAttrsAccessLogsS3Prefix            = "access_logs.s3.prefix"
-	lbAttrsLoadBalancingCrossZoneEnabled = "load_balancing.cross_zone.enabled"
-	resourceIDLoadBalancer               = "LoadBalancer"
-	minimalAvailableIPAddressCount       = int64(8)
+	lbAttrsAccessLogsS3Enabled                 = "access_logs.s3.enabled"
+	lbAttrsAccessLogsS3Bucket                  = "access_logs.s3.bucket"
+	lbAttrsAccessLogsS3Prefix                  = "access_logs.s3.prefix"
+	lbAttrsLoadBalancingCrossZoneEnabled       = "load_balancing.cross_zone.enabled"
+	lbAttrsLoadBalancingDnsClientRoutingPolicy = "dns_record.client_routing_policy"
+	availabilityZoneAffinity                   = "availability_zone_affinity"
+	partialAvailabilityZoneAffinity            = "partial_availability_zone_affinity"
+	anyAvailabilityZone                        = "any_availability_zone"
+	resourceIDLoadBalancer                     = "LoadBalancer"
+	minimalAvailableIPAddressCount             = int64(8)
 )
 
 func (t *defaultModelBuildTask) buildLoadBalancer(ctx context.Context, scheme elbv2model.LoadBalancerScheme) error {
@@ -437,6 +441,18 @@ func (t *defaultModelBuildTask) getLoadBalancerAttributes() (map[string]string,
 	if _, err := t.annotationParser.ParseStringMapAnnotation(annotations.SvcLBSuffixLoadBalancerAttributes, &attributes, t.service.Annotations); err != nil {
 		return nil, err
 	}
+	dnsRecordClientRoutingPolicy, exists := attributes[lbAttrsLoadBalancingDnsClientRoutingPolicy]
+	if exists {
+		switch dnsRecordClientRoutingPolicy {
+		case availabilityZoneAffinity:
+		case partialAvailabilityZoneAffinity:
+		case anyAvailabilityZone:
+		default:
+			return nil, errors.Errorf("invalid dns_record.client_routing_policy set in annotation %s: got '%s' expected one of ['%s', '%s', '%s']",
+				annotations.SvcLBSuffixLoadBalancerAttributes, dnsRecordClientRoutingPolicy,
+				anyAvailabilityZone, partialAvailabilityZoneAffinity, availabilityZoneAffinity)
+		}
+	}
 	return attributes, nil
 }
 
diff --git a/pkg/service/model_build_load_balancer_test.go b/pkg/service/model_build_load_balancer_test.go
index 46de842a6..b90a566f3 100644
--- a/pkg/service/model_build_load_balancer_test.go
+++ b/pkg/service/model_build_load_balancer_test.go
@@ -88,7 +88,7 @@ func Test_defaultModelBuilderTask_buildLBAttributes(t *testing.T) {
 				ObjectMeta: metav1.ObjectMeta{
 					Annotations: map[string]string{
 						"service.beta.kubernetes.io/aws-load-balancer-attributes": "access_logs.s3.enabled=true,access_logs.s3.bucket=nlb-bucket," +
-							"access_logs.s3.prefix=bkt-pfx,load_balancing.cross_zone.enabled=true,deletion_protection.enabled=true",
+							"access_logs.s3.prefix=bkt-pfx,load_balancing.cross_zone.enabled=true,deletion_protection.enabled=true,dns_record.client_routing_policy=availability_zone_affinity",
 					},
 				},
 			},
@@ -114,6 +114,10 @@ func Test_defaultModelBuilderTask_buildLBAttributes(t *testing.T) {
 					Key:   lbAttrsDeletionProtectionEnabled,
 					Value: "true",
 				},
+				{
+					Key:   lbAttrsLoadBalancingDnsClientRoutingPolicy,
+					Value: availabilityZoneAffinity,
+				},
 			},
 		},
 		{