fix(retry): prioritize cancellation in Next() method and update test expectations

Author: sanki92

pkg/util/retry/retry.go

@@ -172,6 +172,15 @@ func (r *Retry) Next() bool {
 		return false
 	}
 
+	// Check for cancellation first to prioritize over timer.
+	select {
+	case <-r.opts.Closer:
+		return false
+	case <-r.ctx.Done():
+		return false
+	default:
+	}
+
 	backoff, actualWait, shouldAttempt := r.calcDurationScopedBackoff()
 
 	if !shouldAttempt && r.opts.PreemptivelyCancel {
@@ -192,15 +201,6 @@ func (r *Retry) Next() bool {
 		r.backingOffHook(actualWait)
 	}
 
-	// Check for cancellation first to prioritize over timer.
-	select {
-	case <-r.opts.Closer:
-		return false
-	case <-r.ctx.Done():
-		return false
-	default:
-	}
-
 	select {
 	case <-r.opts.Closer:
 		return false

pkg/util/retry/retry_test.go

@@ -469,7 +469,7 @@ func TestRetryWithMaxDuration(t *testing.T) {
 			preRetryFunc: func() {
 				cancelCtxFunc()
 			},
-			expectedTimeSpent: time.Millisecond,
+			expectedTimeSpent: 0,
 			expectedErr:       true,
 		},
 		{
@@ -486,7 +486,7 @@ func TestRetryWithMaxDuration(t *testing.T) {
 			preRetryFunc: func() {
 				close(closeCh)
 			},
-			expectedTimeSpent: time.Millisecond,
+			expectedTimeSpent: 0,
 			expectedErr:       true,
 		},
 	}