fix UpdatePodWithRetry

KusionStack · Oct 8, 2024 · 2d99853 · 2d99853
1 parent ba768a9
commit 2d99853
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 17 deletions.
diff --git a/pkg/controllers/operationjob/replace/replace.go b/pkg/controllers/operationjob/replace/replace.go
@@ -73,11 +73,12 @@ func (p *PodReplaceHandler) OperateTargets(ctx context.Context, candidates []*Op
 		// label pod to trigger replace
 		replaceTriggered := replaceIndicated || replaceByReplaceUpdate || replaceNewPodExists
 		if !replaceTriggered {
-			// add replace indicate label on origin pod to trigger replace
-			candidate.Pod.Labels[appsv1alpha1.PodReplaceIndicationLabelKey] = "true"
-			// add finalizer on origin pod before trigger replace
-			controllerutil.AddFinalizer(candidate.Pod, OperationJobReplacePodFinalizer)
-			if err := ojutils.UpdatePodWithRetry(ctx, p.client, candidate.Pod); err != nil {
+			if err := ojutils.UpdatePodWithRetry(ctx, p.client, candidate.Pod, func(pod *corev1.Pod) {
+				// add replace indicate label on origin pod to trigger replace
+				candidate.Pod.Labels[appsv1alpha1.PodReplaceIndicationLabelKey] = "true"
+				// add finalizer on origin pod before trigger replace
+				controllerutil.AddFinalizer(candidate.Pod, OperationJobReplacePodFinalizer)
+			}); err != nil {
 				retErr := fmt.Errorf("fail to add %s finalizer or replace indicate label to origin pod %s/%s : %s", OperationJobReplacePodFinalizer, candidate.Pod.Namespace, candidate.Pod.Name, err.Error())
 				ojutils.SetOpsStatusError(candidate, ojutils.ReasonUpdateObjectFailed, retErr.Error())
 				return retErr
@@ -173,11 +174,12 @@ func (p *PodReplaceHandler) ReleaseTargets(ctx context.Context, candidates []*Op
 			return nil
 		}
 
-		// try to remove replace label from origin pod
-		delete(candidate.Pod.Labels, appsv1alpha1.PodReplaceIndicationLabelKey)
-		// try to remove finalizer from origin pod
-		controllerutil.RemoveFinalizer(candidate.Pod, OperationJobReplacePodFinalizer)
-		if err := ojutils.UpdatePodWithRetry(ctx, p.client, candidate.Pod); err != nil {
+		if err := ojutils.UpdatePodWithRetry(ctx, p.client, candidate.Pod, func(pod *corev1.Pod) {
+			// try to remove replace label from origin pod
+			delete(candidate.Pod.Labels, appsv1alpha1.PodReplaceIndicationLabelKey)
+			// try to remove finalizer from origin pod
+			controllerutil.RemoveFinalizer(candidate.Pod, OperationJobReplacePodFinalizer)
+		}); err != nil {
 			retErr := fmt.Errorf("fail to remove %s finalizer or replace indicate label from origin pod %s/%s : %s", OperationJobReplacePodFinalizer, candidate.Pod.Namespace, candidate.Pod.Name, err.Error())
 			ojutils.SetOpsStatusError(candidate, ojutils.ReasonUpdateObjectFailed, retErr.Error())
 			return retErr

diff --git a/pkg/controllers/operationjob/utils/common.go b/pkg/controllers/operationjob/utils/common.go
@@ -88,15 +88,13 @@ func SetOpsStatusError(candidate *opscore.OpsCandidate, reason string, message s
 	}
 }
 
-func UpdatePodWithRetry(ctx context.Context, c client.Client, obj client.Object) error {
+func UpdatePodWithRetry(ctx context.Context, c client.Client, obj client.Object, updateFn func(*corev1.Pod)) error {
 	return retry.RetryOnConflict(retry.DefaultRetry, func() error {
-		var updateErr error
-		if updateErr = c.Update(ctx, obj); updateErr == nil {
-			return nil
-		}
-		if err := c.Get(ctx, types.NamespacedName{Namespace: obj.GetNamespace(), Name: obj.GetName()}, obj); err != nil {
+		pod := &corev1.Pod{}
+		if err := c.Get(context.TODO(), types.NamespacedName{Namespace: obj.GetNamespace(), Name: obj.GetName()}, pod); err != nil {
 			return err
 		}
-		return updateErr
+		updateFn(pod)
+		return c.Update(ctx, obj)
 	})
 }