From 34a41e86a88006818aba4477a9fb43a69338c3cc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ren=C3=A1n=20Del=20Valle?= <commit@ridv.xyz>
Date: Thu, 29 Apr 2021 20:08:19 -0700
Subject: [PATCH] POC on how verification function would work. Update function
 should have something similar.

---
 realis.go | 18 +++++++++++++++++-
 retry.go  |  7 +++++--
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/realis.go b/realis.go
index ce9b3fa..07385d5 100644
--- a/realis.go
+++ b/realis.go
@@ -672,7 +672,23 @@ func (r *realisClient) CreateJob(auroraJob Job) (*aurora.Response, error) {
 		func() (*aurora.Response, error) {
 			return r.client.CreateJob(context.TODO(), auroraJob.JobConfig())
 		},
-		nil,
+		func() (*aurora.Response, bool) {
+			getTaskResp, err := r.client.GetTasksStatus(
+				context.TODO(),
+				&aurora.TaskQuery{JobKeys: []*aurora.JobKey{auroraJob.JobKey()}},
+			)
+
+			if err != nil {
+				return nil, false
+			}
+
+			tasks := response.ScheduleStatusResult(getTaskResp).GetTasks()
+			if len(tasks) != int(auroraJob.GetInstanceCount()) {
+				return nil, false
+			}
+
+			return nil, true
+		},
 	)
 
 	if retryErr != nil {
diff --git a/retry.go b/retry.go
index f0257ba..4a2ffbc 100644
--- a/retry.go
+++ b/retry.go
@@ -197,9 +197,12 @@ func (r *realisClient) thriftCallWithRetries(
 						// Allow caller to provide a function which checks if the original call was successful before
 						// it timed out.
 						if verifyOnTimeout != nil {
-							resp, ok := verifyOnTimeout()
+							verifyResp, ok := verifyOnTimeout()
 							if ok {
-								return resp, nil
+								r.logger.debugPrint("verified that the call went through successfully")
+								// Response might be different than the original.
+								return verifyResp, nil
+
 							}
 						}
 					}