Merge pull request #24923 from mwielgus/drain-refactor

Automatic merge from submit-queue

Small refactoring around drain - move drain logic to separate function.

cc: @piosz @fgrzadkowski

pkg/kubectl/cmd/drain.go

@@ -183,7 +183,7 @@ func (o *DrainOptions) RunDrain() error {
 		return err
 	}
 
-	pods, err := o.GetPodsForDeletion()
+	pods, err := o.getPodsForDeletion()
 	if err != nil {
 		return err
 	}
@@ -195,87 +195,19 @@ func (o *DrainOptions) RunDrain() error {
 	return nil
 }
 
-// GetPodsForDeletion returns all the pods we're going to delete.  If there are
+// getPodsForDeletion returns all the pods we're going to delete.  If there are
 // any unmanaged pods and the user didn't pass --force, we return that list in
 // an error.
-func (o *DrainOptions) GetPodsForDeletion() ([]api.Pod, error) {
+func (o *DrainOptions) getPodsForDeletion() ([]api.Pod, error) {
-	pods := []api.Pod{}
+	pods, unreplicatedPodNames, daemonSetPodNames, err := GetPodsForDeletionOnNodeDrain(
-	podList, err := o.client.Pods(api.NamespaceAll).List(api.ListOptions{FieldSelector: fields.SelectorFromSet(fields.Set{"spec.nodeName": o.nodeInfo.Name})})
+		o.client,
+		o.nodeInfo.Name,
+		o.factory.Decoder(true),
+		o.Force,
+		o.IgnoreDaemonsets,
+	)
 	if err != nil {
-		return pods, err
+		return []api.Pod{}, err
-	}
-	unreplicatedPodNames := []string{}
-	daemonSetPodNames := []string{}
-
-	for _, pod := range podList.Items {
-		_, found := pod.ObjectMeta.Annotations[types.ConfigMirrorAnnotationKey]
-		if found {
-			// Skip mirror pod
-			continue
-		}
-		replicated := false
-		daemonset_pod := false
-
-		creatorRef, found := pod.ObjectMeta.Annotations[controller.CreatedByAnnotation]
-		if found {
-			// Now verify that the specified creator actually exists.
-			var sr api.SerializedReference
-			if err := runtime.DecodeInto(o.factory.Decoder(true), []byte(creatorRef), &sr); err != nil {
-				return pods, err
-			}
-			if sr.Reference.Kind == "ReplicationController" {
-				rc, err := o.client.ReplicationControllers(sr.Reference.Namespace).Get(sr.Reference.Name)
-				// Assume the only reason for an error is because the RC is
-				// gone/missing, not for any other cause.  TODO(mml): something more
-				// sophisticated than this
-				if err == nil && rc != nil {
-					replicated = true
-				}
-			} else if sr.Reference.Kind == "DaemonSet" {
-				ds, err := o.client.DaemonSets(sr.Reference.Namespace).Get(sr.Reference.Name)
-
-				// Assume the only reason for an error is because the DaemonSet is
-				// gone/missing, not for any other cause.  TODO(mml): something more
-				// sophisticated than this
-				if err == nil && ds != nil {
-					// Otherwise, treat daemonset-managed pods as unmanaged since
-					// DaemonSet Controller currently ignores the unschedulable bit.
-					// FIXME(mml): Add link to the issue concerning a proper way to drain
-					// daemonset pods, probably using taints.
-					daemonset_pod = true
-				}
-			} else if sr.Reference.Kind == "Job" {
-				job, err := o.client.ExtensionsClient.Jobs(sr.Reference.Namespace).Get(sr.Reference.Name)
-
-				// Assume the only reason for an error is because the Job is
-				// gone/missing, not for any other cause.  TODO(mml): something more
-				// sophisticated than this
-				if err == nil && job != nil {
-					replicated = true
-				}
-			} else if sr.Reference.Kind == "ReplicaSet" {
-				rs, err := o.client.ExtensionsClient.ReplicaSets(sr.Reference.Namespace).Get(sr.Reference.Name)
-
-				// Assume the only reason for an error is because the RS is
-				// gone/missing, not for any other cause.  TODO(mml): something more
-				// sophisticated than this
-				if err == nil && rs != nil {
-					replicated = true
-				}
-			}
-		}
-
-		switch {
-		case daemonset_pod:
-			daemonSetPodNames = append(daemonSetPodNames, pod.Name)
-		case !replicated:
-			unreplicatedPodNames = append(unreplicatedPodNames, pod.Name)
-			if o.Force {
-				pods = append(pods, pod)
-			}
-		default:
-			pods = append(pods, pod)
-		}
 	}
 
 	daemonSetErrors := !o.IgnoreDaemonsets && len(daemonSetPodNames) > 0
@@ -300,6 +232,92 @@ func (o *DrainOptions) GetPodsForDeletion() ([]api.Pod, error) {
 	return pods, nil
 }
 
+// GetPodsForDeletionOnNodeDrain returns pods that should be deleted on node drain as well as some extra information
+// about possibly problematic pods (unreplicated and deamon sets).
+func GetPodsForDeletionOnNodeDrain(client *client.Client, nodename string, decoder runtime.Decoder, force bool,
+	ignoreDeamonSet bool) (pods []api.Pod, unreplicatedPodNames []string, daemonSetPodNames []string, finalError error) {
+
+	pods = []api.Pod{}
+	unreplicatedPodNames = []string{}
+	daemonSetPodNames = []string{}
+	podList, err := client.Pods(api.NamespaceAll).List(api.ListOptions{FieldSelector: fields.SelectorFromSet(fields.Set{"spec.nodeName": nodename})})
+	if err != nil {
+		return []api.Pod{}, []string{}, []string{}, err
+	}
+
+	for _, pod := range podList.Items {
+		_, found := pod.ObjectMeta.Annotations[types.ConfigMirrorAnnotationKey]
+		if found {
+			// Skip mirror pod
+			continue
+		}
+		replicated := false
+		daemonset_pod := false
+
+		creatorRef, found := pod.ObjectMeta.Annotations[controller.CreatedByAnnotation]
+		if found {
+			// Now verify that the specified creator actually exists.
+			var sr api.SerializedReference
+			if err := runtime.DecodeInto(decoder, []byte(creatorRef), &sr); err != nil {
+				return []api.Pod{}, []string{}, []string{}, err
+			}
+			if sr.Reference.Kind == "ReplicationController" {
+				rc, err := client.ReplicationControllers(sr.Reference.Namespace).Get(sr.Reference.Name)
+				// Assume the only reason for an error is because the RC is
+				// gone/missing, not for any other cause.  TODO(mml): something more
+				// sophisticated than this
+				if err == nil && rc != nil {
+					replicated = true
+				}
+			} else if sr.Reference.Kind == "DaemonSet" {
+				ds, err := client.DaemonSets(sr.Reference.Namespace).Get(sr.Reference.Name)
+
+				// Assume the only reason for an error is because the DaemonSet is
+				// gone/missing, not for any other cause.  TODO(mml): something more
+				// sophisticated than this
+				if err == nil && ds != nil {
+					// Otherwise, treat daemonset-managed pods as unmanaged since
+					// DaemonSet Controller currently ignores the unschedulable bit.
+					// FIXME(mml): Add link to the issue concerning a proper way to drain
+					// daemonset pods, probably using taints.
+					daemonset_pod = true
+				}
+			} else if sr.Reference.Kind == "Job" {
+				job, err := client.ExtensionsClient.Jobs(sr.Reference.Namespace).Get(sr.Reference.Name)
+
+				// Assume the only reason for an error is because the Job is
+				// gone/missing, not for any other cause.  TODO(mml): something more
+				// sophisticated than this
+				if err == nil && job != nil {
+					replicated = true
+				}
+			} else if sr.Reference.Kind == "ReplicaSet" {
+				rs, err := client.ExtensionsClient.ReplicaSets(sr.Reference.Namespace).Get(sr.Reference.Name)
+
+				// Assume the only reason for an error is because the RS is
+				// gone/missing, not for any other cause.  TODO(mml): something more
+				// sophisticated than this
+				if err == nil && rs != nil {
+					replicated = true
+				}
+			}
+		}
+
+		switch {
+		case daemonset_pod:
+			daemonSetPodNames = append(daemonSetPodNames, pod.Name)
+		case !replicated:
+			unreplicatedPodNames = append(unreplicatedPodNames, pod.Name)
+			if force {
+				pods = append(pods, pod)
+			}
+		default:
+			pods = append(pods, pod)
+		}
+	}
+	return pods, unreplicatedPodNames, daemonSetPodNames, nil
+}
+
 // Helper for generating errors or warnings about unmanaged pods.
 func unmanagedMsg(unreplicatedNames []string, daemonSetNames []string, include_guidance bool) string {
 	msgs := []string{}