Extract drainability rules into packages

cluster-autoscaler/simulator/drain.go

@@ -26,6 +26,7 @@ import (
 	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/autoscaler/cluster-autoscaler/config"
 	"k8s.io/autoscaler/cluster-autoscaler/simulator/drainability"
+	"k8s.io/autoscaler/cluster-autoscaler/simulator/drainability/rules"
 	"k8s.io/autoscaler/cluster-autoscaler/utils/drain"
 	kube_util "k8s.io/autoscaler/cluster-autoscaler/utils/kubernetes"
 	pod_util "k8s.io/autoscaler/cluster-autoscaler/utils/pod"
@@ -44,7 +45,7 @@ type NodeDeleteOptions struct {
 	// to allow their pods deletion in scale down
 	MinReplicaCount int
 	// DrainabilityRules contain a list of checks that are used to verify whether a pod can be drained from node.
-	DrainabilityRules []drainability.Rule
+	DrainabilityRules rules.Rules
 }
 
 // NewNodeDeleteOptions returns new node delete options extracted from autoscaling options
@@ -54,7 +55,7 @@ func NewNodeDeleteOptions(opts config.AutoscalingOptions) NodeDeleteOptions {
 		SkipNodesWithLocalStorage:         opts.SkipNodesWithLocalStorage,
 		MinReplicaCount:                   opts.MinReplicaCount,
 		SkipNodesWithCustomControllerPods: opts.SkipNodesWithCustomControllerPods,
-		DrainabilityRules:                 drainability.DefaultRules(),
+		DrainabilityRules:                 rules.Default(),
 	}
 }
 
@@ -71,12 +72,15 @@ func GetPodsToMove(nodeInfo *schedulerframework.NodeInfo, deleteOptions NodeDele
 	var drainPods, drainDs []*apiv1.Pod
 	drainabilityRules := deleteOptions.DrainabilityRules
 	if drainabilityRules == nil {
-		drainabilityRules = drainability.DefaultRules()
+		drainabilityRules = rules.Default()
+	}
+	drainCtx := &drainability.DrainContext{
+		Pdbs: pdbs,
 	}
 	for _, podInfo := range nodeInfo.Pods {
 		pod := podInfo.Pod
-		d := drainabilityStatus(pod, drainabilityRules)
-		switch d.Outcome {
+		status := drainabilityRules.Drainable(drainCtx, pod)
+		switch status.Outcome {
 		case drainability.UndefinedOutcome:
 			pods = append(pods, podInfo.Pod)
 		case drainability.DrainOk:
@@ -86,12 +90,15 @@ func GetPodsToMove(nodeInfo *schedulerframework.NodeInfo, deleteOptions NodeDele
 				drainPods = append(drainPods, pod)
 			}
 		case drainability.BlockDrain:
-			blockingPod = &drain.BlockingPod{pod, d.BlockingReason}
-			err = d.Error
+			blockingPod = &drain.BlockingPod{
+				Pod:    pod,
+				Reason: status.BlockingReason,
+			}
+			err = status.Error
 			return
-		case drainability.SkipDrain:
 		}
 	}
+
 	pods, daemonSetPods, blockingPod, err = drain.GetPodsForDeletionOnNodeDrain(
 		pods,
 		pdbs,
@@ -131,14 +138,3 @@ func checkPdbs(pods []*apiv1.Pod, pdbs []*policyv1.PodDisruptionBudget) (*drain.
 	}
 	return nil, nil
 }
-
-func drainabilityStatus(pod *apiv1.Pod, dr []drainability.Rule) drainability.Status {
-	for _, f := range dr {
-		if d := f.Drainable(pod); d.Outcome != drainability.UndefinedOutcome {
-			return d
-		}
-	}
-	return drainability.Status{
-		Outcome: drainability.UndefinedOutcome,
-	}
-}

cluster-autoscaler/simulator/drain_test.go

@@ -26,6 +26,7 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/intstr"
 	"k8s.io/autoscaler/cluster-autoscaler/simulator/drainability"
+	"k8s.io/autoscaler/cluster-autoscaler/simulator/drainability/rules"
 	"k8s.io/autoscaler/cluster-autoscaler/utils/drain"
 	. "k8s.io/autoscaler/cluster-autoscaler/utils/test"
 	"k8s.io/kubernetes/pkg/kubelet/types"
@@ -179,7 +180,7 @@ func TestGetPodsToMove(t *testing.T) {
 		desc         string
 		pods         []*apiv1.Pod
 		pdbs         []*policyv1.PodDisruptionBudget
-		rules        []drainability.Rule
+		rules        []rules.Rule
 		wantPods     []*apiv1.Pod
 		wantDs       []*apiv1.Pod
 		wantBlocking *drain.BlockingPod
@@ -256,19 +257,19 @@ func TestGetPodsToMove(t *testing.T) {
 		{
 			desc:     "Rule allows",
 			pods:     []*apiv1.Pod{unreplicatedPod},
-			rules:    []drainability.Rule{alwaysDrain{}},
+			rules:    []rules.Rule{alwaysDrain{}},
 			wantPods: []*apiv1.Pod{unreplicatedPod},
 		},
 		{
 			desc:     "Second rule allows",
 			pods:     []*apiv1.Pod{unreplicatedPod},
-			rules:    []drainability.Rule{cantDecide{}, alwaysDrain{}},
+			rules:    []rules.Rule{cantDecide{}, alwaysDrain{}},
 			wantPods: []*apiv1.Pod{unreplicatedPod},
 		},
 		{
 			desc:    "Rule blocks",
 			pods:    []*apiv1.Pod{rsPod},
-			rules:   []drainability.Rule{neverDrain{}},
+			rules:   []rules.Rule{neverDrain{}},
 			wantErr: true,
 			wantBlocking: &drain.BlockingPod{
 				Pod:    rsPod,
@@ -278,7 +279,7 @@ func TestGetPodsToMove(t *testing.T) {
 		{
 			desc:    "Second rule blocks",
 			pods:    []*apiv1.Pod{rsPod},
-			rules:   []drainability.Rule{cantDecide{}, neverDrain{}},
+			rules:   []rules.Rule{cantDecide{}, neverDrain{}},
 			wantErr: true,
 			wantBlocking: &drain.BlockingPod{
 				Pod:    rsPod,
@@ -288,7 +289,7 @@ func TestGetPodsToMove(t *testing.T) {
 		{
 			desc:    "Undecisive rule fallback to default logic: Unreplicated pod",
 			pods:    []*apiv1.Pod{unreplicatedPod},
-			rules:   []drainability.Rule{cantDecide{}},
+			rules:   []rules.Rule{cantDecide{}},
 			wantErr: true,
 			wantBlocking: &drain.BlockingPod{
 				Pod:    unreplicatedPod,
@@ -298,7 +299,7 @@ func TestGetPodsToMove(t *testing.T) {
 		{
 			desc:     "Undecisive rule fallback to default logic: Replicated pod",
 			pods:     []*apiv1.Pod{rsPod},
-			rules:    []drainability.Rule{cantDecide{}},
+			rules:    []rules.Rule{cantDecide{}},
 			wantPods: []*apiv1.Pod{rsPod},
 		},
 	}
@@ -326,18 +327,18 @@ func TestGetPodsToMove(t *testing.T) {
 
 type alwaysDrain struct{}
 
-func (a alwaysDrain) Drainable(*apiv1.Pod) drainability.Status {
+func (a alwaysDrain) Drainable(*drainability.DrainContext, *apiv1.Pod) drainability.Status {
 	return drainability.NewDrainableStatus()
 }
 
 type neverDrain struct{}
 
-func (n neverDrain) Drainable(*apiv1.Pod) drainability.Status {
+func (n neverDrain) Drainable(*drainability.DrainContext, *apiv1.Pod) drainability.Status {
 	return drainability.NewBlockedStatus(drain.UnexpectedError, fmt.Errorf("nope"))
 }
 
 type cantDecide struct{}
 
-func (c cantDecide) Drainable(*apiv1.Pod) drainability.Status {
+func (c cantDecide) Drainable(*drainability.DrainContext, *apiv1.Pod) drainability.Status {
 	return drainability.NewUndefinedStatus()
 }

cluster-autoscaler/simulator/drainability/context.go

@@ -0,0 +1,26 @@
+/*
+Copyright 2023 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package drainability
+
+import (
+	policyv1 "k8s.io/api/policy/v1"
+)
+
+// DrainContext contains parameters for drainability rules.
+type DrainContext struct {
+	Pdbs []*policyv1.PodDisruptionBudget
+}

cluster-autoscaler/simulator/drainability/mirror.go → cluster-autoscaler/simulator/drainability/rules/mirror/mirror.go

@@ -14,26 +14,26 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package drainability
+package mirror
 
 import (
-	"k8s.io/autoscaler/cluster-autoscaler/utils/pod"
-
 	apiv1 "k8s.io/api/core/v1"
+	"k8s.io/autoscaler/cluster-autoscaler/simulator/drainability"
+	pod_util "k8s.io/autoscaler/cluster-autoscaler/utils/pod"
 )
 
-// MirrorPodRule is a drainability rule on how to handle mirror pods.
-type MirrorPodRule struct{}
+// Rule is a drainability rule on how to handle mirror pods.
+type Rule struct{}
 
-// NewMirrorPodRule creates a new MirrorPodRule.
-func NewMirrorPodRule() *MirrorPodRule {
-	return &MirrorPodRule{}
+// New creates a new Rule.
+func New() *Rule {
+	return &Rule{}
 }
 
 // Drainable decides what to do with mirror pods on node drain.
-func (m *MirrorPodRule) Drainable(p *apiv1.Pod) Status {
-	if pod.IsMirrorPod(p) {
-		return NewSkipStatus()
+func (Rule) Drainable(drainCtx *drainability.DrainContext, pod *apiv1.Pod) drainability.Status {
+	if pod_util.IsMirrorPod(pod) {
+		return drainability.NewSkipStatus()
 	}
-	return NewUndefinedStatus()
+	return drainability.NewUndefinedStatus()
 }

cluster-autoscaler/simulator/drainability/mirror_test.go → cluster-autoscaler/simulator/drainability/rules/mirror/mirror_test.go

@@ -14,21 +14,22 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package drainability
+package mirror
 
 import (
 	"testing"
 
 	apiv1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/autoscaler/cluster-autoscaler/simulator/drainability"
 	"k8s.io/kubernetes/pkg/kubelet/types"
 )
 
-func TestMirrorPodRule(t *testing.T) {
+func TestRule(t *testing.T) {
 	testCases := []struct {
 		desc string
 		pod  *apiv1.Pod
-		want Status
+		want drainability.Status
 	}{
 		{
 			desc: "non mirror pod",
@@ -38,7 +39,7 @@ func TestMirrorPodRule(t *testing.T) {
 					Namespace: "ns",
 				},
 			},
-			want: NewUndefinedStatus(),
+			want: drainability.NewUndefinedStatus(),
 		},
 		{
 			desc: "mirror pod",
@@ -51,15 +52,14 @@ func TestMirrorPodRule(t *testing.T) {
 					},
 				},
 			},
-			want: NewSkipStatus(),
+			want: drainability.NewSkipStatus(),
 		},
 	}
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
-			m := NewMirrorPodRule()
-			got := m.Drainable(tc.pod)
+			got := New().Drainable(nil, tc.pod)
 			if tc.want != got {
-				t.Errorf("MirrorPodRule.Drainable(%v) = %v, want %v", tc.pod.Name, got, tc.want)
+				t.Errorf("Rule.Drainable(%v) = %v, want %v", tc.pod.Name, got, tc.want)
 			}
 		})
 	}

cluster-autoscaler/simulator/drainability/rules/rules.go

@@ -0,0 +1,57 @@
+/*
+Copyright 2023 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package rules
+
+import (
+	apiv1 "k8s.io/api/core/v1"
+	"k8s.io/autoscaler/cluster-autoscaler/simulator/drainability"
+	"k8s.io/autoscaler/cluster-autoscaler/simulator/drainability/rules/mirror"
+)
+
+// Rule determines whether a given pod can be drained or not.
+type Rule interface {
+	// Drainable determines whether a given pod is drainable according to
+	// the specific Rule.
+	//
+	// DrainContext cannot be nil.
+	Drainable(*drainability.DrainContext, *apiv1.Pod) drainability.Status
+}
+
+// Default returns the default list of Rules.
+func Default() Rules {
+	return []Rule{
+		mirror.New(),
+	}
+}
+
+// Rules defines operations on a collections of rules.
+type Rules []Rule
+
+// Drainable determines whether a given pod is drainable according to the
+// specified set of rules.
+func (rs Rules) Drainable(drainCtx *drainability.DrainContext, pod *apiv1.Pod) drainability.Status {
+	if drainCtx == nil {
+		drainCtx = &drainability.DrainContext{}
+	}
+
+	for _, r := range rs {
+		if d := r.Drainable(drainCtx, pod); d.Outcome != drainability.UndefinedOutcome {
+			return d
+		}
+	}
+	return drainability.NewUndefinedStatus()
+}

cluster-autoscaler/simulator/drainability/rule.go → cluster-autoscaler/simulator/drainability/status.go

@@ -18,8 +18,6 @@ package drainability
 
 import (
 	"k8s.io/autoscaler/cluster-autoscaler/utils/drain"
-
-	apiv1 "k8s.io/api/core/v1"
 )
 
 // OutcomeType identifies the action that should be taken when it comes to
@@ -79,17 +77,3 @@ func NewSkipStatus() Status {
 func NewUndefinedStatus() Status {
 	return Status{}
 }
-
-// Rule determines whether a given pod can be drained or not.
-type Rule interface {
-	// Drainable determines whether a given pod is drainable according to
-	// the specific Rule.
-	Drainable(*apiv1.Pod) Status
-}
-
-// DefaultRules returns the default list of Rules.
-func DefaultRules() []Rule {
-	return []Rule{
-		NewMirrorPodRule(),
-	}
-}