From 7629583b8b9cbd84871be2be0ebad08cb3035bb7 Mon Sep 17 00:00:00 2001
From: Daniil Fedotov <daniil.fedotov@kasten.io>
Date: Tue, 19 Dec 2023 14:04:26 -0500
Subject: [PATCH] Limit minimal time poll.Wait* functions can wait between
 attempts (#2543)

Currently calculation of sleep time can make it negative, which
makes `t.C` channel return right away and does not give ctx time
to finish on deadline.
This means that even if context is expired, retries still happen until
there are no more retries.

The fix is to cap minimal sleep time at 5ms to make sure that context
can perform expiration.

NOTE: we could make the minimal sleep smaller, but it's wouldn't make sense
getting less than OS scheduling frequency.

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
---
 pkg/poll/poll.go      | 12 ++++--------
 pkg/poll/poll_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/pkg/poll/poll.go b/pkg/poll/poll.go
index 398513b9ee..a807e421cd 100644
--- a/pkg/poll/poll.go
+++ b/pkg/poll/poll.go
@@ -91,7 +91,10 @@ func WaitWithBackoffWithRetries(ctx context.Context, b backoff.Backoff, numRetri
 		sleep := b.Duration()
 		if deadline, ok := ctx.Deadline(); ok {
 			ctxSleep := time.Until(deadline)
-			sleep = minDuration(sleep, ctxSleep)
+			// We want to wait for smaller of backoff sleep and context sleep
+			// but it has to be > 0 to give ctx.Done() a chance
+			// to return below
+			sleep = max(min(sleep, ctxSleep), 5*time.Millisecond)
 		}
 		t.Reset(sleep)
 		select {
@@ -101,10 +104,3 @@ func WaitWithBackoffWithRetries(ctx context.Context, b backoff.Backoff, numRetri
 		}
 	}
 }
-
-func minDuration(a, b time.Duration) time.Duration {
-	if a < b {
-		return a
-	}
-	return b
-}
diff --git a/pkg/poll/poll_test.go b/pkg/poll/poll_test.go
index c7de205675..4a4983e570 100644
--- a/pkg/poll/poll_test.go
+++ b/pkg/poll/poll_test.go
@@ -16,7 +16,9 @@ package poll
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"runtime"
 	"testing"
 	"time"
 
@@ -130,6 +132,30 @@ func (s *PollSuite) TestWaitWithBackoffCancellation(c *C) {
 	c.Check(err, NotNil)
 }
 
+func (s *PollSuite) TestWaitWithRetriesTimeout(c *C) {
+	// There's a better chance of catching a race condition
+	// if there is only one thread
+	maxprocs := runtime.GOMAXPROCS(1)
+	defer runtime.GOMAXPROCS(maxprocs)
+
+	f := func(context.Context) (bool, error) {
+		return false, errors.New("retryable")
+	}
+	errf := func(err error) bool {
+		return err.Error() == "retryable"
+	}
+	ctx := context.Background()
+	var cancel context.CancelFunc
+	ctx, cancel = context.WithTimeout(ctx, time.Millisecond)
+	defer cancel()
+
+	backoff := backoff.Backoff{}
+	backoff.Min = 2 * time.Millisecond
+	err := WaitWithBackoffWithRetries(ctx, backoff, 10, errf, f)
+	c.Check(err, NotNil)
+	c.Assert(err.Error(), Matches, ".*context deadline exceeded*")
+}
+
 func (s *PollSuite) TestWaitWithBackoffBackoff(c *C) {
 	const numIterations = 10
 	i := 0