Readiness gates implementation for eager mode

api/v1alpha1/rollingupgrade_types.go

@@ -52,6 +52,12 @@ type RollingUpgradeSpec struct {
 	IgnoreDrainFailures bool `json:"ignoreDrainFailures,omitempty"`
 	// ForceRefresh enables draining and terminating the node even if the launch config/template hasn't changed.
 	ForceRefresh bool `json:"forceRefresh,omitempty"`
+	// ReadinessGates allow to specify label selectors that node must match to be considered ready.
+	ReadinessGates []NodeReadinessGate `json:"readinessGates,omitempty"`
+}
+
+type NodeReadinessGate struct {
+	MatchLabels map[string]string `json:"matchLabels,omitempty" protobuf:"bytes,1,rep,name=matchLabels"`
 }
 
 // RollingUpgradeStatus defines the observed state of RollingUpgrade

config/crd/bases/upgrademgr.keikoproj.io_rollingupgrades.yaml

@@ -92,6 +92,17 @@ spec:
                 script:
                   type: string
               type: object
+            readinessGates:
+              description: ReadinessGates allow to specify label selectors that node
+                must match to be considered ready.
+              items:
+                properties:
+                  matchLabels:
+                    additionalProperties:
+                      type: string
+                    type: object
+                type: object
+              type: array
             strategy:
               description: UpdateStrategy holds the information needed to perform
                 update based on different update strategies

controllers/helpers.go

@@ -50,6 +50,21 @@ func isNodeReady(node corev1.Node) bool {
 	return false
 }
 
+func IsNodePassesReadinessGates(node corev1.Node, requiredReadinessGates []upgrademgrv1alpha1.NodeReadinessGate) bool {
+
+	if len(requiredReadinessGates) == 0 {
+		return true
+	}
+	for _, gate := range requiredReadinessGates {
+		for key, value := range gate.MatchLabels {
+			if node.Labels[key] != value {
+				return false
+			}
+		}
+	}
+	return true
+}
+
 func getInServiceCount(instances []*autoscaling.Instance) int64 {
 	var count int64
 	for _, instance := range instances {

controllers/helpers_test.go

@@ -3,6 +3,7 @@ package controllers
 import (
 	"encoding/json"
 	"fmt"
+	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	"github.com/aws/aws-sdk-go/aws"
 
@@ -45,6 +46,105 @@ func TestIsNodeReady(t *testing.T) {
 	}
 }
 
+func TestIsNodePassesReadinessGates(t *testing.T) {
+	g := gomega.NewGomegaWithT(t)
+
+	type test struct {
+		gate   []map[string]string
+		labels map[string]string
+		want   bool
+	}
+	tests := []test{
+		{
+			gate: []map[string]string{
+				{
+					"healthy": "true",
+				},
+			},
+			labels: map[string]string{
+				"healthy": "true",
+			},
+			want: true,
+		},
+
+		{
+			gate: []map[string]string{},
+			labels: map[string]string{
+				"healthy": "true",
+			},
+			want: true,
+		},
+
+		{
+			gate: []map[string]string{
+				{"healthy": "true"},
+			},
+			labels: map[string]string{
+				"healthy": "false",
+			},
+			want: false,
+		},
+
+		{
+			gate: []map[string]string{
+				{"healthy": "true"},
+			},
+			labels: map[string]string{
+
+			},
+			want: false,
+		},
+
+		{
+			gate: []map[string]string{
+				{"healthy": "true"},
+				{"second-check": "true"},
+			},
+			labels: map[string]string{
+				"healthy": "true",
+			},
+			want: false,
+		},
+		{
+			gate: []map[string]string{
+				{"healthy": "true"},
+				{"second-check": "true"},
+			},
+			labels: map[string]string{
+				"healthy": "true",
+				"second-check": "true",
+			},
+			want: true,
+		},
+		{
+			gate: []map[string]string{
+				{"healthy": "true"},
+				{"second-check": "true"},
+			},
+			labels: map[string]string{
+				"healthy": "true",
+				"second-check": "false",
+			},
+			want: false,
+		},	}
+
+	for _, tt := range tests {
+		readinessGates := make([]upgrademgrv1alpha1.NodeReadinessGate, len(tt.gate))
+		for i, g := range tt.gate {
+			readinessGates[i] = upgrademgrv1alpha1.NodeReadinessGate{
+				MatchLabels: g,
+			}
+		}
+		node := corev1.Node{
+			ObjectMeta: v1.ObjectMeta{
+				Labels: tt.labels,
+			},
+		}
+		g.Expect(IsNodePassesReadinessGates(node, readinessGates)).To(gomega.Equal(tt.want))
+	}
+
+}
+
 func TestGetInServiceCount(t *testing.T) {
 	g := gomega.NewGomegaWithT(t)
 	tt := map[*autoscaling.Instance]int64{

controllers/rollingupgrade_controller.go

@@ -308,7 +308,7 @@ func (r *RollingUpgradeReconciler) WaitForDesiredNodes(ruObj *upgrademgrv1alpha1
 		for _, node := range r.NodeList.Items {
 			tokens := strings.Split(node.Spec.ProviderID, "/")
 			instanceID := tokens[len(tokens)-1]
-			if contains(inServiceInstances, instanceID) && isNodeReady(node) {
+			if contains(inServiceInstances, instanceID) && isNodeReady(node) && IsNodePassesReadinessGates(node, ruObj.Spec.ReadinessGates) {
 				foundCount++
 			}
 		}