-
Notifications
You must be signed in to change notification settings - Fork 3.2k
/
retry_tweak.go
39 lines (34 loc) · 1.3 KB
/
retry_tweak.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
package controller
import (
apiv1 "k8s.io/api/core/v1"
"k8s.io/utils/env"
wfv1 "github.com/argoproj/argo-workflows/v3/pkg/apis/workflow/v1alpha1"
wfretry "github.com/argoproj/argo-workflows/v3/workflow/util/retry"
)
// RetryTweak is a 2nd order function interface for tweaking the retry
type RetryTweak = func(retryStrategy wfv1.RetryStrategy, nodes wfv1.Nodes, pod *apiv1.Pod)
// FindRetryNode locates the closes retry node ancestor to nodeID
func FindRetryNode(nodes wfv1.Nodes, nodeID string) *wfv1.NodeStatus {
boundaryID := nodes[nodeID].BoundaryID
boundaryNode := nodes[boundaryID]
templateName := boundaryNode.TemplateName
for _, node := range nodes {
if node.Type == wfv1.NodeTypeRetry && node.TemplateName == templateName {
return &node
}
}
return nil
}
// RetryOnDifferentHost append affinity with fail host to pod
func RetryOnDifferentHost(retryNodeName string) RetryTweak {
return func(retryStrategy wfv1.RetryStrategy, nodes wfv1.Nodes, pod *apiv1.Pod) {
if retryStrategy.Affinity == nil {
return
}
hostNames := wfretry.GetFailHosts(nodes, retryNodeName)
hostLabel := env.GetString("RETRY_HOST_NAME_LABEL_KEY", "kubernetes.io/hostname")
if hostLabel != "" && len(hostNames) > 0 {
pod.Spec.Affinity = wfretry.AddHostnamesToAffinity(hostLabel, hostNames, pod.Spec.Affinity)
}
}
}