From 2e0c8d2189f229952d53f2e61ace731be93c9f72 Mon Sep 17 00:00:00 2001
From: rohitagg2020 <55523204+rohitagg2020@users.noreply.github.com>
Date: Fri, 27 Sep 2024 10:10:50 +0530
Subject: [PATCH] Providing startTime and currentTime to facilitate timeout in
 case of ytt wait rules (#1023)

* Providing startTime and currentTime to facilitate timeout in case of ytt wait rules

Signed-off-by: Rohit Aggarwal <rohit.aggarwal@broadcom.com>

* Update test/e2e/custom_wait_rules_test.go

Co-authored-by: Praveen Rewar <8457124+praveenrewar@users.noreply.github.com>

---------

Signed-off-by: Rohit Aggarwal <rohit.aggarwal@broadcom.com>
Co-authored-by: Rohit Aggarwal <rohit.aggarwal@broadcom.com>
Co-authored-by: Praveen Rewar <8457124+praveenrewar@users.noreply.github.com>
---
 .../resourcesmisc/custom_waiting_resource.go  |  9 ++-
 .../resourcesmisc/wait_rule_contract_v1.go    |  3 +
 test/e2e/custom_wait_rules_test.go            | 62 +++++++++++++++++++
 3 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/pkg/kapp/resourcesmisc/custom_waiting_resource.go b/pkg/kapp/resourcesmisc/custom_waiting_resource.go
index 110ea5689..283a994eb 100644
--- a/pkg/kapp/resourcesmisc/custom_waiting_resource.go
+++ b/pkg/kapp/resourcesmisc/custom_waiting_resource.go
@@ -65,10 +65,17 @@ func (s CustomWaitingResource) IsDoneApplying() DoneApplyState {
 	}
 
 	if s.waitRule.Ytt != nil {
+		startTime, found := timeoutMap.Load(s.resource.Description())
+		if !found {
+			startTime = time.Now().Unix()
+			timeoutMap.Store(s.resource.Description(), startTime)
+		}
 		configObj, err := WaitRuleContractV1{
 			ResourceMatcher: ctlres.AnyMatcher{
 				Matchers: ctlconf.ResourceMatchers(s.waitRule.ResourceMatchers).AsResourceMatchers()},
-			Starlark: s.waitRule.Ytt.FuncContractV1.Resource,
+			Starlark:    s.waitRule.Ytt.FuncContractV1.Resource,
+			CurrentTime: time.Now().Unix(),
+			StartTime:   startTime.(int64),
 		}.Apply(s.resource)
 		if err != nil {
 			return DoneApplyState{Done: true, Successful: false, Message: fmt.Sprintf(
diff --git a/pkg/kapp/resourcesmisc/wait_rule_contract_v1.go b/pkg/kapp/resourcesmisc/wait_rule_contract_v1.go
index e1344d9ac..077247eeb 100644
--- a/pkg/kapp/resourcesmisc/wait_rule_contract_v1.go
+++ b/pkg/kapp/resourcesmisc/wait_rule_contract_v1.go
@@ -16,6 +16,8 @@ import (
 type WaitRuleContractV1 struct {
 	ResourceMatcher ctlres.ResourceMatcher
 	Starlark        string
+	CurrentTime     int64
+	StartTime       int64
 }
 
 type waitRuleContractV1Result struct {
@@ -48,6 +50,7 @@ func (t WaitRuleContractV1) evalYtt(res ctlres.Resource) (*WaitRuleContractV1Res
 		}
 		return yaml.Marshal(res.DeepCopyRaw())
 	}
+	opts.DataValuesFlags.KVsFromStrings = []string{fmt.Sprintf("startTime=%d", t.StartTime), fmt.Sprintf("currentTime=%d", t.CurrentTime)}
 
 	filesToProcess := []*files.File{
 		files.MustNewFileFromSource(files.NewBytesSource("resource.star", []byte(t.Starlark))),
diff --git a/test/e2e/custom_wait_rules_test.go b/test/e2e/custom_wait_rules_test.go
index f969916ba..668e168cb 100644
--- a/test/e2e/custom_wait_rules_test.go
+++ b/test/e2e/custom_wait_rules_test.go
@@ -260,3 +260,65 @@ Succeeded`))
 		require.Equal(t, expectedOutput, out)
 	})
 }
+
+func TestYTTWaitRuleWithTimeout(t *testing.T) {
+	env := BuildEnv(t)
+	logger := Logger{}
+	kapp := Kapp{t, env.Namespace, env.KappBinaryPath, logger}
+	name := "test-custom-wait-rule-timeout-contract-v1"
+
+	yaml := `
+apiVersion: v1
+kind: Pod
+metadata:
+  name: nginx
+spec:
+  containers:
+  - name: nginx
+    image: %s
+    ports:
+    - containerPort: 80
+---
+apiVersion: kapp.k14s.io/v1alpha1
+kind: Config
+waitRules:
+- ytt:
+    funcContractV1:
+      resource.star: |
+        def is_done(resource):
+          for condition in resource.status.conditions:
+            if condition.type == "ContainersReady" and condition.status == "False":
+              expiryTime = getExpiryTime(resource.startTime)
+              if int(resource.currentTime) - expiryTime > 0:
+                  return {"done": True, "successful": False, "message": "Continuously failed for 50s with " + condition.message}
+              else:
+                  return {"done": False, "successful": False, "message": condition.message}
+              end
+            elif condition.type == "Ready" and condition.status == "True":
+              return {"done": True, "successful": True, "message": condition.message}
+            end
+          end
+          return {"done": False, "successful": False, "message": "Not in Failed or Running state"}
+        end
+    
+        def getExpiryTime(startTime):
+          return int(startTime)+50 
+        end
+  resourceMatchers:
+  - apiVersionKindMatcher: {apiVersion: v1, kind: Pod}
+`
+
+	cleanUp := func() {
+		kapp.Run([]string{"delete", "-a", name})
+	}
+	cleanUp()
+	defer cleanUp()
+
+	logger.Section("Deploy timeout after staying in a condition for certain time with ytt wait rules", func() {
+		_, err := kapp.RunWithOpts([]string{"deploy", "-f", "-", "-a", name, "--json"},
+			RunOpts{IntoNs: true, AllowError: true, StdinReader: strings.NewReader(fmt.Sprintf(yaml, "nginx:200"))})
+
+		require.Error(t, err)
+		require.Contains(t, err.Error(), "Continuously failed for 50s with")
+	})
+}