Merge pull request #6196 from artemvmin/drain-blocking-rule

Add mechanism to override drainability status

cluster-autoscaler/simulator/drain_test.go

@@ -790,18 +790,30 @@ func TestGetPodsToMove(t *testing.T) {
 
 type alwaysDrain struct{}
 
+func (a alwaysDrain) Name() string {
+	return "AlwaysDrain"
+}
+
 func (a alwaysDrain) Drainable(*drainability.DrainContext, *apiv1.Pod) drainability.Status {
 	return drainability.NewDrainableStatus()
 }
 
 type neverDrain struct{}
 
+func (n neverDrain) Name() string {
+	return "NeverDrain"
+}
+
 func (n neverDrain) Drainable(*drainability.DrainContext, *apiv1.Pod) drainability.Status {
 	return drainability.NewBlockedStatus(drain.UnexpectedError, fmt.Errorf("nope"))
 }
 
 type cantDecide struct{}
 
+func (c cantDecide) Name() string {
+	return "CantDecide"
+}
+
 func (c cantDecide) Drainable(*drainability.DrainContext, *apiv1.Pod) drainability.Status {
 	return drainability.NewUndefinedStatus()
 }

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

@@ -30,6 +30,11 @@ func New() *Rule {
 	return &Rule{}
 }
 
+// Name returns the name of the rule.
+func (r *Rule) Name() string {
+	return "DaemonSet"
+}
+
 // Drainable decides what to do with daemon set pods on node drain.
 func (r *Rule) Drainable(drainCtx *drainability.DrainContext, pod *apiv1.Pod) drainability.Status {
 	if pod_util.IsDaemonSetPod(pod) {

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

@@ -19,6 +19,7 @@ package daemonset
 import (
 	"testing"
 
+	"github.com/google/go-cmp/cmp"
 	apiv1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/autoscaler/cluster-autoscaler/simulator/drainability"
@@ -52,8 +53,8 @@ func TestDrainable(t *testing.T) {
 	} {
 		t.Run(desc, func(t *testing.T) {
 			got := New().Drainable(nil, tc.pod)
-			if tc.want != got {
-				t.Errorf("Rule.Drainable(%v) = %v, want %v", tc.pod.Name, got, tc.want)
+			if diff := cmp.Diff(tc.want, got); diff != "" {
+				t.Errorf("Rule.Drainable(%v): got status diff (-want +got):\n%s", tc.pod.Name, diff)
 			}
 		})
 	}

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

@@ -32,6 +32,11 @@ func New() *Rule {
 	return &Rule{}
 }
 
+// Name returns the name of the rule.
+func (r *Rule) Name() string {
+	return "LocalStorage"
+}
+
 // Drainable decides what to do with local storage pods on node drain.
 func (r *Rule) Drainable(drainCtx *drainability.DrainContext, pod *apiv1.Pod) drainability.Status {
 	if drain.HasBlockingLocalStorage(pod) {

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

@@ -30,6 +30,11 @@ func New() *Rule {
 	return &Rule{}
 }
 
+// Name returns the name of the rule.
+func (r *Rule) Name() string {
+	return "LongTerminating"
+}
+
 // Drainable decides what to do with long terminating pods on node drain.
 func (r *Rule) Drainable(drainCtx *drainability.DrainContext, pod *apiv1.Pod) drainability.Status {
 	if drain.IsPodLongTerminating(pod, drainCtx.Timestamp) {

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

@@ -20,6 +20,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/go-cmp/cmp"
 	apiv1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/autoscaler/cluster-autoscaler/simulator/drainability"
@@ -120,8 +121,8 @@ func TestDrainable(t *testing.T) {
 				Timestamp: testTime,
 			}
 			got := New().Drainable(drainCtx, tc.pod)
-			if tc.want != got {
-				t.Errorf("Rule.Drainable(%v) = %v, want %v", tc.pod.Name, got, tc.want)
+			if diff := cmp.Diff(tc.want, got); diff != "" {
+				t.Errorf("Rule.Drainable(%v): got status diff (-want +got):\n%s", tc.pod.Name, diff)
 			}
 		})
 	}

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

@@ -30,6 +30,11 @@ func New() *Rule {
 	return &Rule{}
 }
 
+// Name returns the name of the rule.
+func (r *Rule) Name() string {
+	return "Mirror"
+}
+
 // Drainable decides what to do with mirror pods on node drain.
 func (Rule) Drainable(drainCtx *drainability.DrainContext, pod *apiv1.Pod) drainability.Status {
 	if pod_util.IsMirrorPod(pod) {

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

@@ -19,6 +19,7 @@ package mirror
 import (
 	"testing"
 
+	"github.com/google/go-cmp/cmp"
 	apiv1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/autoscaler/cluster-autoscaler/simulator/drainability"
@@ -54,8 +55,8 @@ func TestDrainable(t *testing.T) {
 	} {
 		t.Run(desc, func(t *testing.T) {
 			got := New().Drainable(nil, tc.pod)
-			if tc.want != got {
-				t.Errorf("Rule.Drainable(%v) = %v, want %v", tc.pod.Name, got, tc.want)
+			if diff := cmp.Diff(tc.want, got); diff != "" {
+				t.Errorf("Rule.Drainable(%v): got status diff (-want +got):\n%s", tc.pod.Name, diff)
 			}
 		})
 	}

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

@@ -32,6 +32,11 @@ func New() *Rule {
 	return &Rule{}
 }
 
+// Name returns the name of the rule.
+func (r *Rule) Name() string {
+	return "NotSafeToEvict"
+}
+
 // Drainable decides what to do with not safe to evict pods on node drain.
 func (Rule) Drainable(drainCtx *drainability.DrainContext, pod *apiv1.Pod) drainability.Status {
 	if drain.HasNotSafeToEvictAnnotation(pod) {

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

@@ -32,6 +32,11 @@ func New() *Rule {
 	return &Rule{}
 }
 
+// Name returns the name of the rule.
+func (r *Rule) Name() string {
+	return "PDB"
+}
+
 // Drainable decides how to handle pods with pdbs on node drain.
 func (Rule) Drainable(drainCtx *drainability.DrainContext, pod *apiv1.Pod) drainability.Status {
 	for _, pdb := range drainCtx.RemainingPdbTracker.MatchingPdbs(pod) {

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

@@ -26,6 +26,8 @@ import (
 	"k8s.io/autoscaler/cluster-autoscaler/core/scaledown/pdb"
 	"k8s.io/autoscaler/cluster-autoscaler/simulator/drainability"
 	"k8s.io/autoscaler/cluster-autoscaler/utils/drain"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestDrainable(t *testing.T) {
@@ -142,9 +144,8 @@ func TestDrainable(t *testing.T) {
 			}
 
 			got := New().Drainable(drainCtx, tc.pod)
-			if got.Outcome != tc.wantOutcome || got.BlockingReason != tc.wantReason {
-				t.Errorf("Rule.Drainable(%s) = (outcome: %v, reason: %v), want (outcome: %v, reason: %v)", tc.pod.Name, got.Outcome, got.BlockingReason, tc.wantOutcome, tc.wantReason)
-			}
+			assert.Equal(t, tc.wantReason, got.BlockingReason)
+			assert.Equal(t, tc.wantOutcome, got.Outcome)
 		})
 	}
 }

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

@@ -38,6 +38,11 @@ func New(minReplicaCount int) *Rule {
 	}
 }
 
+// Name returns the name of the rule.
+func (r *Rule) Name() string {
+	return "ReplicaCount"
+}
+
 // Drainable decides what to do with replicated pods on node drain.
 func (r *Rule) Drainable(drainCtx *drainability.DrainContext, pod *apiv1.Pod) drainability.Status {
 	if drainCtx.Listers == nil {

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

@@ -36,6 +36,11 @@ func New(skipNodesWithCustomControllerPods bool) *Rule {
 	}
 }
 
+// Name returns the name of the rule.
+func (r *Rule) Name() string {
+	return "Replicated"
+}
+
 // Drainable decides what to do with replicated pods on node drain.
 func (r *Rule) Drainable(drainCtx *drainability.DrainContext, pod *apiv1.Pod) drainability.Status {
 	controllerRef := drain.ControllerRef(pod)

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

@@ -32,10 +32,13 @@ import (
 	"k8s.io/autoscaler/cluster-autoscaler/simulator/drainability/rules/system"
 	"k8s.io/autoscaler/cluster-autoscaler/simulator/drainability/rules/terminal"
 	"k8s.io/autoscaler/cluster-autoscaler/simulator/options"
+	"k8s.io/klog/v2"
 )
 
 // Rule determines whether a given pod can be drained or not.
 type Rule interface {
+	// The name of the rule.
+	Name() string
 	// Drainable determines whether a given pod is drainable according to
 	// the specific Rule.
 	//
@@ -86,11 +89,30 @@ func (rs Rules) Drainable(drainCtx *drainability.DrainContext, pod *apiv1.Pod) d
 		drainCtx.RemainingPdbTracker = pdb.NewBasicRemainingPdbTracker()
 	}
 
+	var candidates []overrideCandidate
+
 	for _, r := range rs {
-		d := r.Drainable(drainCtx, pod)
-		if d.Outcome != drainability.UndefinedOutcome {
-			return d
+		status := r.Drainable(drainCtx, pod)
+		if len(status.Overrides) > 0 {
+			candidates = append(candidates, overrideCandidate{r.Name(), status})
+			continue
+		}
+		for _, candidate := range candidates {
+			for _, override := range candidate.status.Overrides {
+				if status.Outcome == override {
+					klog.V(5).Info("Overriding pod %s/%s drainability rule %s with rule %s, outcome %v", pod.GetNamespace(), pod.GetName(), r.Name(), candidate.name, candidate.status.Outcome)
+					return candidate.status
+				}
+			}
+		}
+		if status.Outcome != drainability.UndefinedOutcome {
+			return status
 		}
 	}
 	return drainability.NewUndefinedStatus()
 }
+
+type overrideCandidate struct {
+	name   string
+	status drainability.Status
+}

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

@@ -0,0 +1,132 @@
+/*
+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 (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	apiv1 "k8s.io/api/core/v1"
+	"k8s.io/autoscaler/cluster-autoscaler/simulator/drainability"
+	"k8s.io/autoscaler/cluster-autoscaler/utils/drain"
+)
+
+func TestDrainable(t *testing.T) {
+	for desc, tc := range map[string]struct {
+		rules Rules
+		want  drainability.Status
+	}{
+		"no rules": {
+			want: drainability.NewUndefinedStatus(),
+		},
+		"first non-undefined rule returned": {
+			rules: Rules{
+				fakeRule{drainability.NewUndefinedStatus()},
+				fakeRule{drainability.NewDrainableStatus()},
+				fakeRule{drainability.NewSkipStatus()},
+			},
+			want: drainability.NewDrainableStatus(),
+		},
+		"override match": {
+			rules: Rules{
+				fakeRule{drainability.Status{
+					Outcome:   drainability.DrainOk,
+					Overrides: []drainability.OutcomeType{drainability.BlockDrain},
+				}},
+				fakeRule{drainability.NewBlockedStatus(drain.NotEnoughPdb, nil)},
+			},
+			want: drainability.Status{
+				Outcome:   drainability.DrainOk,
+				Overrides: []drainability.OutcomeType{drainability.BlockDrain},
+			},
+		},
+		"override no match": {
+			rules: Rules{
+				fakeRule{drainability.Status{
+					Outcome:   drainability.DrainOk,
+					Overrides: []drainability.OutcomeType{drainability.SkipDrain},
+				}},
+				fakeRule{drainability.NewBlockedStatus(drain.NotEnoughPdb, nil)},
+			},
+			want: drainability.NewBlockedStatus(drain.NotEnoughPdb, nil),
+		},
+		"override unreachable": {
+			rules: Rules{
+				fakeRule{drainability.NewSkipStatus()},
+				fakeRule{drainability.Status{
+					Outcome:   drainability.DrainOk,
+					Overrides: []drainability.OutcomeType{drainability.BlockDrain},
+				}},
+				fakeRule{drainability.NewBlockedStatus(drain.NotEnoughPdb, nil)},
+			},
+			want: drainability.NewSkipStatus(),
+		},
+		"multiple overrides all run": {
+			rules: Rules{
+				fakeRule{drainability.Status{
+					Outcome:   drainability.DrainOk,
+					Overrides: []drainability.OutcomeType{drainability.SkipDrain},
+				}},
+				fakeRule{drainability.Status{
+					Outcome:   drainability.SkipDrain,
+					Overrides: []drainability.OutcomeType{drainability.BlockDrain},
+				}},
+				fakeRule{drainability.NewBlockedStatus(drain.NotEnoughPdb, nil)},
+			},
+			want: drainability.Status{
+				Outcome:   drainability.SkipDrain,
+				Overrides: []drainability.OutcomeType{drainability.BlockDrain},
+			},
+		},
+		"multiple overrides respects order": {
+			rules: Rules{
+				fakeRule{drainability.Status{
+					Outcome:   drainability.SkipDrain,
+					Overrides: []drainability.OutcomeType{drainability.BlockDrain},
+				}},
+				fakeRule{drainability.Status{
+					Outcome:   drainability.DrainOk,
+					Overrides: []drainability.OutcomeType{drainability.BlockDrain},
+				}},
+				fakeRule{drainability.NewBlockedStatus(drain.NotEnoughPdb, nil)},
+			},
+			want: drainability.Status{
+				Outcome:   drainability.SkipDrain,
+				Overrides: []drainability.OutcomeType{drainability.BlockDrain},
+			},
+		},
+	} {
+		t.Run(desc, func(t *testing.T) {
+			got := tc.rules.Drainable(nil, &apiv1.Pod{})
+			if diff := cmp.Diff(tc.want, got); diff != "" {
+				t.Errorf("Drainable(): got status diff (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
+
+type fakeRule struct {
+	status drainability.Status
+}
+
+func (r fakeRule) Name() string {
+	return "FakeRule"
+}
+
+func (r fakeRule) Drainable(*drainability.DrainContext, *apiv1.Pod) drainability.Status {
+	return r.status
+}

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

@@ -30,6 +30,11 @@ func New() *Rule {
 	return &Rule{}
 }
 
+// Name returns the name of the rule.
+func (r *Rule) Name() string {
+	return "SafeToEvict"
+}
+
 // Drainable decides what to do with safe to evict pods on node drain.
 func (r *Rule) Drainable(drainCtx *drainability.DrainContext, pod *apiv1.Pod) drainability.Status {
 	if drain.HasSafeToEvictAnnotation(pod) {