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())