From 1961facfaf524b42e5a3337ce22648cc79797ed5 Mon Sep 17 00:00:00 2001
From: Jason Deal <jmdeal@amazon.com>
Date: Fri, 26 Jan 2024 14:07:22 -0800
Subject: [PATCH] chore: update for empty domain tracking

---
 .../scheduling/topologydomaingroup.go             |  2 +-
 .../provisioning/scheduling/topologygroup.go      | 15 ++++++++++-----
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/pkg/controllers/provisioning/scheduling/topologydomaingroup.go b/pkg/controllers/provisioning/scheduling/topologydomaingroup.go
index 8fa783a99b..8f99193685 100644
--- a/pkg/controllers/provisioning/scheduling/topologydomaingroup.go
+++ b/pkg/controllers/provisioning/scheduling/topologydomaingroup.go
@@ -28,7 +28,7 @@ import (
 type TopologyDomainGroup map[string][][]v1.Taint
 
 func NewTopologyDomainGroup() TopologyDomainGroup {
-	return map[string][][]v1.Taint{}
+	return TopologyDomainGroup{}
 }
 
 // Insert either adds a new domain to the TopologyDomainGroup or updates an existing domain.
diff --git a/pkg/controllers/provisioning/scheduling/topologygroup.go b/pkg/controllers/provisioning/scheduling/topologygroup.go
index e1bad1a624..e71f7945c0 100644
--- a/pkg/controllers/provisioning/scheduling/topologygroup.go
+++ b/pkg/controllers/provisioning/scheduling/topologygroup.go
@@ -83,17 +83,22 @@ func NewTopologyGroup(topologyType TopologyType, topologyKey string, pod *v1.Pod
 		nodeSelector = MakeTopologyNodeFilter(pod, nodeTaintsPolicy, nodeAffinityPolicy)
 	}
 
-	domainCounts := map[string]int32{}
+	domains := map[string]int32{}
 	if nodeSelector.TaintPolicy == v1.NodeInclusionPolicyHonor {
 		domainGroup.ForEachToleratedDomain(pod, func(domain string) {
-			domainCounts[domain] = 0
+			domains[domain] = 0
 		})
 	} else {
 		domainGroup.ForEachDomain(func(domain string) {
-			domainCounts[domain] = 0
+			domains[domain] = 0
 		})
 	}
 
+	emptyDomains := sets.New[string]()
+	for domain := range domains {
+		emptyDomains.Insert(domain)
+	}
+
 	return &TopologyGroup{
 		Type:         topologyType,
 		Key:          topologyKey,
@@ -101,8 +106,8 @@ func NewTopologyGroup(topologyType TopologyType, topologyKey string, pod *v1.Pod
 		selector:     labelSelector,
 		nodeFilter:   nodeSelector,
 		maxSkew:      maxSkew,
-		domains:      domainCounts,
-		emptyDomains: domains.Clone(),
+		domains:      domains,
+		emptyDomains: emptyDomains,
 		owners:       map[types.UID]struct{}{},
 		minDomains:   minDomains,
 	}