fix: optimise pod finalizers with merge patch and resourceVersion

workflow/controller/controller.go

@@ -32,11 +32,9 @@ import (
 	"k8s.io/client-go/dynamic"
 	v1 "k8s.io/client-go/informers/core/v1"
 	"k8s.io/client-go/kubernetes"
-	typedv1 "k8s.io/client-go/kubernetes/typed/core/v1"
 	"k8s.io/client-go/rest"
 	"k8s.io/client-go/tools/cache"
 	apiwatch "k8s.io/client-go/tools/watch"
-	"k8s.io/client-go/util/retry"
 	"k8s.io/client-go/util/workqueue"
 
 	"github.com/argoproj/argo-workflows/v3"
@@ -52,6 +50,7 @@ import (
 	"github.com/argoproj/argo-workflows/v3/util/diff"
 	"github.com/argoproj/argo-workflows/v3/util/env"
 	errorsutil "github.com/argoproj/argo-workflows/v3/util/errors"
+	"github.com/argoproj/argo-workflows/v3/util/slice"
 	"github.com/argoproj/argo-workflows/v3/workflow/artifactrepositories"
 	"github.com/argoproj/argo-workflows/v3/workflow/common"
 	controllercache "github.com/argoproj/argo-workflows/v3/workflow/controller/cache"
@@ -153,12 +152,6 @@ type WorkflowController struct {
 	recentCompletions recentCompletions
 }
 
-type PatchOperation struct {
-	Operation string      `json:"op"`
-	Path      string      `json:"path"`
-	Value     interface{} `json:"value,omitempty"`
-}
-
 const (
 	workflowResyncPeriod                = 20 * time.Minute
 	workflowTemplateResyncPeriod        = 20 * time.Minute
@@ -522,6 +515,58 @@ func (wfc *WorkflowController) runPodCleanup(ctx context.Context) {
 	}
 }
 
+func (wfc *WorkflowController) getPodCleanupPatch(pod *apiv1.Pod, labelPodCompleted bool) ([]byte, error) {
+	un := unstructured.Unstructured{}
+	if labelPodCompleted {
+		un.SetLabels(map[string]string{common.LabelKeyCompleted: "true"})
+	}
+
+	finalizerEnabled := os.Getenv("ARGO_POD_STATUS_CAPTURE_FINALIZER") == "true"
+	if finalizerEnabled && slice.ContainsString(pod.Finalizers, common.FinalizerPodStatus) {
+		un.SetFinalizers(slice.RemoveString(pod.Finalizers, common.FinalizerPodStatus))
+		un.SetResourceVersion(pod.ObjectMeta.ResourceVersion)
+	}
+
+	if len(un.Object) == 0 {
+		return nil, nil
+	}
+
+	return un.MarshalJSON()
+}
+
+func (wfc *WorkflowController) patchPodForCleanup(ctx context.Context, namespace, podName string, labelPodCompleted bool) error {
+	pod, err := wfc.getPod(namespace, podName)
+	// err is always nil in all kind of caches for now
+	if err != nil {
+		return err
+	}
+	// if pod is nil, it must have been deleted
+	if pod == nil {
+		return nil
+	}
+
+	patch, err := wfc.getPodCleanupPatch(pod, labelPodCompleted)
+	if err != nil {
+		return err
+	}
+	if patch == nil {
+		return nil
+	}
+
+	_, err = wfc.kubeclientset.CoreV1().Pods(namespace).Patch(
+		ctx,
+		podName,
+		types.MergePatchType,
+		patch,
+		metav1.PatchOptions{},
+	)
+	if err != nil && !apierr.IsNotFound(err) {
+		return err
+	}
+
+	return nil
+}
+
 // all pods will ultimately be cleaned up by either deleting them, or labelling them
 func (wfc *WorkflowController) processNextPodCleanupItem(ctx context.Context) bool {
 	key, quit := wfc.podCleanupQueue.Get()
@@ -538,9 +583,10 @@ func (wfc *WorkflowController) processNextPodCleanupItem(ctx context.Context) bo
 	logCtx := log.WithFields(log.Fields{"key": key, "action": action})
 	logCtx.Info("cleaning up pod")
 	err := func() error {
+		pods := wfc.kubeclientset.CoreV1().Pods(namespace)
 		switch action {
 		case terminateContainers:
-			pod, err := wfc.getPodFromCache(namespace, podName)
+			pod, err := wfc.getPod(namespace, podName)
 			if err == nil && pod != nil && pod.Status.Phase == apiv1.PodPending {
 				wfc.queuePodForCleanup(namespace, podName, deletePod)
 			} else if terminationGracePeriod, err := wfc.signalContainers(ctx, namespace, podName, syscall.SIGTERM); err != nil {
@@ -553,15 +599,14 @@ func (wfc *WorkflowController) processNextPodCleanupItem(ctx context.Context) bo
 				return err
 			}
 		case labelPodCompleted:
-			pods := wfc.kubeclientset.CoreV1().Pods(namespace)
-			if err := wfc.enablePodForDeletion(ctx, pods, namespace, podName, true); err != nil {
+			if err := wfc.patchPodForCleanup(ctx, namespace, podName, true); err != nil {
 				return err
 			}
 		case deletePod:
-			pods := wfc.kubeclientset.CoreV1().Pods(namespace)
-			if err := wfc.enablePodForDeletion(ctx, pods, namespace, podName, false); err != nil {
+			if err := wfc.patchPodForCleanup(ctx, namespace, podName, false); err != nil {
 				return err
 			}
+
 			propagation := metav1.DeletePropagationBackground
 			err := pods.Delete(ctx, podName, metav1.DeleteOptions{
 				PropagationPolicy:  &propagation,
@@ -575,22 +620,14 @@ func (wfc *WorkflowController) processNextPodCleanupItem(ctx context.Context) bo
 	}()
 	if err != nil {
 		logCtx.WithError(err).Warn("failed to clean-up pod")
-		if errorsutil.IsTransientErr(err) {
+		if errorsutil.IsTransientErr(err) || apierr.IsConflict(err) {
 			wfc.podCleanupQueue.AddRateLimited(key)
 		}
 	}
 	return true
 }
 
-func (wfc *WorkflowController) getPodFromAPI(ctx context.Context, namespace string, podName string) (*apiv1.Pod, error) {
-	pod, err := wfc.kubeclientset.CoreV1().Pods(namespace).Get(ctx, podName, metav1.GetOptions{})
-	if err != nil {
-		return nil, err
-	}
-	return pod, nil
-}
-
-func (wfc *WorkflowController) getPodFromCache(namespace string, podName string) (*apiv1.Pod, error) {
+func (wfc *WorkflowController) getPod(namespace string, podName string) (*apiv1.Pod, error) {
 	obj, exists, err := wfc.podInformer.GetStore().GetByKey(namespace + "/" + podName)
 	if err != nil {
 		return nil, err
@@ -605,46 +642,8 @@ func (wfc *WorkflowController) getPodFromCache(namespace string, podName string)
 	return pod, nil
 }
 
-func (wfc *WorkflowController) enablePodForDeletion(ctx context.Context, pods typedv1.PodInterface, namespace string, podName string, isCompleted bool) error {
-	// Get current Pod from K8S and update it to remove finalizer, and if the Pod was completed, set the Label
-	// In the case that the Pod changed in between Get and Update, we'll get a conflict error and can try again
-	err := retry.RetryOnConflict(retry.DefaultRetry, func() error {
-		currentPod, err := wfc.getPodFromAPI(ctx, namespace, podName)
-		if err != nil {
-			return err
-		}
-		updatedPod := currentPod.DeepCopy()
-
-		if isCompleted {
-			if updatedPod.Labels == nil {
-				updatedPod.Labels = make(map[string]string)
-			}
-			updatedPod.Labels[common.LabelKeyCompleted] = "true"
-		}
-
-		updatedPod.Finalizers = removeFinalizer(updatedPod.Finalizers, common.FinalizerPodStatus)
-
-		_, err = pods.Update(ctx, updatedPod, metav1.UpdateOptions{})
-		return err
-	})
-	if err != nil {
-		return err
-	}
-	return nil
-}
-
-func removeFinalizer(finalizers []string, targetFinalizer string) []string {
-	var updatedFinalizers []string
-	for _, finalizer := range finalizers {
-		if finalizer != targetFinalizer {
-			updatedFinalizers = append(updatedFinalizers, finalizer)
-		}
-	}
-	return updatedFinalizers
-}
-
 func (wfc *WorkflowController) signalContainers(ctx context.Context, namespace string, podName string, sig syscall.Signal) (time.Duration, error) {
-	pod, err := wfc.getPodFromCache(namespace, podName)
+	pod, err := wfc.getPod(namespace, podName)
 	if pod == nil || err != nil {
 		return 0, err
 	}
@@ -1054,8 +1053,8 @@ func (wfc *WorkflowController) addWorkflowInformerHandlers(ctx context.Context)
 						log.WithError(err).Error("Failed to list pods")
 					}
 					for _, p := range podList.Items {
-						if err := wfc.enablePodForDeletion(ctx, pods, p.Namespace, p.Name, false); err != nil {
-							log.WithError(err).Error("Failed to enable pod for deletion")
+						if slice.ContainsString(p.Finalizers, common.FinalizerPodStatus) {
+							wfc.queuePodForCleanup(p.Namespace, p.Name, deletePod)
 						}
 					}
 

workflow/controller/controller_test.go

@@ -2,6 +2,8 @@ package controller
 
 import (
 	"context"
+	"encoding/json"
+	"os"
 	"testing"
 	"time"
 
@@ -1256,3 +1258,53 @@ func TestWorkflowWithLongArguments(t *testing.T) {
 	assert.True(t, controller.processNextPodCleanupItem(ctx))
 	assert.Equal(t, wfv1.WorkflowSucceeded, woc.wf.Status.Phase)
 }
+
+func TestGetPodCleaupCache(t *testing.T) {
+	wfc := &WorkflowController{}
+
+	pod := &apiv1.Pod{}
+	pod.SetLabels(map[string]string{common.LabelKeyCompleted: "false"})
+	pod.SetFinalizers([]string{common.FinalizerPodStatus})
+	pod.SetResourceVersion("123456")
+
+	// pod finalizer enabled, patch label
+	os.Setenv("ARGO_POD_STATUS_CAPTURE_FINALIZER", "true")
+	expected := &map[string]interface{}{}
+	err := json.Unmarshal([]byte(`{"metadata":{"resourceVersion":"123456","finalizers":[],"labels":{"workflows.argoproj.io/completed":"true"}}}`), expected)
+	assert.Nil(t, err)
+	patch, err := wfc.getPodCleanupPatch(pod, true)
+	assert.Nil(t, err)
+	actual := &map[string]interface{}{}
+	err = json.Unmarshal(patch, actual)
+	assert.Nil(t, err)
+	assert.Equal(t, expected, actual)
+	assert.Len(t, pod.Finalizers, 1)
+
+	// pod finalizer enabled, not patch label
+	expected = &map[string]interface{}{}
+	err = json.Unmarshal([]byte(`{"metadata":{"resourceVersion":"123456","finalizers":[]}}`), expected)
+	assert.Nil(t, err)
+	patch, err = wfc.getPodCleanupPatch(pod, false)
+	assert.Nil(t, err)
+	actual = &map[string]interface{}{}
+	err = json.Unmarshal(patch, actual)
+	assert.Nil(t, err)
+	assert.Equal(t, expected, actual)
+
+	// pod finalizer disabled, patch both
+	os.Setenv("ARGO_POD_STATUS_CAPTURE_FINALIZER", "false")
+	expected = &map[string]interface{}{}
+	err = json.Unmarshal([]byte(`{"metadata":{"labels":{"workflows.argoproj.io/completed":"true"}}}`), expected)
+	assert.Nil(t, err)
+	patch, err = wfc.getPodCleanupPatch(pod, true)
+	assert.Nil(t, err)
+	actual = &map[string]interface{}{}
+	err = json.Unmarshal(patch, actual)
+	assert.Nil(t, err)
+	assert.Equal(t, expected, actual)
+
+	// pod finalizer disabled, not patch label
+	patch, err = wfc.getPodCleanupPatch(pod, false)
+	assert.Nil(t, err)
+	assert.Nil(t, patch)
+}

workflow/controller/operator.go

@@ -2629,7 +2629,7 @@ func (woc *wfOperationCtx) getPodByNode(node *wfv1.NodeStatus) (*apiv1.Pod, erro
 	}
 
 	podName := woc.getPodName(node.Name, node.TemplateName)
-	return woc.controller.getPodFromCache(woc.wf.GetNamespace(), podName)
+	return woc.controller.getPod(woc.wf.GetNamespace(), podName)
 }
 
 func (woc *wfOperationCtx) recordNodePhaseEvent(node *wfv1.NodeStatus) {