From 854afa7c739a6e843f78a6fcce6a3a28e23df245 Mon Sep 17 00:00:00 2001
From: lixiang <lixiang.cdc@gmail.com>
Date: Mon, 1 Mar 2021 11:13:44 +0800
Subject: [PATCH] PodsListing: Add WithLabelSelector option.

---
 pkg/descheduler/pod/pods.go      | 27 +++++++++++++++++++++++++--
 pkg/descheduler/pod/pods_test.go | 31 ++++++++++++++++++++++++++++++-
 2 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/pkg/descheduler/pod/pods.go b/pkg/descheduler/pod/pods.go
index 04b154310..db9790691 100644
--- a/pkg/descheduler/pod/pods.go
+++ b/pkg/descheduler/pod/pods.go
@@ -31,6 +31,7 @@ type Options struct {
 	filter             func(pod *v1.Pod) bool
 	includedNamespaces []string
 	excludedNamespaces []string
+	labelSelector      *metav1.LabelSelector
 }
 
 // WithFilter sets a pod filter.
@@ -55,6 +56,13 @@ func WithoutNamespaces(namespaces []string) func(opts *Options) {
 	}
 }
 
+// WithLabelSelector sets a pod label selector
+func WithLabelSelector(labelSelector *metav1.LabelSelector) func(opts *Options) {
+	return func(opts *Options) {
+		opts.labelSelector = labelSelector
+	}
+}
+
 // ListPodsOnANode lists all of the pods on a node
 // It also accepts an optional "filter" function which can be used to further limit the pods that are returned.
 // (Usually this is podEvictor.Evictable().IsEvictable, in order to only list the evictable pods on a node, but can
@@ -74,6 +82,15 @@ func ListPodsOnANode(
 
 	fieldSelectorString := "spec.nodeName=" + node.Name + ",status.phase!=" + string(v1.PodSucceeded) + ",status.phase!=" + string(v1.PodFailed)
 
+	labelSelectorString := ""
+	if options.labelSelector != nil {
+		selector, err := metav1.LabelSelectorAsSelector(options.labelSelector)
+		if err != nil {
+			return []*v1.Pod{}, err
+		}
+		labelSelectorString = selector.String()
+	}
+
 	if len(options.includedNamespaces) > 0 {
 		fieldSelector, err := fields.ParseSelector(fieldSelectorString)
 		if err != nil {
@@ -82,7 +99,10 @@ func ListPodsOnANode(
 
 		for _, namespace := range options.includedNamespaces {
 			podList, err := client.CoreV1().Pods(namespace).List(ctx,
-				metav1.ListOptions{FieldSelector: fieldSelector.String()})
+				metav1.ListOptions{
+					FieldSelector: fieldSelector.String(),
+					LabelSelector: labelSelectorString,
+				})
 			if err != nil {
 				return []*v1.Pod{}, err
 			}
@@ -111,7 +131,10 @@ func ListPodsOnANode(
 	// Once the descheduler switches to pod listers (through informers),
 	// We need to flip to client-side filtering.
 	podList, err := client.CoreV1().Pods(v1.NamespaceAll).List(ctx,
-		metav1.ListOptions{FieldSelector: fieldSelector.String()})
+		metav1.ListOptions{
+			FieldSelector: fieldSelector.String(),
+			LabelSelector: labelSelectorString,
+		})
 	if err != nil {
 		return []*v1.Pod{}, err
 	}
diff --git a/pkg/descheduler/pod/pods_test.go b/pkg/descheduler/pod/pods_test.go
index d67e042f2..8fa2b2cfb 100644
--- a/pkg/descheduler/pod/pods_test.go
+++ b/pkg/descheduler/pod/pods_test.go
@@ -24,6 +24,7 @@ import (
 	"testing"
 
 	v1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/client-go/kubernetes/fake"
 	core "k8s.io/client-go/testing"
@@ -40,6 +41,7 @@ func TestListPodsOnANode(t *testing.T) {
 		name             string
 		pods             map[string][]v1.Pod
 		node             *v1.Node
+		labelSelector    *metav1.LabelSelector
 		expectedPodCount int
 	}{
 		{
@@ -52,6 +54,33 @@ func TestListPodsOnANode(t *testing.T) {
 				"n2": {*test.BuildTestPod("pod3", 100, 0, "n2", nil)},
 			},
 			node:             test.BuildTestNode("n1", 2000, 3000, 10, nil),
+			labelSelector:    nil,
+			expectedPodCount: 2,
+		},
+		{
+			name: "test listing pods with label selector",
+			pods: map[string][]v1.Pod{
+				"n1": {
+					*test.BuildTestPod("pod1", 100, 0, "n1", nil),
+					*test.BuildTestPod("pod2", 100, 0, "n1", func(pod *v1.Pod) {
+						pod.Labels = map[string]string{"foo": "bar"}
+					}),
+					*test.BuildTestPod("pod3", 100, 0, "n1", func(pod *v1.Pod) {
+						pod.Labels = map[string]string{"foo": "bar1"}
+					}),
+				},
+				"n2": {*test.BuildTestPod("pod4", 100, 0, "n2", nil)},
+			},
+			node: test.BuildTestNode("n1", 2000, 3000, 10, nil),
+			labelSelector: &metav1.LabelSelector{
+				MatchExpressions: []metav1.LabelSelectorRequirement{
+					{
+						Key:      "foo",
+						Operator: metav1.LabelSelectorOpIn,
+						Values:   []string{"bar", "bar1"},
+					},
+				},
+			},
 			expectedPodCount: 2,
 		},
 	}
@@ -67,7 +96,7 @@ func TestListPodsOnANode(t *testing.T) {
 			}
 			return true, nil, fmt.Errorf("Failed to list: %v", list)
 		})
-		pods, _ := ListPodsOnANode(context.TODO(), fakeClient, testCase.node)
+		pods, _ := ListPodsOnANode(context.TODO(), fakeClient, testCase.node, WithLabelSelector(testCase.labelSelector))
 		if len(pods) != testCase.expectedPodCount {
 			t.Errorf("expected %v pods on node %v, got %+v", testCase.expectedPodCount, testCase.node.Name, len(pods))
 		}