From e2e2487f40db63f46e11a677df4ae3b629af1180 Mon Sep 17 00:00:00 2001
From: wei liu <wei.liu@zilliz.com>
Date: Thu, 2 Nov 2023 09:02:15 +0800
Subject: [PATCH] avoid sleep to ctx done during retry (#27992) (#28092)

Signed-off-by: Wei Liu <wei.liu@zilliz.com>
---
 pkg/util/retry/retry.go      | 6 ++++++
 pkg/util/retry/retry_test.go | 6 ++++++
 2 files changed, 12 insertions(+)

diff --git a/pkg/util/retry/retry.go b/pkg/util/retry/retry.go
index 3f48a18d7e..cb1a8a2092 100644
--- a/pkg/util/retry/retry.go
+++ b/pkg/util/retry/retry.go
@@ -53,6 +53,12 @@ func Do(ctx context.Context, fn func() error, opts ...Option) error {
 				return el
 			}
 
+			deadline, ok := ctx.Deadline()
+			if ok && time.Until(deadline) < c.sleep {
+				// to avoid sleep until ctx done
+				return el
+			}
+
 			select {
 			case <-time.After(c.sleep):
 			case <-ctx.Done():
diff --git a/pkg/util/retry/retry_test.go b/pkg/util/retry/retry_test.go
index d21522482e..afc0861838 100644
--- a/pkg/util/retry/retry_test.go
+++ b/pkg/util/retry/retry_test.go
@@ -63,6 +63,12 @@ func TestMaxSleepTime(t *testing.T) {
 	err := Do(ctx, testFn, Attempts(3), MaxSleepTime(200*time.Millisecond))
 	assert.Error(t, err)
 	t.Log(err)
+
+	ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
+	defer cancel()
+	err = Do(ctx, testFn, Attempts(10), MaxSleepTime(200*time.Millisecond))
+	assert.Error(t, err)
+	assert.Nil(t, ctx.Err())
 }
 
 func TestSleep(t *testing.T) {