Make drain package reusable for controller-runtime based callers

pkg/util/provider/drain/drain.go

@@ -77,6 +77,25 @@ type Options struct {
 	volumeAttachmentHandler      *VolumeAttachmentHandler
 	Timeout                      time.Duration
 	podSynced                    cache.InformerSynced
+
+	// AdditionalPodFilters are extra predicates evaluated alongside the built-in filters when selecting pods to
+	// drain.
+	AdditionalPodFilters []AdditionalPodFilter
+	// podProvider, if set, is used to list the pods on the node instead of podLister. It allows callers that do not
+	// run client-go informers (e.g. controller-runtime based controllers) to supply their own pod source.
+	podProvider PodProvider
+	// SkipVolumeHandling, if true, skips the PersistentVolume detach/reattach handling during eviction.
+	SkipVolumeHandling bool
+}
+
+// AdditionalPodFilter takes a pod and returns whether the pod should be included for draining (true) or excluded
+// (false). It is a caller-supplied predicate evaluated in addition to the built-in filters.
+type AdditionalPodFilter func(corev1.Pod) bool
+
+// PodProvider lists the pods running on a given node. It abstracts the pod source (client-go lister vs. a
+// controller-runtime client) used during drain.
+type PodProvider interface {
+	PodsForNode(ctx context.Context, nodeName string) ([]corev1.Pod, error)
 }
 
 // Takes a pod and returns a bool indicating whether or not to operate on the
@@ -214,6 +233,13 @@ func NewDrainOptions(
 	}
 }
 
+// SetPodProvider sets the PodProvider used to list the pods on the node instead of the client-go pod lister. It lets
+// callers that do not run client-go informers (e.g. controller-runtime based controllers) supply their own pod
+// source when using NewDrainOptions/RunDrain.
+func (o *Options) SetPodProvider(podProvider PodProvider) {
+	o.podProvider = podProvider
+}
+
 // RunDrain runs the 'drain' command
 func (o *Options) RunDrain(ctx context.Context) error {
 	o.drainStartedOn = time.Now()
@@ -239,7 +265,7 @@ func (o *Options) RunDrain(ctx context.Context) error {
 		klog.Errorf("Drain Error: Cordoning of node failed with error: %v", err)
 		return err
 	}
-	if !cache.WaitForCacheSync(drainContext.Done(), o.podSynced) {
+	if o.podSynced != nil && !cache.WaitForCacheSync(drainContext.Done(), o.podSynced) {
 		err := fmt.Errorf("timed out waiting for pod cache to sync")
 		return err
 	}
@@ -249,14 +275,14 @@ func (o *Options) RunDrain(ctx context.Context) error {
 }
 
 func (o *Options) deleteOrEvictPodsSimple(ctx context.Context) error {
-	pods, err := o.getPodsForDeletion()
+	pods, err := o.getPodsForDeletion(ctx)
 	if err != nil {
 		return err
 	}
 
 	err = o.deleteOrEvictPods(ctx, pods)
 	if err != nil {
-		pendingPods, newErr := o.getPodsForDeletion()
+		pendingPods, newErr := o.getPodsForDeletion(ctx)
 		if newErr != nil {
 			return newErr
 		}
@@ -348,10 +374,21 @@ func (ps podStatuses) Message() string {
 
 // getPodsForDeletion returns all the pods we're going to delete.  If there are
 // any pods preventing us from deleting, we return that list in an error.
-func (o *Options) getPodsForDeletion() (pods []corev1.Pod, err error) {
-	podList, err := o.podLister.List(labels.Everything())
-	if err != nil {
-		return
+func (o *Options) getPodsForDeletion(ctx context.Context) (pods []corev1.Pod, err error) {
+	var podList []*corev1.Pod
+	if o.podProvider != nil {
+		nodePods, providerErr := o.podProvider.PodsForNode(ctx, o.nodeName)
+		if providerErr != nil {
+			return nil, providerErr
+		}
+		for _, pod := range nodePods {
+			podList = append(podList, &pod)
+		}
+	} else {
+		podList, err = o.podLister.List(labels.Everything())
+		if err != nil {
+			return
+		}
 	}
 	if len(podList) == 0 {
 		klog.Infof("no pods found in store")
@@ -375,6 +412,9 @@ func (o *Options) getPodsForDeletion() (pods []corev1.Pod, err error) {
 				fs[f.string] = append(fs[f.string], pod.Name)
 			}
 		}
+		for _, filt := range o.AdditionalPodFilters {
+			podOk = podOk && filt(*pod)
+		}
 		if podOk {
 			pods = append(pods, *pod)
 		}
@@ -491,15 +531,15 @@ func (o *Options) evictPods(ctx context.Context, attemptEvict bool, pods []corev
 	returnCh := make(chan error, len(pods))
 	defer close(returnCh)
 
-	if o.ForceDeletePods {
+	if o.ForceDeletePods || o.SkipVolumeHandling {
 		podsToDrain := make([]*corev1.Pod, len(pods))
 		for i := range pods {
 			podsToDrain[i] = &pods[i]
 		}
 
-		klog.V(3).Infof("Forceful eviction of pods on the node: %q", o.nodeName)
+		klog.V(3).Infof("Eviction of pods on the node without volume handling: %q", o.nodeName)
 
-		// evict all pods in parallel without waiting for pods or volume detachment
+		// Evict all pods without waiting for volume detachment.
 		go o.evictPodsWithoutPv(ctx, attemptEvict, podsToDrain, policyGroupVersion, getPodFn, returnCh)
 	} else {
 		podsWithPv, podsWithoutPv := filterPodsWithPv(pods)
@@ -1193,6 +1233,9 @@ func (o *Options) GetDrainDuration() time.Duration {
 }
 
 func getPdbForPod(pdbLister policyv1listers.PodDisruptionBudgetLister, pod *corev1.Pod) *policyv1.PodDisruptionBudget {
+	if pdbLister == nil {
+		return nil
+	}
 	// GetPodPodDisruptionBudgets returns an error only if no PodDisruptionBudgets are found.
 	// We don't return that as an error to the caller.
 	pdbs, err := pdbLister.GetPodPodDisruptionBudgets(pod)

pkg/util/provider/drain/drain_test.go

@@ -60,6 +60,7 @@ var _ = Describe("drain", func() {
 		terminationGracePeriod    time.Duration
 		pvReattachTimeout         time.Duration
 		force                     bool
+		skipVolumeHandling        bool
 		evictError                error
 		deleteError               error
 	}
@@ -163,6 +164,7 @@ var _ = Describe("drain", func() {
 			Timeout:                      2 * time.Minute,
 			volumeAttachmentHandler:      volumeAttachmentHandler,
 			podSynced:                    podSynced,
+			SkipVolumeHandling:           setup.skipVolumeHandling,
 		}
 
 		// Get the pod directly from the ObjectTracker to avoid locking issues in the Fake object.
@@ -298,12 +300,10 @@ var _ = Describe("drain", func() {
 					}
 
 					// Delete the pod asyncronously to work around the lock problems in testing.Fake
-					wg.Add(1)
-					go func() {
-						defer wg.Done()
+					wg.Go(func() {
 						runPodDrainHandlers(pod)
 						fmt.Fprintf(GinkgoWriter, "Drained pod %s/%s in %s\n", pod.Namespace, pod.Name, time.Since(start).String())
-					}()
+					})
 
 					nEvictions++
 					return
@@ -330,12 +330,10 @@ var _ = Describe("drain", func() {
 					}
 
 					// Delete the pod asyncronously to work around the lock problems in testing.Fake
-					wg.Add(1)
-					go func() {
-						defer wg.Done()
+					wg.Go(func() {
 						runPodDrainHandlers(pod)
 						fmt.Fprintf(GinkgoWriter, "Drained pod %s/%s in %s\n", pod.Namespace, pod.Name, time.Since(start).String())
-					}()
+					})
 				default:
 					err = fmt.Errorf("Expected type k8stesting.GetAction but got %T", action)
 				}
@@ -858,8 +856,185 @@ var _ = Describe("drain", func() {
 				// Because waitForDetach polling Interval is equal to terminationGracePeriodShort
 				minDrainDuration: terminationGracePeriodMedium,
 			}),
+		Entry("Successful drain with SkipVolumeHandling and eviction of pods with exclusive volumes",
+			&setup{
+				stats: stats{
+					nPodsWithoutPV:                0,
+					nPodsWithOnlyExclusivePV:      2,
+					nPodsWithOnlySharedPV:         0,
+					nPodsWithExclusiveAndSharedPV: 0,
+					nPVsPerPodWithExclusivePV:     1,
+				},
+				attemptEviction:        true,
+				terminationGracePeriod: terminationGracePeriodShort,
+				skipVolumeHandling:     true,
+			},
+			// SkipVolumeHandling routes pods through evictPodsWithoutPv, which does NOT wait for volume detach.
+			// The test only needs deletePod so that the fake eviction reactor finds and removes the pod.
+			[]podDrainHandler{deletePod},
+			&expectation{
+				stats: stats{
+					nPodsWithoutPV:                0,
+					nPodsWithOnlyExclusivePV:      0,
+					nPodsWithOnlySharedPV:         0,
+					nPodsWithExclusiveAndSharedPV: 0,
+				},
+				// No PV detach/reattach wait — completes quickly.
+				timeout:          terminationGracePeriodMedium,
+				drainTimeout:     false,
+				drainError:       nil,
+				nEvictions:       2,
+				minDrainDuration: 0,
+			}),
+		Entry("Successful drain with SkipVolumeHandling without eviction of pods with exclusive volumes",
+			&setup{
+				stats: stats{
+					nPodsWithoutPV:                0,
+					nPodsWithOnlyExclusivePV:      2,
+					nPodsWithOnlySharedPV:         0,
+					nPodsWithExclusiveAndSharedPV: 0,
+					nPVsPerPodWithExclusivePV:     1,
+				},
+				attemptEviction:        false,
+				terminationGracePeriod: terminationGracePeriodShort,
+				skipVolumeHandling:     true,
+			},
+			nil,
+			&expectation{
+				stats: stats{
+					nPodsWithoutPV:                0,
+					nPodsWithOnlyExclusivePV:      0,
+					nPodsWithOnlySharedPV:         0,
+					nPodsWithExclusiveAndSharedPV: 0,
+				},
+				timeout:          terminationGracePeriodShort,
+				drainTimeout:     false,
+				drainError:       nil,
+				nEvictions:       0,
+				minDrainDuration: 0,
+			}),
 	)
 
+	Describe("getPodsForDeletion", func() {
+		var (
+			ctx      context.Context
+			nodeName string
+			drain    *Options
+		)
+
+		BeforeEach(func() {
+			ctx = context.Background()
+			nodeName = oldNodeName
+			drain = &Options{
+				ErrOut:                       GinkgoWriter,
+				Out:                          GinkgoWriter,
+				nodeName:                     nodeName,
+				IgnorePodsWithoutControllers: true,
+				IgnoreDaemonsets:             true,
+				DeleteLocalData:              true,
+			}
+		})
+
+		Context("PodProvider", func() {
+			It("uses PodProvider instead of podLister when set", func() {
+				pod := getPodWithoutPV(testNamespace, "pod-0", nodeName, terminationGracePeriodDefault, nil)
+				drain.podProvider = &fakePodProvider{pods: []corev1.Pod{*pod}}
+
+				pods, err := drain.getPodsForDeletion(ctx)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(pods).To(HaveLen(1))
+				Expect(pods[0].Name).To(Equal("pod-0"))
+			})
+
+			It("returns only pods on the target node from PodProvider", func() {
+				podOnNode := getPodWithoutPV(testNamespace, "on-node", nodeName, terminationGracePeriodDefault, nil)
+				podOtherNode := getPodWithoutPV(testNamespace, "other-node", "different-node", terminationGracePeriodDefault, nil)
+				drain.podProvider = &fakePodProvider{pods: []corev1.Pod{*podOnNode, *podOtherNode}}
+
+				pods, err := drain.getPodsForDeletion(ctx)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(pods).To(HaveLen(1))
+				Expect(pods[0].Name).To(Equal("on-node"))
+			})
+
+			It("returns error when PodProvider fails", func() {
+				drain.podProvider = &fakePodProvider{err: fmt.Errorf("provider error")}
+
+				_, err := drain.getPodsForDeletion(ctx)
+				Expect(err).To(MatchError("provider error"))
+			})
+
+			It("returns empty list when PodProvider returns no pods", func() {
+				drain.podProvider = &fakePodProvider{pods: []corev1.Pod{}}
+
+				pods, err := drain.getPodsForDeletion(ctx)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(pods).To(BeEmpty())
+			})
+		})
+
+		Context("AdditionalPodFilters", func() {
+			var podProvider *fakePodProvider
+
+			BeforeEach(func() {
+				pods := getPodsWithoutPV(4, testNamespace, "pod-", nodeName, terminationGracePeriodDefault, nil)
+				podProvider = &fakePodProvider{}
+				for _, p := range pods {
+					podProvider.pods = append(podProvider.pods, *p)
+				}
+				drain.podProvider = podProvider
+			})
+
+			It("includes all pods when no additional filters are set", func() {
+				pods, err := drain.getPodsForDeletion(ctx)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(pods).To(HaveLen(4))
+			})
+
+			It("filters pods with a single additional filter", func() {
+				// Keep only pods whose name ends with "0" or "1"
+				drain.AdditionalPodFilters = []AdditionalPodFilter{
+					func(pod corev1.Pod) bool {
+						return pod.Name == "pod-0" || pod.Name == "pod-1"
+					},
+				}
+
+				pods, err := drain.getPodsForDeletion(ctx)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(pods).To(HaveLen(2))
+				Expect(pods).To(ConsistOf(
+					HaveField("Name", "pod-0"),
+					HaveField("Name", "pod-1"),
+				))
+			})
+
+			It("combines multiple additional filters with AND semantics", func() {
+				drain.AdditionalPodFilters = []AdditionalPodFilter{
+					func(pod corev1.Pod) bool { return pod.Name != "pod-0" },
+					func(pod corev1.Pod) bool { return pod.Name != "pod-1" },
+				}
+
+				pods, err := drain.getPodsForDeletion(ctx)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(pods).To(HaveLen(2))
+				Expect(pods).To(ConsistOf(
+					HaveField("Name", "pod-2"),
+					HaveField("Name", "pod-3"),
+				))
+			})
+
+			It("returns empty list when an additional filter excludes all pods", func() {
+				drain.AdditionalPodFilters = []AdditionalPodFilter{
+					func(_ corev1.Pod) bool { return false },
+				}
+
+				pods, err := drain.getPodsForDeletion(ctx)
+				Expect(err).NotTo(HaveOccurred())
+				Expect(pods).To(BeEmpty())
+			})
+		})
+	})
+
 	Describe("getPodVolumeInfos", func() {
 		var (
 			ctx context.Context
@@ -1313,3 +1488,15 @@ func appendSuffixToVolumeHandles(pvs []*corev1.PersistentVolume, suffix string)
 	}
 	return pvs
 }
+
+type fakePodProvider struct {
+	pods []corev1.Pod
+	err  error
+}
+
+func (f *fakePodProvider) PodsForNode(_ context.Context, _ string) ([]corev1.Pod, error) {
+	if f.err != nil {
+		return nil, f.err
+	}
+	return f.pods, nil
+}