feat(ehancement): make fair context deadline for all hedged requests

Author: jeevatkm

hedging.go

@@ -166,7 +166,30 @@ func (ht *Hedging) RoundTrip(req *http.Request) (*http.Response, error) {
 	}
 
 	ctx := req.Context()
-	hedgeCtx, cancel := context.WithCancel(ctx)
+	deadline, hasDeadline := ctx.Deadline()
+
+	// Derive hedgeCtx from the original request context to respect cancellations
+	var (
+		hedgeCtx context.Context
+		cancel   context.CancelFunc
+	)
+	if hasDeadline {
+		// Use original deadline for the race (first to complete wins)
+		remaining := time.Until(deadline)
+		if remaining > 0 {
+			hedgeCtx, cancel = context.WithTimeout(ctx, remaining)
+		} else {
+			// Deadline already expired, use context with cancel
+			hedgeCtx, cancel = context.WithCancel(ctx)
+		}
+	} else {
+		// No deadline in original context, create cancellable context from it
+		hedgeCtx, cancel = context.WithCancel(ctx)
+	}
+
+	// defer cancel() ensures cleanup on all paths (timeout, cancellation, or normal return)
+	// cancel() may also be called inside once.Do() when a request wins, but calling it
+	// multiple times is safe and ensures the context is canceled as soon as any goroutine completes
 	defer cancel()
 
 	type result struct {
@@ -177,21 +200,24 @@ func (ht *Hedging) RoundTrip(req *http.Request) (*http.Response, error) {
 	resultCh := make(chan result, ht.MaxRequest())
 	var once sync.Once
 
-	for i := 0; i < ht.MaxRequest(); i++ {
+	ht.lock.RLock()
+	maxReq := ht.maxRequest
+	delay := ht.delay
+	rateDelay := ht.rateDelay
+	ht.lock.RUnlock()
+
+	for i := 0; i < maxReq; i++ {
 		if i > 0 {
-			if ht.Delay() > 0 {
+			if delay > 0 {
 				select {
-				case <-time.After(ht.Delay()):
+				case <-time.After(delay):
 				case <-hedgeCtx.Done():
 					break
 				}
 			}
 
 			// Rate limiting: add delay between requests based on maxPerSecond
 			// to prevent overwhelming the server.
-			ht.lock.RLock()
-			rateDelay := ht.rateDelay
-			ht.lock.RUnlock()
 			if rateDelay > 0 {
 				select {
 				case <-time.After(rateDelay):
@@ -202,13 +228,16 @@ func (ht *Hedging) RoundTrip(req *http.Request) (*http.Response, error) {
 		}
 
 		go func() {
-			hedgedReq := req.Clone(ctx)
+			hedgedReq := req.Clone(hedgeCtx)
 			resp, err := ht.transport.RoundTrip(hedgedReq)
 
 			won := false
 			once.Do(func() {
 				won = true
 				resultCh <- result{resp: resp, err: err}
+
+				// Cancel inside once.Do() to stop other goroutines immediately when a request wins
+				// defer cancel() ensures cleanup even if no request completes successfully
 				cancel()
 			})
 

hedging_test.go

@@ -704,3 +704,25 @@ func TestHedgingNoDoubleWrap(t *testing.T) {
 	// Verify the configuration is the new one
 	assertEqual(t, hedging2.Delay(), 100*time.Millisecond, "Expected 100ms delay")
 }
+
+func TestHedgingRoundTripDeadlineExpired(t *testing.T) {
+	var attemptCount int32
+	ts := createHedgingTestServer(t, &attemptCount, 0, 0)
+	defer ts.Close()
+
+	h := NewHedging().
+		SetDelay(10 * time.Millisecond).
+		SetMaxRequest(3).
+		SetMaxRequestPerSecond(0)
+
+	c := dcnl().SetHedging(h)
+
+	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(-1*time.Millisecond))
+	defer cancel()
+
+	_, err := c.R().SetContext(ctx).Get(ts.URL + "/")
+	assertErrorIs(t, context.DeadlineExceeded, err, "Expected context deadline expired error")
+
+	time.Sleep(50 * time.Millisecond)
+	assertEqual(t, int32(0), atomic.LoadInt32(&attemptCount))
+}