Merge pull request #85574 from mgugino-upstream-stage/kubectl-drain-ctx

kubectl/drain: Add context support

Kubernetes-commit: 3e7c787d3ae8d03de55c1c3e64fd1d41437c6bd4

Godeps/Godeps.json

@@ -604,7 +604,7 @@
 		},
 		{
 			"ImportPath": "k8s.io/client-go",
-			"Rev": "8c19b9f4a642"
+			"Rev": "e9644b2e3edc"
 		},
 		{
 			"ImportPath": "k8s.io/code-generator",

go.mod

@@ -39,7 +39,7 @@ require (
 	k8s.io/api v0.0.0-20191121015604-11707872ac1c
 	k8s.io/apimachinery v0.0.0-20191128180518-03184f823e28
 	k8s.io/cli-runtime v0.0.0-20191204091246-2f9002c43394
-	k8s.io/client-go v0.0.0-20191204082517-8c19b9f4a642
+	k8s.io/client-go v0.0.0-20191204082519-e9644b2e3edc
 	k8s.io/component-base v0.0.0-20191204083903-0d4d24e738e4
 	k8s.io/klog v1.0.0
 	k8s.io/kube-openapi v0.0.0-20191107075043-30be4d16710a
@@ -56,7 +56,7 @@ replace (
 	k8s.io/api => k8s.io/api v0.0.0-20191121015604-11707872ac1c
 	k8s.io/apimachinery => k8s.io/apimachinery v0.0.0-20191128180518-03184f823e28
 	k8s.io/cli-runtime => k8s.io/cli-runtime v0.0.0-20191204091246-2f9002c43394
-	k8s.io/client-go => k8s.io/client-go v0.0.0-20191204082517-8c19b9f4a642
+	k8s.io/client-go => k8s.io/client-go v0.0.0-20191204082519-e9644b2e3edc
 	k8s.io/code-generator => k8s.io/code-generator v0.0.0-20191121015212-c4c8f8345c7e
 	k8s.io/component-base => k8s.io/component-base v0.0.0-20191204083903-0d4d24e738e4
 	k8s.io/metrics => k8s.io/metrics v0.0.0-20191121021546-b1134fd1210c

go.sum

@@ -326,7 +326,7 @@ honnef.co/go/tools v0.0.0-20190106161140-3f1c8253044a/go.mod h1:rf3lG4BRIbNafJWh
 k8s.io/api v0.0.0-20191121015604-11707872ac1c/go.mod h1:R/s4gKT0V/cWEnbQa9taNRJNbWUK57/Dx6cPj6MD3A0=
 k8s.io/apimachinery v0.0.0-20191128180518-03184f823e28/go.mod h1:b9qmWdKlLuU9EBh+06BtLcSf/Mu89rWL33naRxs1uZg=
 k8s.io/cli-runtime v0.0.0-20191204091246-2f9002c43394/go.mod h1:/Blo2zCEC7hipSlcsvbJtx6nRoStWgS9Wd0d4v8Y094=
-k8s.io/client-go v0.0.0-20191204082517-8c19b9f4a642/go.mod h1:HMVIZ0dPop3WCrPEaJ+v5/94cjt56avdDFshpX0Fjvo=
+k8s.io/client-go v0.0.0-20191204082519-e9644b2e3edc/go.mod h1:5lSG1yeDZVwDYAHe9VK48SCe5zmcnkAcf2Mx59TuhmM=
 k8s.io/code-generator v0.0.0-20191121015212-c4c8f8345c7e/go.mod h1:DVmfPQgxQENqDIzVR2ddLXMH34qeszkKSdH/N+s+38s=
 k8s.io/component-base v0.0.0-20191204083903-0d4d24e738e4/go.mod h1:8VIh1jErItC4bg9hLBkPneyS77Tin8KwSzbYepHJnQI=
 k8s.io/gengo v0.0.0-20190128074634-0689ccc1d7d6/go.mod h1:ezvh/TsK7cY6rbqRK0oQQ8IAqLxYwwyPxAX1Pzy0ii0=

pkg/drain/drain.go

@@ -43,6 +43,7 @@ const (
 
 // Helper contains the parameters to control the behaviour of drainer
 type Helper struct {
+	Ctx                 context.Context
 	Client              kubernetes.Interface
 	Force               bool
 	GracePeriodSeconds  int
@@ -203,7 +204,7 @@ func (d *Helper) evictPods(pods []corev1.Pod, policyGroupVersion string, getPodF
 	} else {
 		globalTimeout = d.Timeout
 	}
-	ctx, cancel := context.WithTimeout(context.TODO(), globalTimeout)
+	ctx, cancel := context.WithTimeout(d.getContext(), globalTimeout)
 	defer cancel()
 	for _, pod := range pods {
 		go func(pod corev1.Pod, returnCh chan error) {
@@ -271,7 +272,7 @@ func (d *Helper) deletePods(pods []corev1.Pod, getPodFn func(namespace, name str
 			return err
 		}
 	}
-	ctx := context.TODO()
+	ctx := d.getContext()
 	_, err := waitForDelete(ctx, pods, 1*time.Second, globalTimeout, false, getPodFn, d.OnPodDeletedOrEvicted, globalTimeout)
 	return err
 }
@@ -306,3 +307,13 @@ func waitForDelete(ctx context.Context, pods []corev1.Pod, interval, timeout tim
 	})
 	return pods, err
 }
+
+// Since Helper does not have a constructor, we can't enforce Helper.Ctx != nil
+// Multiple public methods prevent us from initializing the context in a single
+// place as well.
+func (d *Helper) getContext() context.Context {
+	if d.Ctx != nil {
+		return d.Ctx
+	}
+	return context.Background()
+}

pkg/drain/drain_test.go

@@ -46,6 +46,7 @@ func TestDeletePods(t *testing.T) {
 		description       string
 		interval          time.Duration
 		timeout           time.Duration
+		ctxTimeoutEarly   bool
 		expectPendingPods bool
 		expectError       bool
 		expectedError     *error
@@ -91,6 +92,22 @@ func TestDeletePods(t *testing.T) {
 				return nil, fmt.Errorf("%q: not found", name)
 			},
 		},
+		{
+			description:       "Context Canceled",
+			interval:          1000 * time.Millisecond,
+			timeout:           5 * time.Second,
+			ctxTimeoutEarly:   true,
+			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,
@@ -107,14 +124,22 @@ func TestDeletePods(t *testing.T) {
 	for _, test := range tests {
 		t.Run(test.description, func(t *testing.T) {
 			_, pods := createPods(false)
-			ctx := context.TODO()
+			ctx := context.Background()
+			if test.ctxTimeoutEarly {
+				ctx, _ = context.WithTimeout(ctx, 100*time.Millisecond)
+			}
+			start := time.Now()
 			pendingPods, err := waitForDelete(ctx, pods, test.interval, test.timeout, false, test.getPodFn, nil, time.Duration(math.MaxInt64))
-
+			elapsed := time.Since(start)
 			if test.expectError {
 				if err == nil {
 					t.Fatalf("%s: unexpected non-error", test.description)
 				} else if test.expectedError != nil {
-					if *test.expectedError != err {
+					if test.ctxTimeoutEarly {
+						if elapsed >= test.timeout {
+							t.Fatalf("%s: the supplied context did not effectively cancel the waitForDelete", test.description)
+						}
+					} else if *test.expectedError != err {
 						t.Fatalf("%s: the error does not match expected error", test.description)
 					}
 				}