✨ Add MachineDrainRule "WaitCompleted"

Signed-off-by: Vince Prignano <[email protected]>

api/v1beta1/machinedrainrules_types.go

@@ -27,8 +27,8 @@ const (
 	PodDrainLabel = "cluster.x-k8s.io/drain"
 )
 
-// MachineDrainRuleDrainBehavior defines the drain behavior. Can be either "Drain" or "Skip".
-// +kubebuilder:validation:Enum=Drain;Skip
+// MachineDrainRuleDrainBehavior defines the drain behavior. Can be either "Drain", "Skip", or "WaitCompleted".
+// +kubebuilder:validation:Enum=Drain;Skip;Wait
 type MachineDrainRuleDrainBehavior string
 
 const (
@@ -37,6 +37,10 @@ const (
 
 	// MachineDrainRuleDrainBehaviorSkip means the drain for a Pod should be skipped.
 	MachineDrainRuleDrainBehaviorSkip MachineDrainRuleDrainBehavior = "Skip"
+
+	// MachineDrainRuleDrainBehaviorWaitCompleted means the Pod should not be drained,
+	// but overall drain should wait until the Pod completes.
+	MachineDrainRuleDrainBehaviorWaitCompleted MachineDrainRuleDrainBehavior = "WaitCompleted"
 )
 
 // MachineDrainRuleSpec defines the spec of a MachineDrainRule.

internal/controllers/machine/drain/drain.go

@@ -281,6 +281,7 @@ func (d *Helper) EvictPods(ctx context.Context, podDeleteList *PodDeleteList) Ev
 	var podsToTriggerEvictionLater []PodDelete
 	var podsWithDeletionTimestamp []PodDelete
 	var podsToBeIgnored []PodDelete
+	var podsToWait []PodDelete
 	for _, pod := range podDeleteList.items {
 		switch {
 		case pod.Status.DrainBehavior == clusterv1.MachineDrainRuleDrainBehaviorDrain && pod.Pod.DeletionTimestamp.IsZero():
@@ -289,6 +290,8 @@ func (d *Helper) EvictPods(ctx context.Context, podDeleteList *PodDeleteList) Ev
 			} else {
 				podsToTriggerEvictionLater = append(podsToTriggerEvictionLater, pod)
 			}
+		case pod.Status.DrainBehavior == clusterv1.MachineDrainRuleDrainBehaviorWaitCompleted:
+			podsToWait = append(podsToWait, pod)
 		case pod.Status.DrainBehavior == clusterv1.MachineDrainRuleDrainBehaviorDrain:
 			podsWithDeletionTimestamp = append(podsWithDeletionTimestamp, pod)
 		default:
@@ -394,6 +397,10 @@ evictionLoop:
 		res.PodsToTriggerEvictionLater = append(res.PodsToTriggerEvictionLater, pd.Pod)
 	}
 
+	for _, pd := range podsToWait {
+		res.PodsToWait = append(res.PodsToWait, pd.Pod)
+	}
+
 	return res
 }
 
@@ -431,13 +438,15 @@ type EvictionResult struct {
 	PodsDeletionTimestampSet   []*corev1.Pod
 	PodsFailedEviction         map[string][]*corev1.Pod
 	PodsToTriggerEvictionLater []*corev1.Pod
+	PodsToWait                 []*corev1.Pod
 	PodsNotFound               []*corev1.Pod
 	PodsIgnored                []*corev1.Pod
 }
 
 // DrainCompleted returns if a Node is entirely drained, i.e. if all relevant Pods have gone away.
 func (r EvictionResult) DrainCompleted() bool {
-	return len(r.PodsDeletionTimestampSet) == 0 && len(r.PodsFailedEviction) == 0 && len(r.PodsToTriggerEvictionLater) == 0
+	return len(r.PodsDeletionTimestampSet) == 0 && len(r.PodsFailedEviction) == 0 &&
+		len(r.PodsToTriggerEvictionLater) == 0 && len(r.PodsToWait) == 0
 }
 
 // ConditionMessage returns a condition message for the case where a drain is not completed.
@@ -498,6 +507,10 @@ func (r EvictionResult) ConditionMessage(nodeDrainStartTime *metav1.Time) string
 		conditionMessage = fmt.Sprintf("%s\nAfter above Pods have been removed from the Node, the following Pods will be evicted: %s",
 			conditionMessage, PodListToString(r.PodsToTriggerEvictionLater, 3))
 	}
+	if len(r.PodsToWait) > 0 {
+		conditionMessage = fmt.Sprintf("%s\nWaiting for the following Pods to complete without drain: %s",
+			conditionMessage, PodListToString(r.PodsToWait, 3))
+	}
 	return conditionMessage
 }
 

internal/controllers/machine/drain/filters.go

@@ -61,7 +61,8 @@ type PodDeleteList struct {
 func (l *PodDeleteList) Pods() []*corev1.Pod {
 	pods := []*corev1.Pod{}
 	for _, i := range l.items {
-		if i.Status.DrainBehavior == clusterv1.MachineDrainRuleDrainBehaviorDrain {
+		if i.Status.DrainBehavior == clusterv1.MachineDrainRuleDrainBehaviorDrain ||
+			i.Status.DrainBehavior == clusterv1.MachineDrainRuleDrainBehaviorWaitCompleted {
 			pods = append(pods, i.Pod)
 		}
 	}
@@ -124,6 +125,8 @@ const (
 	PodDeleteStatusTypeOkay = "Okay"
 	// PodDeleteStatusTypeSkip is "Skip".
 	PodDeleteStatusTypeSkip = "Skip"
+	// PodDeleteStatusTypeWaitCompleted is "WaitCompleted".
+	PodDeleteStatusTypeWaitCompleted = "WaitCompleted"
 	// PodDeleteStatusTypeWarning is "Warning".
 	PodDeleteStatusTypeWarning = "Warning"
 	// PodDeleteStatusTypeError is "Error".
@@ -156,6 +159,14 @@ func MakePodDeleteStatusSkip() PodDeleteStatus {
 	}
 }
 
+// MakePodDeleteStatusWaitCompleted is a helper method to return the corresponding PodDeleteStatus.
+func MakePodDeleteStatusWaitCompleted() PodDeleteStatus {
+	return PodDeleteStatus{
+		DrainBehavior: clusterv1.MachineDrainRuleDrainBehaviorWaitCompleted,
+		Reason:        PodDeleteStatusTypeWaitCompleted,
+	}
+}
+
 // MakePodDeleteStatusWithWarning is a helper method to return the corresponding PodDeleteStatus.
 func MakePodDeleteStatusWithWarning(behavior clusterv1.MachineDrainRuleDrainBehavior, message string) PodDeleteStatus {
 	var order *int32
@@ -275,6 +286,11 @@ func (d *Helper) drainLabelFilter(ctx context.Context, pod *corev1.Pod) PodDelet
 		log.V(4).Info(fmt.Sprintf("Skip evicting Pod, because Pod has %s label", clusterv1.PodDrainLabel))
 		return MakePodDeleteStatusSkip()
 	}
+	if labelValue, found := pod.ObjectMeta.Labels[clusterv1.PodDrainLabel]; found && strings.EqualFold(labelValue, string(clusterv1.MachineDrainRuleDrainBehaviorWaitCompleted)) {
+		log := ctrl.LoggerFrom(ctx, "Pod", klog.KObj(pod))
+		log.V(4).Info(fmt.Sprintf("Skip evicting Pod, because Pod has %s label", clusterv1.PodDrainLabel))
+		return MakePodDeleteStatusWaitCompleted()
+	}
 	return MakePodDeleteStatusOkay()
 }
 
@@ -300,6 +316,8 @@ func (d *Helper) machineDrainRulesFilter(machineDrainRules []*clusterv1.MachineD
 				log := ctrl.LoggerFrom(ctx, "Pod", klog.KObj(pod))
 				log.V(4).Info(fmt.Sprintf("Skip evicting Pod, because MachineDrainRule %s with behavior %s applies to the Pod", mdr.Name, clusterv1.MachineDrainRuleDrainBehaviorSkip))
 				return MakePodDeleteStatusSkip()
+			case clusterv1.MachineDrainRuleDrainBehaviorWaitCompleted:
+				return MakePodDeleteStatusWaitCompleted()
 			default:
 				return MakePodDeleteStatusWithError(
 					fmt.Sprintf("MachineDrainRule %q has unknown spec.drain.behavior: %q",

internal/controllers/machine/machine_controller.go

@@ -873,7 +873,6 @@ func (r *Reconciler) drainNode(ctx context.Context, s *scope) (ctrl.Result, erro
 	}
 
 	podsToBeDrained := podDeleteList.Pods()
-
 	if len(podsToBeDrained) == 0 {
 		log.Info("Drain completed")
 		return ctrl.Result{}, nil
@@ -903,6 +902,7 @@ func (r *Reconciler) drainNode(ctx context.Context, s *scope) (ctrl.Result, erro
 		"podsFailedEviction", drain.PodListToString(podsFailedEviction, 5),
 		"podsWithDeletionTimestamp", drain.PodListToString(evictionResult.PodsDeletionTimestampSet, 5),
 		"podsToTriggerEvictionLater", drain.PodListToString(evictionResult.PodsToTriggerEvictionLater, 5),
+		"podsToWaitCompleted", drain.PodListToString(evictionResult.PodsToWait, 5),
 	)
 	return ctrl.Result{RequeueAfter: drainRetryInterval}, nil
 }