From 690b60fe661c337c9ea9a9b51093a44c28b5191c Mon Sep 17 00:00:00 2001
From: Jonathan Innis <joinnis@amazon.com>
Date: Wed, 5 Feb 2025 15:25:08 -0800
Subject: [PATCH] Only select DaemonSet pods that match DaemonSet labels

---
 pkg/controllers/state/cluster.go    | 10 ++++++----
 pkg/controllers/state/suite_test.go | 20 ++++++++++++++++++++
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/pkg/controllers/state/cluster.go b/pkg/controllers/state/cluster.go
index 56f71a09b7..07f990532d 100644
--- a/pkg/controllers/state/cluster.go
+++ b/pkg/controllers/state/cluster.go
@@ -478,22 +478,24 @@ func (c *Cluster) GetDaemonSetPod(daemonset *appsv1.DaemonSet) *corev1.Pod {
 
 func (c *Cluster) UpdateDaemonSet(ctx context.Context, daemonset *appsv1.DaemonSet) error {
 	pods := &corev1.PodList{}
-	err := c.kubeClient.List(ctx, pods, client.InNamespace(daemonset.Namespace))
+	// Scope down this call to only select the pods in this namespace that specifically match the DaemonSet
+	opts := []client.ListOption{client.InNamespace(daemonset.Namespace)}
+	if selector, err := metav1.LabelSelectorAsSelector(daemonset.Spec.Selector); err == nil {
+		opts = append(opts, client.MatchingLabelsSelector{Selector: selector})
+	}
+	err := c.kubeClient.List(ctx, pods, opts...)
 	if err != nil {
 		return err
 	}
-
 	sort.Slice(pods.Items, func(i, j int) bool {
 		return pods.Items[i].CreationTimestamp.Unix() > pods.Items[j].CreationTimestamp.Unix()
 	})
-
 	for i := range pods.Items {
 		if metav1.IsControlledBy(&pods.Items[i], daemonset) {
 			c.daemonSetPods.Store(client.ObjectKeyFromObject(daemonset), &pods.Items[i])
 			break
 		}
 	}
-
 	return nil
 }
 
diff --git a/pkg/controllers/state/suite_test.go b/pkg/controllers/state/suite_test.go
index 9e52e503a3..27f3bd97ce 100644
--- a/pkg/controllers/state/suite_test.go
+++ b/pkg/controllers/state/suite_test.go
@@ -1376,8 +1376,12 @@ var _ = Describe("Cluster State Sync", func() {
 
 var _ = Describe("DaemonSet Controller", func() {
 	It("should not update daemonsetCache when daemonset pod is not present", func() {
+		labels := map[string]string{"key": "value"}
 		daemonset := test.DaemonSet(
 			test.DaemonSetOptions{PodOptions: test.PodOptions{
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: labels,
+				},
 				ResourceRequirements: corev1.ResourceRequirements{Requests: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("1"), corev1.ResourceMemory: resource.MustParse("1Gi")}},
 			}},
 		)
@@ -1387,8 +1391,12 @@ var _ = Describe("DaemonSet Controller", func() {
 		Expect(daemonsetPod).To(BeNil())
 	})
 	It("should update daemonsetCache when daemonset pod is created", func() {
+		labels := map[string]string{"key": "value"}
 		daemonset := test.DaemonSet(
 			test.DaemonSetOptions{PodOptions: test.PodOptions{
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: labels,
+				},
 				ResourceRequirements: corev1.ResourceRequirements{Requests: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("1"), corev1.ResourceMemory: resource.MustParse("1Gi")}},
 			}},
 		)
@@ -1396,6 +1404,7 @@ var _ = Describe("DaemonSet Controller", func() {
 		daemonsetPod := test.UnschedulablePod(
 			test.PodOptions{
 				ObjectMeta: metav1.ObjectMeta{
+					Labels: labels,
 					OwnerReferences: []metav1.OwnerReference{
 						{
 							APIVersion:         "apps/v1",
@@ -1415,8 +1424,12 @@ var _ = Describe("DaemonSet Controller", func() {
 		Expect(cluster.GetDaemonSetPod(daemonset)).To(Equal(daemonsetPod))
 	})
 	It("should update daemonsetCache with the newest created pod", func() {
+		labels := map[string]string{"key": "value"}
 		daemonset := test.DaemonSet(
 			test.DaemonSetOptions{PodOptions: test.PodOptions{
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: labels,
+				},
 				ResourceRequirements: corev1.ResourceRequirements{Requests: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("1"), corev1.ResourceMemory: resource.MustParse("1Gi")}},
 			}},
 		)
@@ -1424,6 +1437,7 @@ var _ = Describe("DaemonSet Controller", func() {
 		daemonsetPod1 := test.UnschedulablePod(
 			test.PodOptions{
 				ObjectMeta: metav1.ObjectMeta{
+					Labels: labels,
 					OwnerReferences: []metav1.OwnerReference{
 						{
 							APIVersion:         "apps/v1",
@@ -1445,6 +1459,7 @@ var _ = Describe("DaemonSet Controller", func() {
 		daemonsetPod2 := test.UnschedulablePod(
 			test.PodOptions{
 				ObjectMeta: metav1.ObjectMeta{
+					Labels: labels,
 					OwnerReferences: []metav1.OwnerReference{
 						{
 							APIVersion:         "apps/v1",
@@ -1464,8 +1479,12 @@ var _ = Describe("DaemonSet Controller", func() {
 		Expect(cluster.GetDaemonSetPod(daemonset)).To(Equal(daemonsetPod2))
 	})
 	It("should delete daemonset in cache when daemonset is deleted", func() {
+		labels := map[string]string{"key": "value"}
 		daemonset := test.DaemonSet(
 			test.DaemonSetOptions{PodOptions: test.PodOptions{
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: labels,
+				},
 				ResourceRequirements: corev1.ResourceRequirements{Requests: corev1.ResourceList{corev1.ResourceCPU: resource.MustParse("1"), corev1.ResourceMemory: resource.MustParse("1Gi")}},
 			}},
 		)
@@ -1473,6 +1492,7 @@ var _ = Describe("DaemonSet Controller", func() {
 		daemonsetPod := test.UnschedulablePod(
 			test.PodOptions{
 				ObjectMeta: metav1.ObjectMeta{
+					Labels: labels,
 					OwnerReferences: []metav1.OwnerReference{
 						{
 							APIVersion:         "apps/v1",