From 72f141c577670c918eee6e3254cf3beef59aa5db Mon Sep 17 00:00:00 2001
From: Sagar Muchhal <muchhals@vmware.com>
Date: Fri, 4 Dec 2020 15:13:03 -0800
Subject: [PATCH] Passes a context to the drain helper object

Enables the PatchOrReplace call for the cordon helper to accept a
user-supplied context to be used for the client.Patch/Update calls to
alter the node status while cordoning/uncordoning the node.

Signed-off-by: Sagar Muchhal <muchhals@vmware.com>

Kubernetes-commit: 017eaa519de5926fd75a9ddc61bedf2398b69653
---
 pkg/drain/cordon.go  | 10 ++++++++--
 pkg/drain/default.go |  2 +-
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/pkg/drain/cordon.go b/pkg/drain/cordon.go
index cfe2b8b9..006eef76 100644
--- a/pkg/drain/cordon.go
+++ b/pkg/drain/cordon.go
@@ -73,6 +73,12 @@ func (c *CordonHelper) UpdateIfRequired(desired bool) bool {
 // JSON, or if either patch or update calls fail; it will also return a second error
 // whenever creating a patch has failed
 func (c *CordonHelper) PatchOrReplace(clientset kubernetes.Interface, serverDryRun bool) (error, error) {
+	return c.PatchOrReplaceWithContext(context.TODO(), clientset, serverDryRun)
+}
+
+// PatchOrReplaceWithContext provides the option to pass a custom context while updating
+// the node status
+func (c *CordonHelper) PatchOrReplaceWithContext(clientCtx context.Context, clientset kubernetes.Interface, serverDryRun bool) (error, error) {
 	client := clientset.CoreV1().Nodes()
 
 	oldData, err := json.Marshal(c.node)
@@ -93,13 +99,13 @@ func (c *CordonHelper) PatchOrReplace(clientset kubernetes.Interface, serverDryR
 		if serverDryRun {
 			patchOptions.DryRun = []string{metav1.DryRunAll}
 		}
-		_, err = client.Patch(context.TODO(), c.node.Name, types.StrategicMergePatchType, patchBytes, patchOptions)
+		_, err = client.Patch(clientCtx, c.node.Name, types.StrategicMergePatchType, patchBytes, patchOptions)
 	} else {
 		updateOptions := metav1.UpdateOptions{}
 		if serverDryRun {
 			updateOptions.DryRun = []string{metav1.DryRunAll}
 		}
-		_, err = client.Update(context.TODO(), c.node, updateOptions)
+		_, err = client.Update(clientCtx, c.node, updateOptions)
 	}
 	return err, patchErr
 }
diff --git a/pkg/drain/default.go b/pkg/drain/default.go
index a6b005de..43f5dc42 100644
--- a/pkg/drain/default.go
+++ b/pkg/drain/default.go
@@ -57,7 +57,7 @@ func RunCordonOrUncordon(drainer *Helper, node *corev1.Node, desired bool) error
 		return nil
 	}
 
-	err, patchErr := c.PatchOrReplace(drainer.Client, false)
+	err, patchErr := c.PatchOrReplaceWithContext(drainer.Ctx, drainer.Client, false)
 	if err != nil {
 		if patchErr != nil {
 			return fmt.Errorf("cordon error: %s; merge patch error: %s", err.Error(), patchErr.Error())