feat: allow last retry variables in expressions. Fixes #10364

docs/variables.md

@@ -222,6 +222,7 @@ Note: These variables evaluate to a string type. If using advanced expressions,
 |----------|------------|
 | `pod.name` | Pod name of the container/script |
 | `retries` | The retry number of the container/script if `retryStrategy` is specified |
+| `lastRetry` | The last retry is a structure that contains the fields `exitCode`, `status`, `duration` and `message` of the last retry |
 | `inputs.artifacts.<NAME>.path` | Local path of the input artifact |
 | `outputs.artifacts.<NAME>.path` | Local path of the output artifact |
 | `outputs.parameters.<NAME>.path` | Local path of the output parameter |

util/template/expression_template.go

@@ -21,6 +21,15 @@ func init() {
 	}
 }
 
+func anyVarNotInEnv(expression string, variables []string, env map[string]interface{}) bool {
+	for _, variable := range variables {
+		if hasVariableInExpression(expression, variable) && !hasVarInEnv(env, variable) {
+			return true
+		}
+	}
+	return false
+}
+
 func expressionReplace(w io.Writer, expression string, env map[string]interface{}, allowUnresolved bool) (int, error) {
 	// The template is JSON-marshaled. This JSON-unmarshals the expression to undo any character escapes.
 	var unmarshalledExpression string
@@ -33,21 +42,28 @@ func expressionReplace(w io.Writer, expression string, env map[string]interface{
 		return 0, fmt.Errorf("failed to unmarshall JSON expression: %w", err)
 	}
 
-	if _, ok := env["retries"]; !ok && hasRetries(unmarshalledExpression) && allowUnresolved {
+	if anyVarNotInEnv(unmarshalledExpression, []string{"retries"}, env) && allowUnresolved {
 		// this is to make sure expressions like `sprig.int(retries)` don't get resolved to 0 when `retries` don't exist in the env
 		// See https://github.com/argoproj/argo-workflows/issues/5388
 		log.WithError(err).Debug("Retries are present and unresolved is allowed")
 		return fmt.Fprintf(w, "{{%s%s}}", kindExpression, expression)
 	}
 
+	lastRetryVariables := []string{"lastRetry.exitCode", "lastRetry.status", "lastRetry.duration", "lastRetry.message"}
+	if anyVarNotInEnv(unmarshalledExpression, lastRetryVariables, env) && allowUnresolved {
+		// This is to make sure expressions which contains `lastRetry.*` don't get resolved to nil
+		// when they don't exist in the env.
+		log.WithError(err).Debug("LastRetry variables are present and unresolved is allowed")
+		return fmt.Fprintf(w, "{{%s%s}}", kindExpression, expression)
+	}
+
 	// This is to make sure expressions which contains `workflow.status` and `work.failures` don't get resolved to nil
 	// when `workflow.status` and `workflow.failures` don't exist in the env.
 	// See https://github.com/argoproj/argo-workflows/issues/10393, https://github.com/expr-lang/expr/issues/330
 	// This issue doesn't happen to other template parameters since `workflow.status` and `workflow.failures` only exist in the env
 	// when the exit handlers complete.
-	if ((hasWorkflowStatus(unmarshalledExpression) && !hasVarInEnv(env, "workflow.status")) ||
-		(hasWorkflowFailures(unmarshalledExpression) && !hasVarInEnv(env, "workflow.failures"))) &&
-		allowUnresolved {
+	if anyVarNotInEnv(unmarshalledExpression, []string{"workflow.status", "workflow.failures"}, env) && allowUnresolved {
+		log.WithError(err).Debug("workflow.status or workflow.failures are present and unresolved is allowed")
 		return fmt.Fprintf(w, "{{%s%s}}", kindExpression, expression)
 	}
 
@@ -102,54 +118,58 @@ func EnvMap(replaceMap map[string]string) map[string]interface{} {
 
 // hasRetries checks if the variable `retries` exists in the expression template
 func hasRetries(expression string) bool {
-	tokens, err := lexer.Lex(file.NewSource(expression))
-	if err != nil {
+	return hasVariableInExpression(expression, "retries")
+}
+
+func searchTokens(haystack []lexer.Token, needle []lexer.Token) bool {
+	if len(needle) > len(haystack) {
 		return false
 	}
-	for _, token := range tokens {
-		if token.Kind == lexer.Identifier && token.Value == "retries" {
-			return true
+	if len(needle) == 0 {
+		return true
+	}
+outer:
+	for i := 0; i <= len(haystack)-len(needle); i++ {
+		for j := 0; j < len(needle); j++ {
+			if haystack[i+j].String() != needle[j].String() {
+				continue outer
+			}
 		}
+		return true
 	}
 	return false
 }
 
-// hasWorkflowStatus checks if expression contains `workflow.status`
-func hasWorkflowStatus(expression string) bool {
-	if !strings.Contains(expression, "workflow.status") {
-		return false
-	}
-	// Even if the expression contains `workflow.status`, it could be the case that it represents a string (`"workflow.status"`),
-	// not a variable, so we need to parse it and handle filter the string case.
-	tokens, err := lexer.Lex(file.NewSource(expression))
-	if err != nil {
-		return false
-	}
-	for i := 0; i < len(tokens)-2; i++ {
-		if tokens[i].Value+tokens[i+1].Value+tokens[i+2].Value == "workflow.status" {
-			return true
+func filterEOF(toks []lexer.Token) []lexer.Token {
+	newToks := []lexer.Token{}
+	for _, tok := range toks {
+		if tok.Kind != lexer.EOF {
+			newToks = append(newToks, tok)
 		}
 	}
-	return false
+	return newToks
 }
 
-// hasWorkflowFailures checks if expression contains `workflow.failures`
-func hasWorkflowFailures(expression string) bool {
-	if !strings.Contains(expression, "workflow.failures") {
+// hasVariableInExpression checks if an expression contains a variable.
+// This function is somewhat cursed and I have attempted my best to
+// remove this curse, but it still exists.
+// The strings.Contains is needed because the lexer doesn't do
+// any whitespace processing (workflow .status will be seen as workflow.status)
+func hasVariableInExpression(expression, variable string) bool {
+	if !strings.Contains(expression, variable) {
 		return false
 	}
-	// Even if the expression contains `workflow.failures`, it could be the case that it represents a string (`"workflow.failures"`),
-	// not a variable, so we need to parse it and handle filter the string case.
 	tokens, err := lexer.Lex(file.NewSource(expression))
 	if err != nil {
 		return false
 	}
-	for i := 0; i < len(tokens)-2; i++ {
-		if tokens[i].Value+tokens[i+1].Value+tokens[i+2].Value == "workflow.failures" {
-			return true
-		}
+	variableTokens, err := lexer.Lex(file.NewSource(variable))
+	if err != nil {
+		return false
 	}
-	return false
+	variableTokens = filterEOF(variableTokens)
+
+	return searchTokens(tokens, variableTokens)
 }
 
 // hasVarInEnv checks if a parameter is in env or not

util/template/expression_template_test.go

@@ -3,6 +3,8 @@ package template
 import (
 	"testing"
 
+	"github.com/expr-lang/expr/file"
+	"github.com/expr-lang/expr/parser/lexer"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -18,19 +20,18 @@ func Test_hasRetries(t *testing.T) {
 }
 
 func Test_hasWorkflowParameters(t *testing.T) {
-	t.Run("hasWorkflowStatusInExpression", func(t *testing.T) {
-		assert.True(t, hasWorkflowStatus("workflow.status"))
-		assert.True(t, hasWorkflowStatus(`workflow.status == "Succeeded" ? "SUCCESSFUL" : "FAILED"`))
-		assert.False(t, hasWorkflowStatus(`"workflow.status" == "Succeeded" ? "SUCCESSFUL" : "FAILED"`))
-		assert.False(t, hasWorkflowStatus("workflow status"))
-		assert.False(t, hasWorkflowStatus("workflow .status"))
-	})
+	t.Run("hasVariableInExpression", func(t *testing.T) {
+		expression := `workflow.status == "Succeeded" ? "SUCCESSFUL" : "FAILED"`
+		exprToks, _ := lexer.Lex(file.NewSource(expression))
+		variableToks, _ := lexer.Lex(file.NewSource("workflow.status"))
+		variableToks = filterEOF(variableToks)
+		assert.True(t, searchTokens(exprToks, variableToks))
+		assert.True(t, hasVariableInExpression(expression, "workflow.status"))
+
+		assert.False(t, hasVariableInExpression(expression, "workflow status"))
+		assert.False(t, hasVariableInExpression(expression, "workflow .status"))
 
-	t.Run("hasWorkflowFailuresInExpression", func(t *testing.T) {
-		assert.True(t, hasWorkflowFailures("workflow.failures"))
-		assert.True(t, hasWorkflowFailures(`workflow.failures == "Succeeded" ? "SUCCESSFUL" : "FAILED"`))
-		assert.False(t, hasWorkflowFailures(`"workflow.failures" == "Succeeded" ? "SUCCESSFUL" : "FAILED"`))
-		assert.False(t, hasWorkflowFailures("workflow failures"))
-		assert.False(t, hasWorkflowFailures("workflow .failures"))
+		expression = `"workflow.status" == "Succeeded" ? "SUCCESSFUL" : "FAILED"`
+		assert.False(t, hasVariableInExpression(expression, "workflow .status"))
 	})
 }

workflow/controller/operator.go

@@ -2221,6 +2221,23 @@ func (woc *wfOperationCtx) executeTemplate(ctx context.Context, nodeName string,
 		// Inject the retryAttempt number
 		localParams[common.LocalVarRetries] = strconv.Itoa(retryNum)
 
+		// Inject lastRetry variables
+		// the first node will not have "lastRetry" variables so they must have default values
+		// for the expression to resolve
+		lastRetryExitCode, lastRetryDuration := "0", "0"
+		var lastRetryStatus, lastRetryMessage string
+		if lastChildNode != nil {
+			if lastChildNode.Outputs != nil && lastChildNode.Outputs.ExitCode != nil {
+				lastRetryExitCode = *lastChildNode.Outputs.ExitCode
+			}
+			lastRetryStatus = string(lastChildNode.Phase)
+			lastRetryDuration = fmt.Sprint(lastChildNode.GetDuration().Seconds())
+			lastRetryMessage = lastChildNode.Message
+		}
+		localParams[common.LocalVarRetriesLastExitCode] = lastRetryExitCode
+		localParams[common.LocalVarRetriesLastDuration] = lastRetryDuration
+		localParams[common.LocalVarRetriesLastStatus] = lastRetryStatus
+		localParams[common.LocalVarRetriesLastMessage] = lastRetryMessage
 		processedTmpl, err = common.SubstituteParams(processedTmpl, woc.globalParams, localParams)
 		if errorsutil.IsTransientErr(err) {
 			return node, err

workflow/controller/operator_test.go

@@ -1508,6 +1508,59 @@ func TestRetriesVariableWithGlobalVariableInPodSpecPatch(t *testing.T) {
 	assert.ElementsMatch(t, actual, expected)
 }
 
+var lastretryVariableInPodSpecPatchTemplate = `
+apiVersion: argoproj.io/v1alpha1
+kind: Workflow
+metadata:
+  name: whalesay
+spec:
+  entrypoint: whalesay
+  templates:
+  - name: whalesay
+    retryStrategy:
+      limit: 10
+    podSpecPatch: |
+      containers:
+        - name: main
+          resources:
+            limits:
+              memory: "{{= (sprig.int(lastRetry.exitCode)==1 ? (sprig.int(retries)+1) : 1)* 100}}Mi"
+    container:
+      image: python:alpine3.6
+      command: ["python", -c]
+      args: ["import sys; sys.exit(1)"]
+`
+
+func TestLastRetryVariableInPodSpecPatch(t *testing.T) {
+	wf := wfv1.MustUnmarshalWorkflow(lastretryVariableInPodSpecPatchTemplate)
+	cancel, controller := newController(wf)
+	defer cancel()
+	ctx := context.Background()
+	iterations := 5
+	var woc *wfOperationCtx
+	for i := 1; i <= iterations; i++ {
+		woc = newWorkflowOperationCtx(wf, controller)
+		if i != 1 {
+			makePodsPhase(ctx, woc, apiv1.PodFailed, withExitCode(1))
+		}
+		woc.operate(ctx)
+		wf = woc.wf
+	}
+
+	pods, err := listPods(woc)
+	require.NoError(t, err)
+	assert.Len(t, pods.Items, iterations)
+	expected := []string{}
+	actual := []string{}
+	for i := 0; i < iterations; i++ {
+		actual = append(actual, pods.Items[i].Spec.Containers[1].Resources.Limits.Memory().String())
+		expected = append(expected, fmt.Sprintf("%dMi", (i+1)*100))
+	}
+	// expecting memory limit to increase after each retry: "100Mi", "200Mi", "300Mi", "400Mi", "500Mi"
+	// ordering not preserved
+	assert.ElementsMatch(t, actual, expected)
+}
+
 var stepsRetriesVariableTemplate = `
 apiVersion: argoproj.io/v1alpha1
 kind: Workflow