Add mechanism to override drainability status

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/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_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/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/rules.go

@@ -86,10 +86,23 @@ func (rs Rules) Drainable(drainCtx *drainability.DrainContext, pod *apiv1.Pod) d
 		drainCtx.RemainingPdbTracker = pdb.NewBasicRemainingPdbTracker()
 	}
 
+	var candidates []drainability.Status
+
 	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, status)
+			continue
+		}
+		for _, candidate := range candidates {
+			for _, override := range candidate.Overrides {
+				if status.Outcome == override {
+					return candidate
+				}
+			}
+		}
+		if status.Outcome != drainability.UndefinedOutcome {
+			return status
 		}
 	}
 	return drainability.NewUndefinedStatus()

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

@@ -0,0 +1,128 @@
+/*
+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, nil)
+			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) Drainable(*drainability.DrainContext, *apiv1.Pod) drainability.Status {
+	return r.status
+}

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

@@ -19,6 +19,7 @@ package safetoevict
 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/terminal/rule_test.go

@@ -19,6 +19,7 @@ package terminal
 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"
@@ -71,8 +72,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/status.go

@@ -43,6 +43,12 @@ type Status struct {
 	// Outcome indicates what can happen when it comes to draining a
 	// specific pod.
 	Outcome OutcomeType
+	// Overrides specifies Outcomes that should be trumped by this Status.
+	// If Overrides is empty, this Status is returned immediately.
+	// If Overrides is non-empty, we continue running the remaining Rules. If a
+	// Rule is encountered that matches one of the Outcomes specified by this
+	// field, this Status will will be returned instead.
+	Overrides []OutcomeType
 	// Reason contains the reason why a pod is blocking node drain. It is
 	// set only when Outcome is BlockDrain.
 	BlockingReason drain.BlockingPodReason