From c582bbeb399ca91c1cf87fd4959596ecbe334de6 Mon Sep 17 00:00:00 2001
From: thiyyakat <meghana.thiyyakat@gmail.com>
Date: Tue, 22 Oct 2024 11:13:25 +0530
Subject: [PATCH] Address review comments

---
 pkg/util/provider/drain/drain.go | 31 +++++++++++++++----------------
 1 file changed, 15 insertions(+), 16 deletions(-)

diff --git a/pkg/util/provider/drain/drain.go b/pkg/util/provider/drain/drain.go
index 31e1cb4e4..a77e15407 100644
--- a/pkg/util/provider/drain/drain.go
+++ b/pkg/util/provider/drain/drain.go
@@ -239,9 +239,7 @@ func (o *Options) RunDrain(ctx context.Context) error {
 		klog.Errorf("Drain Error: Cordoning of node failed with error: %v", err)
 		return err
 	}
-	stopCh := make(chan struct{})
-	defer close(stopCh)
-	if !cache.WaitForCacheSync(stopCh, o.podSynced) {
+	if !cache.WaitForCacheSync(drainContext.Done(), o.podSynced) {
 		err := fmt.Errorf("timed out waiting for pod cache to sync")
 		return err
 	}
@@ -363,22 +361,23 @@ func (o *Options) getPodsForDeletion() (pods []corev1.Pod, err error) {
 	fs := podStatuses{}
 
 	for _, pod := range podList {
-		if pod.Spec.NodeName == o.nodeName {
-			podOk := true
-			for _, filt := range []podFilter{mirrorPodFilter, o.localStorageFilter, o.unreplicatedFilter, o.daemonsetFilter} {
-				filterOk, w, f := filt(*pod)
-				podOk = podOk && filterOk
-				if w != nil {
-					ws[w.string] = append(ws[w.string], pod.Name)
-				}
-				if f != nil {
-					fs[f.string] = append(fs[f.string], pod.Name)
-				}
+		if pod.Spec.NodeName != o.nodeName {
+			continue
+		}
+		podOk := true
+		for _, filt := range []podFilter{mirrorPodFilter, o.localStorageFilter, o.unreplicatedFilter, o.daemonsetFilter} {
+			filterOk, w, f := filt(*pod)
+			podOk = podOk && filterOk
+			if w != nil {
+				ws[w.string] = append(ws[w.string], pod.Name)
 			}
-			if podOk {
-				pods = append(pods, *pod)
+			if f != nil {
+				fs[f.string] = append(fs[f.string], pod.Name)
 			}
 		}
+		if podOk {
+			pods = append(pods, *pod)
+		}
 	}
 
 	if len(fs) > 0 {