From e45515c8b3fbc5b8e9bddb1fe36aad4eb2251822 Mon Sep 17 00:00:00 2001
From: jswxstw <jswxstw@gmail.com>
Date: Wed, 3 Jan 2024 19:58:23 +0800
Subject: [PATCH] fix: Apply workflow level PodSpecPatch in agent pod. Fixes
 #12387

Signed-off-by: oninowang <oninowang@tencent.com>
---
 workflow/controller/agent.go      | 21 +++++++++++++++------
 workflow/controller/agent_test.go |  4 ++++
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/workflow/controller/agent.go b/workflow/controller/agent.go
index 86cade8cada7..35552f151c38 100644
--- a/workflow/controller/agent.go
+++ b/workflow/controller/agent.go
@@ -18,6 +18,7 @@ import (
 	wfv1 "github.com/argoproj/argo-workflows/v3/pkg/apis/workflow/v1alpha1"
 	"github.com/argoproj/argo-workflows/v3/util/env"
 	"github.com/argoproj/argo-workflows/v3/workflow/common"
+	"github.com/argoproj/argo-workflows/v3/workflow/util"
 )
 
 func (woc *wfOperationCtx) getAgentPodName() string {
@@ -243,6 +244,14 @@ func (woc *wfOperationCtx) createAgentPod(ctx context.Context) (*apiv1.Pod, erro
 	addSchedulingConstraints(pod, woc.execWf.Spec.DeepCopy(), tmpl)
 	woc.addMetadata(pod, tmpl)
 
+	if woc.execWf.Spec.HasPodSpecPatch() {
+		patchedPodSpec, err := util.ApplyPodSpecPatch(pod.Spec, woc.execWf.Spec.PodSpecPatch)
+		if err != nil {
+			return nil, errors.Wrap(err, "", "Error applying PodSpecPatch")
+		}
+		pod.Spec = *patchedPodSpec
+	}
+
 	if woc.controller.Config.InstanceID != "" {
 		pod.ObjectMeta.Labels[common.LabelKeyControllerInstanceID] = woc.controller.Config.InstanceID
 	}
@@ -293,17 +302,17 @@ func (woc *wfOperationCtx) getExecutorPlugins(ctx context.Context) ([]apiv1.Cont
 }
 
 func addresses(containers []apiv1.Container) []string {
-	var addresses []string
+	var pluginAddresses []string
 	for _, c := range containers {
-		addresses = append(addresses, fmt.Sprintf("http://localhost:%d", c.Ports[0].ContainerPort))
+		pluginAddresses = append(pluginAddresses, fmt.Sprintf("http://localhost:%d", c.Ports[0].ContainerPort))
 	}
-	return addresses
+	return pluginAddresses
 }
 
 func names(containers []apiv1.Container) []string {
-	var addresses []string
+	var pluginNames []string
 	for _, c := range containers {
-		addresses = append(addresses, c.Name)
+		pluginNames = append(pluginNames, c.Name)
 	}
-	return addresses
+	return pluginNames
 }
diff --git a/workflow/controller/agent_test.go b/workflow/controller/agent_test.go
index 04a448806875..e563b318714e 100644
--- a/workflow/controller/agent_test.go
+++ b/workflow/controller/agent_test.go
@@ -20,6 +20,8 @@ metadata:
   name: http-template
   namespace: default
 spec:
+  podSpecPatch: |
+    nodeName: virtual-node
   entrypoint: main
   templates:
     - name: main
@@ -107,6 +109,7 @@ status:
 		for _, pod := range pods.Items {
 			assert.NotNil(t, pod)
 			assert.True(t, strings.HasSuffix(pod.Name, "-agent"))
+			assert.Equal(t, "virtual-node", pod.Spec.NodeName)
 		}
 	})
 	t.Run("CreateTaskSetWithInstanceID", func(t *testing.T) {
@@ -133,6 +136,7 @@ status:
 			assert.NotNil(t, pod)
 			assert.True(t, strings.HasSuffix(pod.Name, "-agent"))
 			assert.Equal(t, "testID", pod.ObjectMeta.Labels[common.LabelKeyControllerInstanceID])
+			assert.Equal(t, "virtual-node", pod.Spec.NodeName)
 		}
 	})
 }