Make drain library more reusable

Move more functionality from the kubectl cmd to a package with
fewer dependencies.

Author: justinsb

staging/src/k8s.io/kubectl/pkg/cmd/drain/BUILD

@@ -8,13 +8,10 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//staging/src/k8s.io/api/core/v1:go_default_library",
-        "//staging/src/k8s.io/apimachinery/pkg/api/errors:go_default_library",
-        "//staging/src/k8s.io/apimachinery/pkg/apis/meta/v1:go_default_library",
         "//staging/src/k8s.io/apimachinery/pkg/labels:go_default_library",
         "//staging/src/k8s.io/apimachinery/pkg/runtime/schema:go_default_library",
         "//staging/src/k8s.io/apimachinery/pkg/util/errors:go_default_library",
         "//staging/src/k8s.io/apimachinery/pkg/util/sets:go_default_library",
-        "//staging/src/k8s.io/apimachinery/pkg/util/wait:go_default_library",
         "//staging/src/k8s.io/cli-runtime/pkg/genericclioptions:go_default_library",
         "//staging/src/k8s.io/cli-runtime/pkg/printers:go_default_library",
         "//staging/src/k8s.io/cli-runtime/pkg/resource:go_default_library",
@@ -36,15 +33,11 @@ go_test(
         "//staging/src/k8s.io/api/batch/v1:go_default_library",
         "//staging/src/k8s.io/api/core/v1:go_default_library",
         "//staging/src/k8s.io/api/policy/v1beta1:go_default_library",
-        "//staging/src/k8s.io/apimachinery/pkg/api/errors:go_default_library",
         "//staging/src/k8s.io/apimachinery/pkg/apis/meta/v1:go_default_library",
         "//staging/src/k8s.io/apimachinery/pkg/runtime:go_default_library",
         "//staging/src/k8s.io/apimachinery/pkg/runtime/schema:go_default_library",
-        "//staging/src/k8s.io/apimachinery/pkg/types:go_default_library",
         "//staging/src/k8s.io/apimachinery/pkg/util/strategicpatch:go_default_library",
-        "//staging/src/k8s.io/apimachinery/pkg/util/wait:go_default_library",
         "//staging/src/k8s.io/cli-runtime/pkg/genericclioptions:go_default_library",
-        "//staging/src/k8s.io/cli-runtime/pkg/printers:go_default_library",
         "//staging/src/k8s.io/client-go/rest/fake:go_default_library",
         "//staging/src/k8s.io/kubectl/pkg/cmd/testing:go_default_library",
         "//staging/src/k8s.io/kubectl/pkg/cmd/util:go_default_library",

staging/src/k8s.io/kubectl/pkg/cmd/drain/drain.go

@@ -19,19 +19,14 @@ package drain
 import (
 	"errors"
 	"fmt"
-	"math"
-	"time"
 
 	"github.com/spf13/cobra"
 
 	corev1 "k8s.io/api/core/v1"
-	apierrors "k8s.io/apimachinery/pkg/api/errors"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	utilerrors "k8s.io/apimachinery/pkg/util/errors"
 	"k8s.io/apimachinery/pkg/util/sets"
-	"k8s.io/apimachinery/pkg/util/wait"
 
 	"k8s.io/cli-runtime/pkg/genericclioptions"
 	"k8s.io/cli-runtime/pkg/printers"
@@ -146,14 +141,34 @@ var (
 )
 
 func NewDrainCmdOptions(f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *DrainCmdOptions {
-	return &DrainCmdOptions{
+	o := &DrainCmdOptions{
 		PrintFlags: genericclioptions.NewPrintFlags("drained").WithTypeSetter(scheme.Scheme),
 		IOStreams:  ioStreams,
 		drainer: &drain.Helper{
 			GracePeriodSeconds: -1,
+			Out:                ioStreams.Out,
 			ErrOut:             ioStreams.ErrOut,
 		},
 	}
+	o.drainer.OnPodDeletedOrEvicted = o.onPodDeletedOrEvicted
+	return o
+}
+
+// onPodDeletedOrEvicted is called by drain.Helper, when the pod has been deleted or evicted
+func (o *DrainCmdOptions) onPodDeletedOrEvicted(pod *corev1.Pod, usingEviction bool) {
+	var verbStr string
+	if usingEviction {
+		verbStr = "evicted"
+	} else {
+		verbStr = "deleted"
+	}
+	printObj, err := o.ToPrinter(verbStr)
+	if err != nil {
+		fmt.Fprintf(o.ErrOut, "error building printer: %v\n", err)
+		fmt.Fprintf(o.Out, "pod %s/%s %s\n", pod.Namespace, pod.Name, verbStr)
+	} else {
+		printObj(pod, o.Out)
+	}
 }
 
 func NewCmdDrain(f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *cobra.Command {
@@ -313,7 +328,7 @@ func (o *DrainCmdOptions) deleteOrEvictPodsSimple(nodeInfo *resource.Info) error
 		fmt.Fprintf(o.ErrOut, "WARNING: %s\n", warnings)
 	}
 
-	if err := o.deleteOrEvictPods(list.Pods()); err != nil {
+	if err := o.drainer.DeleteOrEvictPods(list.Pods()); err != nil {
 		pendingList, newErrs := o.drainer.GetPodsForDeletion(nodeInfo.Name)
 
 		fmt.Fprintf(o.ErrOut, "There are pending pods in node %q when an error occurred: %v\n", nodeInfo.Name, err)
@@ -328,136 +343,6 @@ func (o *DrainCmdOptions) deleteOrEvictPodsSimple(nodeInfo *resource.Info) error
 	return nil
 }
 
-// deleteOrEvictPods deletes or evicts the pods on the api server
-func (o *DrainCmdOptions) deleteOrEvictPods(pods []corev1.Pod) error {
-	if len(pods) == 0 {
-		return nil
-	}
-
-	policyGroupVersion, err := drain.CheckEvictionSupport(o.drainer.Client)
-	if err != nil {
-		return err
-	}
-
-	getPodFn := func(namespace, name string) (*corev1.Pod, error) {
-		return o.drainer.Client.CoreV1().Pods(namespace).Get(name, metav1.GetOptions{})
-	}
-
-	if len(policyGroupVersion) > 0 {
-		return o.evictPods(pods, policyGroupVersion, getPodFn)
-	} else {
-		return o.deletePods(pods, getPodFn)
-	}
-}
-
-func (o *DrainCmdOptions) evictPods(pods []corev1.Pod, policyGroupVersion string, getPodFn func(namespace, name string) (*corev1.Pod, error)) error {
-	returnCh := make(chan error, 1)
-
-	for _, pod := range pods {
-		go func(pod corev1.Pod, returnCh chan error) {
-			for {
-				fmt.Fprintf(o.Out, "evicting pod %q\n", pod.Name)
-				err := o.drainer.EvictPod(pod, policyGroupVersion)
-				if err == nil {
-					break
-				} else if apierrors.IsNotFound(err) {
-					returnCh <- nil
-					return
-				} else if apierrors.IsTooManyRequests(err) {
-					fmt.Fprintf(o.ErrOut, "error when evicting pod %q (will retry after 5s): %v\n", pod.Name, err)
-					time.Sleep(5 * time.Second)
-				} else {
-					returnCh <- fmt.Errorf("error when evicting pod %q: %v", pod.Name, err)
-					return
-				}
-			}
-			_, err := o.waitForDelete([]corev1.Pod{pod}, 1*time.Second, time.Duration(math.MaxInt64), true, getPodFn)
-			if err == nil {
-				returnCh <- nil
-			} else {
-				returnCh <- fmt.Errorf("error when waiting for pod %q terminating: %v", pod.Name, err)
-			}
-		}(pod, returnCh)
-	}
-
-	doneCount := 0
-	var errors []error
-
-	// 0 timeout means infinite, we use MaxInt64 to represent it.
-	var globalTimeout time.Duration
-	if o.drainer.Timeout == 0 {
-		globalTimeout = time.Duration(math.MaxInt64)
-	} else {
-		globalTimeout = o.drainer.Timeout
-	}
-	globalTimeoutCh := time.After(globalTimeout)
-	numPods := len(pods)
-	for doneCount < numPods {
-		select {
-		case err := <-returnCh:
-			doneCount++
-			if err != nil {
-				errors = append(errors, err)
-			}
-		case <-globalTimeoutCh:
-			return fmt.Errorf("drain did not complete within %v", globalTimeout)
-		}
-	}
-	return utilerrors.NewAggregate(errors)
-}
-
-func (o *DrainCmdOptions) deletePods(pods []corev1.Pod, getPodFn func(namespace, name string) (*corev1.Pod, error)) error {
-	// 0 timeout means infinite, we use MaxInt64 to represent it.
-	var globalTimeout time.Duration
-	if o.drainer.Timeout == 0 {
-		globalTimeout = time.Duration(math.MaxInt64)
-	} else {
-		globalTimeout = o.drainer.Timeout
-	}
-	for _, pod := range pods {
-		err := o.drainer.DeletePod(pod)
-		if err != nil && !apierrors.IsNotFound(err) {
-			return err
-		}
-	}
-	_, err := o.waitForDelete(pods, 1*time.Second, globalTimeout, false, getPodFn)
-	return err
-}
-
-func (o *DrainCmdOptions) waitForDelete(pods []corev1.Pod, interval, timeout time.Duration, usingEviction bool, getPodFn func(string, string) (*corev1.Pod, error)) ([]corev1.Pod, error) {
-	var verbStr string
-	if usingEviction {
-		verbStr = "evicted"
-	} else {
-		verbStr = "deleted"
-	}
-	printObj, err := o.ToPrinter(verbStr)
-	if err != nil {
-		return pods, err
-	}
-
-	err = wait.PollImmediate(interval, timeout, func() (bool, error) {
-		pendingPods := []corev1.Pod{}
-		for i, pod := range pods {
-			p, err := getPodFn(pod.Namespace, pod.Name)
-			if apierrors.IsNotFound(err) || (p != nil && p.ObjectMeta.UID != pod.ObjectMeta.UID) {
-				printObj(&pod, o.Out)
-				continue
-			} else if err != nil {
-				return false, err
-			} else {
-				pendingPods = append(pendingPods, pods[i])
-			}
-		}
-		pods = pendingPods
-		if len(pendingPods) > 0 {
-			return false, nil
-		}
-		return true, nil
-	})
-	return pods, err
-}
-
 // RunCordonOrUncordon runs either Cordon or Uncordon.  The desired value for
 // "Unschedulable" is passed as the first arg.
 func (o *DrainCmdOptions) RunCordonOrUncordon(desired bool) error {

staging/src/k8s.io/kubectl/pkg/cmd/drain/drain_test.go

@@ -17,15 +17,11 @@ limitations under the License.
 package drain
 
 import (
-	"errors"
-	"fmt"
-	"io"
 	"io/ioutil"
 	"net/http"
 	"net/url"
 	"os"
 	"reflect"
-	"strconv"
 	"strings"
 	"sync/atomic"
 	"testing"
@@ -38,14 +34,10 @@ import (
 	batchv1 "k8s.io/api/batch/v1"
 	corev1 "k8s.io/api/core/v1"
 	policyv1beta1 "k8s.io/api/policy/v1beta1"
-	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/runtime/schema"
-	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apimachinery/pkg/util/strategicpatch"
-	"k8s.io/apimachinery/pkg/util/wait"
-	"k8s.io/cli-runtime/pkg/printers"
 	"k8s.io/client-go/rest/fake"
 	cmdtesting "k8s.io/kubectl/pkg/cmd/testing"
 	cmdutil "k8s.io/kubectl/pkg/cmd/util"
@@ -907,135 +899,6 @@ func TestDrain(t *testing.T) {
 	}
 }
 
-func TestDeletePods(t *testing.T) {
-	ifHasBeenCalled := map[string]bool{}
-	tests := []struct {
-		description       string
-		interval          time.Duration
-		timeout           time.Duration
-		expectPendingPods bool
-		expectError       bool
-		expectedError     *error
-		getPodFn          func(namespace, name string) (*corev1.Pod, error)
-	}{
-		{
-			description:       "Wait for deleting to complete",
-			interval:          100 * time.Millisecond,
-			timeout:           10 * time.Second,
-			expectPendingPods: false,
-			expectError:       false,
-			expectedError:     nil,
-			getPodFn: func(namespace, name string) (*corev1.Pod, error) {
-				oldPodMap, _ := createPods(false)
-				newPodMap, _ := createPods(true)
-				if oldPod, found := oldPodMap[name]; found {
-					if _, ok := ifHasBeenCalled[name]; !ok {
-						ifHasBeenCalled[name] = true
-						return &oldPod, nil
-					}
-					if oldPod.ObjectMeta.Generation < 4 {
-						newPod := newPodMap[name]
-						return &newPod, nil
-					}
-					return nil, apierrors.NewNotFound(schema.GroupResource{Resource: "pods"}, name)
-
-				}
-				return nil, apierrors.NewNotFound(schema.GroupResource{Resource: "pods"}, name)
-			},
-		},
-		{
-			description:       "Deleting could timeout",
-			interval:          200 * time.Millisecond,
-			timeout:           3 * time.Second,
-			expectPendingPods: true,
-			expectError:       true,
-			expectedError:     &wait.ErrWaitTimeout,
-			getPodFn: func(namespace, name string) (*corev1.Pod, error) {
-				oldPodMap, _ := createPods(false)
-				if oldPod, found := oldPodMap[name]; found {
-					return &oldPod, nil
-				}
-				return nil, fmt.Errorf("%q: not found", name)
-			},
-		},
-		{
-			description:       "Client error could be passed out",
-			interval:          200 * time.Millisecond,
-			timeout:           5 * time.Second,
-			expectPendingPods: true,
-			expectError:       true,
-			expectedError:     nil,
-			getPodFn: func(namespace, name string) (*corev1.Pod, error) {
-				return nil, errors.New("This is a random error for testing")
-			},
-		},
-	}
-
-	for _, test := range tests {
-		t.Run(test.description, func(t *testing.T) {
-			tf := cmdtesting.NewTestFactory()
-			defer tf.Cleanup()
-
-			o := DrainCmdOptions{
-				PrintFlags: genericclioptions.NewPrintFlags("drained").WithTypeSetter(scheme.Scheme),
-			}
-			o.Out = os.Stdout
-
-			o.ToPrinter = func(operation string) (printers.ResourcePrinterFunc, error) {
-				return func(obj runtime.Object, out io.Writer) error {
-					return nil
-				}, nil
-			}
-
-			_, pods := createPods(false)
-			pendingPods, err := o.waitForDelete(pods, test.interval, test.timeout, false, test.getPodFn)
-
-			if test.expectError {
-				if err == nil {
-					t.Fatalf("%s: unexpected non-error", test.description)
-				} else if test.expectedError != nil {
-					if *test.expectedError != err {
-						t.Fatalf("%s: the error does not match expected error", test.description)
-					}
-				}
-			}
-			if !test.expectError && err != nil {
-				t.Fatalf("%s: unexpected error", test.description)
-			}
-			if test.expectPendingPods && len(pendingPods) == 0 {
-				t.Fatalf("%s: unexpected empty pods", test.description)
-			}
-			if !test.expectPendingPods && len(pendingPods) > 0 {
-				t.Fatalf("%s: unexpected pending pods", test.description)
-			}
-		})
-	}
-}
-
-func createPods(ifCreateNewPods bool) (map[string]corev1.Pod, []corev1.Pod) {
-	podMap := make(map[string]corev1.Pod)
-	podSlice := []corev1.Pod{}
-	for i := 0; i < 8; i++ {
-		var uid types.UID
-		if ifCreateNewPods {
-			uid = types.UID(i)
-		} else {
-			uid = types.UID(strconv.Itoa(i) + strconv.Itoa(i))
-		}
-		pod := corev1.Pod{
-			ObjectMeta: metav1.ObjectMeta{
-				Name:       "pod" + strconv.Itoa(i),
-				Namespace:  "default",
-				UID:        uid,
-				Generation: int64(i),
-			},
-		}
-		podMap[pod.Name] = pod
-		podSlice = append(podSlice, pod)
-	}
-	return podMap, podSlice
-}
-
 type MyReq struct {
 	Request *http.Request
 }