From aed9f492df4fd2d188ed972eaa8292e76fc01b37 Mon Sep 17 00:00:00 2001
From: Jason Deal <jmdeal@amazon.com>
Date: Thu, 7 Dec 2023 18:12:12 -0800
Subject: [PATCH] support matchLabelKeys for topology

---
 .../provisioning/scheduling/topology.go       |  5 ++++
 .../provisioning/scheduling/topology_test.go  | 30 +++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/pkg/controllers/provisioning/scheduling/topology.go b/pkg/controllers/provisioning/scheduling/topology.go
index 9ac70ca04d..be8a7c2d60 100644
--- a/pkg/controllers/provisioning/scheduling/topology.go
+++ b/pkg/controllers/provisioning/scheduling/topology.go
@@ -315,6 +315,11 @@ func (t *Topology) countDomains(ctx context.Context, tg *TopologyGroup) error {
 func (t *Topology) newForTopologies(p *v1.Pod) []*TopologyGroup {
 	var topologyGroups []*TopologyGroup
 	for _, cs := range p.Spec.TopologySpreadConstraints {
+		for _, label := range cs.MatchLabelKeys {
+			if value, ok := p.ObjectMeta.Labels[label]; ok {
+				cs.LabelSelector.MatchLabels[label] = value
+			}
+		}
 		topologyGroups = append(topologyGroups, NewTopologyGroup(TopologyTypeSpread, cs.TopologyKey, p, sets.New(p.Namespace),
 			cs.LabelSelector, cs.MaxSkew, cs.MinDomains, cs.NodeTaintsPolicy, cs.NodeAffinityPolicy, t.domains[cs.TopologyKey]))
 	}
diff --git a/pkg/controllers/provisioning/scheduling/topology_test.go b/pkg/controllers/provisioning/scheduling/topology_test.go
index 54cfca9638..4140fe4204 100644
--- a/pkg/controllers/provisioning/scheduling/topology_test.go
+++ b/pkg/controllers/provisioning/scheduling/topology_test.go
@@ -89,6 +89,36 @@ var _ = Describe("Topology", func() {
 		ExpectSkew(ctx, env.Client, "default", &topology[0]).To(ConsistOf(2))
 	})
 
+	It("should support matchLabelKeys", func() {
+		if env.Version.Minor() < 27 {
+			Skip("Feature \"matchLabelKeys\" only available in K8S >= 1.27.x")
+		}
+		matchedLabel := "test-label"
+		topology := []v1.TopologySpreadConstraint{{
+			TopologyKey:       v1.LabelHostname,
+			WhenUnsatisfiable: v1.DoNotSchedule,
+			LabelSelector:     &metav1.LabelSelector{MatchLabels: labels},
+			MatchLabelKeys:    []string{matchedLabel},
+			MaxSkew:           1,
+		}}
+		count := 2
+		pods := test.UnschedulablePods(test.PodOptions{
+			ObjectMeta: metav1.ObjectMeta{
+				Labels: lo.Assign(labels, map[string]string{ matchedLabel: "value-a" }),
+			},
+			TopologySpreadConstraints: topology,
+		}, count)
+		pods = append(pods, test.UnschedulablePods(test.PodOptions{
+			ObjectMeta: metav1.ObjectMeta{
+				Labels: lo.Assign(labels, map[string]string{ matchedLabel: "value-b" }),
+			},
+			TopologySpreadConstraints: topology,
+		}, count)...)
+		ExpectApplied(ctx, env.Client, nodePool)
+		ExpectProvisioned(ctx, env.Client, cluster, cloudProvider, prov, pods...)
+		ExpectSkew(ctx, env.Client, "default", &topology[0]).To(ConsistOf(2, 2))
+	})
+
 	Context("Zonal", func() {
 		It("should balance pods across zones (match labels)", func() {
 			topology := []v1.TopologySpreadConstraint{{