Support addtional filter for extended resource pod

cmd/kured/main.go

@@ -30,6 +30,7 @@ import (
 	"github.com/google/shlex"
 
 	shoutrrr "github.com/containrrr/shoutrrr"
+	"github.com/kubereboot/kured/pkg/additionalfilters"
 	"github.com/kubereboot/kured/pkg/alerts"
 	"github.com/kubereboot/kured/pkg/daemonsetlock"
 	"github.com/kubereboot/kured/pkg/delaytick"
@@ -55,6 +56,7 @@ var (
 	metricsPort                     int
 	drainGracePeriod                int
 	drainPodSelector                string
+	drainResourceNameWithPrefixes   string
 	skipWaitForDeleteTimeoutSeconds int
 	dsNamespace                     string
 	dsName                          string
@@ -149,6 +151,8 @@ func NewRootCommand() *cobra.Command {
 		"time in seconds given to each pod to terminate gracefully, if negative, the default value specified in the pod will be used")
 	rootCmd.PersistentFlags().StringVar(&drainPodSelector, "drain-pod-selector", "",
 		"only drain pods with labels matching the selector (default: '', all pods)")
+	rootCmd.PersistentFlags().StringVar(&drainResourceNameWithPrefixes, "drain-resource-name-with-prefixes", "",
+		"only drain those pods whose resource name has the specified prefix (default: '', all pods)")
 	rootCmd.PersistentFlags().IntVar(&skipWaitForDeleteTimeoutSeconds, "skip-wait-for-delete-timeout", 0,
 		"when seconds is greater than zero, skip waiting for the pods whose deletion timestamp is older than N seconds while draining a node")
 	rootCmd.PersistentFlags().DurationVar(&drainDelay, "drain-delay", 0,
@@ -526,6 +530,12 @@ func drain(client *kubernetes.Clientset, node *v1.Node) error {
 		Timeout:                         drainTimeout,
 	}
 
+	if len(drainResourceNameWithPrefixes) > 0 {
+		// Add resource name filter
+		rtf := additionalfilters.NewResourceNameFilter(drainResourceNameWithPrefixes)
+		drainer.AdditionalFilters = []kubectldrain.PodFilter{rtf.GetPodFilter()}
+	}
+
 	if err := kubectldrain.RunCordonOrUncordon(drainer, node, true); err != nil {
 		log.Errorf("Error cordonning %s: %v", nodename, err)
 		return err

pkg/additionalfilters/resourcenamefilter.go

@@ -0,0 +1,52 @@
+package additionalfilters
+
+import (
+	"strings"
+
+	v1 "k8s.io/api/core/v1"
+	kubectldrain "k8s.io/kubectl/pkg/drain"
+)
+
+type resourceNameFilter struct {
+	filter kubectldrain.PodFilter
+}
+
+func NewResourceNameFilter(resourceNames string) *resourceNameFilter {
+	// Create a custom drain filter which will be passed to the drain helper.
+	// The drain helper will carry out the actual deletion of pods on a node.
+	customDrainFilter := func(pod v1.Pod) kubectldrain.PodDeleteStatus {
+		delete := extendedResourceFilter(pod, strings.Split(resourceNames, ","))
+		if !delete {
+			return kubectldrain.MakePodDeleteStatusSkip()
+		}
+		return kubectldrain.MakePodDeleteStatusOkay()
+	}
+	return &resourceNameFilter{
+		filter: customDrainFilter,
+	}
+}
+
+func (r *resourceNameFilter) GetPodFilter() kubectldrain.PodFilter {
+	return r.filter
+}
+
+func extendedResourceFilter(pod v1.Pod, resourceNamePrefixs []string) bool {
+	podHasExtendedResource := func(rl v1.ResourceList) bool {
+		for resourceName := range rl {
+			for _, prefix := range resourceNamePrefixs {
+				if strings.HasPrefix(string(resourceName), prefix) {
+					return true
+				}
+			}
+
+		}
+		return false
+	}
+
+	for _, c := range pod.Spec.Containers {
+		if podHasExtendedResource(c.Resources.Limits) || podHasExtendedResource(c.Resources.Requests) {
+			return true
+		}
+	}
+	return false
+}