From 6b9e1ed9612ae7b46e43294784ff9aaffdca9be6 Mon Sep 17 00:00:00 2001
From: Anas Khan <83116240+anxkhn@users.noreply.github.com>
Date: Mon, 29 Jun 2026 01:07:10 +0530
Subject: [PATCH] fix(time): prevent overflow in ExponentialDelayWithJitter
 backoff

ExponentialDelayWithJitter computed `baseDelay * (1 << attempt)` before
applying `min(delay, maxDelay)`, so a large attempt overflowed the int64
time.Duration ahead of the cap. The wrapped negative or zero value passed
the cap unchanged, making `delay > 0` false and skipping the sleep.

In the scheduler this `attempt` is a host's ConcurrentRegisterCount, not a
retry counter. With the production constants (baseDelay 30ms, maxDelay 1s)
the backoff breaks once a host has >= 39 concurrent registrations, which is
exactly the thundering herd the delay exists to dampen.

Compute the shift only when it cannot overflow time.Duration (using the
leading-zero count of baseDelay); otherwise the result would exceed maxDelay
anyway, so clamp to maxDelay directly. Also guard a non-positive baseDelay,
matching the existing guard in RandomDelayWithJitter. Behavior is unchanged
for the existing small-attempt cases. Add regression cases covering the
overflow regime.

Signed-off-by: Anas Khan <83116240+anxkhn@users.noreply.github.com>
---
 pkg/time/delay.go      | 14 ++++++++++++--
 pkg/time/delay_test.go | 24 ++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/pkg/time/delay.go b/pkg/time/delay.go
index 0c13eb3b602..956e5750ee9 100644
--- a/pkg/time/delay.go
+++ b/pkg/time/delay.go
@@ -18,6 +18,7 @@ package time
 
 import (
 	"context"
+	"math/bits"
 	"math/rand"
 	"time"
 )
@@ -39,8 +40,17 @@ func LinearDelay(ctx context.Context, attempt uint, increment, maxDelay time.Dur
 // ExponentialDelayWithJitter is an exponential backoff strategy with jitter for retries. It calculates delay based on the attempt number,
 // adds jitter, and sleeps for that duration, capped at maxDelay.
 func ExponentialDelayWithJitter(ctx context.Context, attempt uint, baseDelay, maxDelay time.Duration) error {
-	delay := baseDelay * time.Duration(1<<attempt)
-	delay = min(delay, maxDelay)
+	// Compute baseDelay * 2^attempt, capped at maxDelay. The shift is only applied when it
+	// cannot overflow time.Duration; otherwise the result would exceed maxDelay anyway, so
+	// the cap is used directly. Without this guard a large attempt overflows int64 before the
+	// cap is applied, yielding a negative or zero delay that silently skips the backoff.
+	var delay time.Duration
+	if baseDelay > 0 {
+		delay = maxDelay
+		if attempt < uint(bits.LeadingZeros64(uint64(baseDelay))) {
+			delay = min(baseDelay<<attempt, maxDelay)
+		}
+	}
 
 	if delay > 0 {
 		jitter := time.Duration(rand.Int63n(int64(delay)))
diff --git a/pkg/time/delay_test.go b/pkg/time/delay_test.go
index 288ce105859..19f7c1e15df 100644
--- a/pkg/time/delay_test.go
+++ b/pkg/time/delay_test.go
@@ -149,6 +149,30 @@ func TestExponentialDelayWithJitter(t *testing.T) {
 			expectedMin: 650 * time.Millisecond,
 			expectedMax: 3000 * time.Millisecond,
 		},
+		{
+			name:        "overflow attempt stays capped at maxDelay",
+			attempt:     39,
+			baseDelay:   30 * time.Millisecond,
+			maxDelay:    1 * time.Second,
+			expectedMin: 380 * time.Millisecond,
+			expectedMax: 1500 * time.Millisecond,
+		},
+		{
+			name:        "shift overflow attempt stays capped at maxDelay",
+			attempt:     62,
+			baseDelay:   30 * time.Millisecond,
+			maxDelay:    1 * time.Second,
+			expectedMin: 380 * time.Millisecond,
+			expectedMax: 1500 * time.Millisecond,
+		},
+		{
+			name:        "huge attempt stays capped at maxDelay",
+			attempt:     100,
+			baseDelay:   30 * time.Millisecond,
+			maxDelay:    1 * time.Second,
+			expectedMin: 380 * time.Millisecond,
+			expectedMax: 1500 * time.Millisecond,
+		},
 		{
 			name:        "zero baseDelay with jitter",
 			attempt:     5,