fix: interrupt/skip retry on context cancel (#144)

* fix: interrupt retry sleep if context is cancelled

* add unit tests for sleep interrupt

* Early exit for context cancellation/deadline

---------

Co-authored-by: Anmol Chopra <anmol.chopra@gojek.com>

Author: mohitsethia

httpclient/client.go

@@ -147,7 +147,12 @@ func (c *Client) Do(request *http.Request) (*http.Response, error) {
 			_ = response.Body.Close()
 		}
 		if i > 0 {
-			time.Sleep(c.retrier.NextInterval(i - 1)) // sleep after closing the previous response body
+			if err := internal.SleepInterruptible(request.Context(), c.retrier.NextInterval(i-1)); err != nil {
+				multiErr.Push(err.Error())
+				c.reportError(request, err)
+				// no point of retrying after context has been cancelled
+				break
+			}
 
 			request, err = internal.CloneRequest(request, reqGetBody) // Clone the request to reset the body for retry
 			if err != nil {
@@ -162,11 +167,18 @@ func (c *Client) Do(request *http.Request) (*http.Response, error) {
 		if err != nil {
 			multiErr.Push(err.Error())
 			c.reportError(request, err)
+			if internal.IsCtxDone(request.Context()) {
+				break
+			}
 			continue
 		}
 		c.reportRequestEnd(request, response)
 
 		if response.StatusCode >= http.StatusInternalServerError {
+			if internal.IsCtxDone(request.Context()) {
+				break
+			}
+
 			continue
 		}
 

httpclient/client_test.go

@@ -2,10 +2,12 @@ package httpclient
 
 import (
 	"bytes"
+	"context"
 	"io"
 	"net/http"
 	"net/http/httptest"
 	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -518,3 +520,107 @@ func respBody(t *testing.T, response *http.Response) string {
 
 	return string(respBody)
 }
+
+func TestHTTPClientDoContextCancelledDuringRetry(t *testing.T) {
+	noOfRetries := 3
+	backoffInterval := 100 * time.Millisecond
+	maximumJitterInterval := 10 * time.Millisecond
+
+	client := NewClient(
+		WithHTTPTimeout(10*time.Millisecond),
+		WithRetryCount(noOfRetries),
+		WithRetrier(heimdall.NewRetrier(heimdall.NewConstantBackoff(backoffInterval, maximumJitterInterval))),
+	)
+
+	count := atomic.Int32{}
+	count.Store(0)
+	dummyHandler := func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusInternalServerError)
+		count.Add(1)
+	}
+
+	server := httptest.NewServer(http.HandlerFunc(dummyHandler))
+	defer server.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	req, err := http.NewRequest(http.MethodGet, server.URL, nil)
+	require.NoError(t, err)
+	req = req.WithContext(ctx)
+
+	// Cancel the context after a short delay to simulate context cancellation during sleep
+	go func() {
+		time.Sleep(10 * time.Millisecond)
+		cancel()
+	}()
+
+	_, err = client.Do(req)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), context.Canceled.Error())
+	assert.Less(t, count.Load(), int32(noOfRetries+1), "should not have completed all retries due to context cancellation")
+}
+
+func TestHTTPClientDoContextCancelledBeforeRetry(t *testing.T) {
+	client := NewClient(
+		WithHTTPTimeout(10*time.Millisecond),
+		WithRetryCount(3),
+		WithRetrier(heimdall.NewRetrierFunc(func(retry int) time.Duration {
+			assert.Fail(t, "should not have retrier func due to context cancellation")
+			return 0
+		})),
+	)
+	ctx, cancel := context.WithCancel(context.Background())
+
+	count := atomic.Int32{}
+	count.Store(0)
+	dummyHandler := func(w http.ResponseWriter, r *http.Request) {
+		cancel() // Cancel immediately
+		count.Add(1)
+		w.WriteHeader(http.StatusInternalServerError)
+	}
+
+	server := httptest.NewServer(http.HandlerFunc(dummyHandler))
+	defer server.Close()
+
+	req, err := http.NewRequest(http.MethodGet, server.URL, nil)
+	require.NoError(t, err)
+	req = req.WithContext(ctx)
+
+	_, err = client.Do(req)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), context.Canceled.Error())
+	assert.Equal(t, int32(1), count.Load())
+}
+
+func TestHTTPClientDoContextTimeoutDuringRetry(t *testing.T) {
+	noOfRetries := 3
+	backoffInterval := 100 * time.Millisecond
+	maximumJitterInterval := 10 * time.Millisecond
+
+	client := NewClient(
+		WithHTTPTimeout(10*time.Millisecond),
+		WithRetryCount(noOfRetries),
+		WithRetrier(heimdall.NewRetrier(heimdall.NewConstantBackoff(backoffInterval, maximumJitterInterval))),
+	)
+
+	count := atomic.Int32{}
+	count.Store(0)
+	dummyHandler := func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusInternalServerError)
+		count.Add(1)
+	}
+
+	server := httptest.NewServer(http.HandlerFunc(dummyHandler))
+	defer server.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Millisecond)
+	defer cancel()
+
+	req, err := http.NewRequest(http.MethodGet, server.URL, nil)
+	require.NoError(t, err)
+	req = req.WithContext(ctx)
+
+	_, err = client.Do(req)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), context.DeadlineExceeded.Error())
+	assert.Less(t, count.Load(), int32(noOfRetries+1), "should not have completed all retries due to context timeout")
+}

hystrix/hystrix_client.go

@@ -194,7 +194,10 @@ func (hhc *Client) Do(request *http.Request) (*http.Response, error) {
 		}
 
 		if i > 0 {
-			time.Sleep(hhc.retrier.NextInterval(i - 1)) // sleep after closing the previous response body
+			err = internal.SleepInterruptible(request.Context(), hhc.retrier.NextInterval(i-1))
+			if err != nil {
+				return nil, err
+			}
 
 			request, err = internal.CloneRequest(request, reqGetBody) // Clone the request to reset the body for retry
 			if err != nil {
@@ -203,7 +206,7 @@ func (hhc *Client) Do(request *http.Request) (*http.Response, error) {
 		}
 
 		response, err = hhc.hystrixDo(request)
-		if err == nil {
+		if err == nil || internal.IsCtxDone(request.Context()) {
 			break
 		}
 	}

hystrix/hystrix_client_test.go

@@ -606,7 +606,13 @@ func TestHystrixHTTPClientDoContextCancelled(t *testing.T) {
 	const cmdName = "some_command_name_for_ctx_cncl"
 	r := newSimpleMetricRegistry()
 
-	client := NewClient(WithCommandName(cmdName))
+	client := NewClient(
+		WithCommandName(cmdName),
+		WithRetryCount(3),
+		WithRetrier(heimdall.NewRetrierFunc(func(retry int) time.Duration {
+			assert.Fail(t, "should not invoke retrier func due to context cancellation")
+			return 0
+		})))
 
 	req, err := http.NewRequest(http.MethodGet, "http://localhost/test", nil)
 	require.NoError(t, err)
@@ -632,6 +638,44 @@ func TestHystrixHTTPClientDoContextCancelled(t *testing.T) {
 	assert.Zero(t, m.ContextDeadlineExceeded)
 }
 
+func TestHystrixHTTPClientDoContextCancelledDuringRetryBackoff(t *testing.T) {
+	const cmdName = "some_command_name_for_ctx_cncl"
+	ctx, cnclFN := context.WithCancel(context.Background())
+
+	var retrierFuncCount int
+	var retrierFuncStartTime time.Time
+
+	client := NewClient(
+		WithCommandName(cmdName),
+		WithRetryCount(3),
+		WithRetrier(heimdall.NewRetrierFunc(func(retry int) time.Duration {
+			retrierFuncStartTime = time.Now()
+			retrierFuncCount++
+			cnclFN()
+			return time.Hour
+		})))
+
+	reqCount := atomic.Int32{}
+	reqCount.Store(0)
+	server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		reqCount.Add(1)
+		rw.WriteHeader(http.StatusInternalServerError)
+	}))
+	defer server.Close()
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil)
+	require.NoError(t, err)
+
+	response, err := client.Do(req)
+	endTime := time.Now()
+	require.Contains(t, err.Error(), context.Canceled.Error())
+	require.Nil(t, response)
+	assert.Equal(t, int32(1), reqCount.Load())
+	require.Equal(t, 1, retrierFuncCount)
+	assert.Less(t, endTime.Sub(retrierFuncStartTime), time.Millisecond)
+
+	time.Sleep(time.Second)
+}
+
 func TestResponseBodyStreaming(t *testing.T) {
 	const cmdName = "response_body_streamig"
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

internal/ctx.go

@@ -0,0 +1,12 @@
+package internal
+
+import "context"
+
+func IsCtxDone(ctx context.Context) bool {
+	select {
+	case <-ctx.Done():
+		return true
+	default:
+		return false
+	}
+}

internal/ctx_test.go

@@ -0,0 +1,25 @@
+package internal_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/gojek/heimdall/v7/internal"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestIsCtxDone(t *testing.T) {
+	assert.False(t, internal.IsCtxDone(context.Background()))
+
+	ctx, cancel := context.WithCancel(context.Background())
+	assert.False(t, internal.IsCtxDone(ctx))
+	cancel()
+	assert.True(t, internal.IsCtxDone(ctx))
+
+	ctx, cancel = context.WithTimeout(context.Background(), 10*time.Millisecond)
+	defer cancel()
+	assert.False(t, internal.IsCtxDone(ctx))
+	time.Sleep(12 * time.Millisecond)
+	assert.True(t, internal.IsCtxDone(ctx))
+}

internal/sleep.go

@@ -0,0 +1,16 @@
+package internal
+
+import (
+	"context"
+	"time"
+)
+
+// SleepInterruptible sleeps until either the timer triggers or context is cancelled
+func SleepInterruptible(ctx context.Context, d time.Duration) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-time.After(d):
+	}
+	return nil
+}

internal/sleep_test.go

@@ -0,0 +1,32 @@
+package internal_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/gojek/heimdall/v7/internal"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSleepInterruptible_CancelledContext(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	// Cancel the context immediately
+	cancel()
+
+	err := internal.SleepInterruptible(ctx, 10*time.Second) // Long duration to ensure cancellation is what stops it
+	assert.Error(t, err)
+	assert.Equal(t, context.Canceled, err)
+}
+
+func TestSleepInterruptible_CompletesWithoutCancel(t *testing.T) {
+	ctx := context.Background()
+	start := time.Now()
+
+	err := internal.SleepInterruptible(ctx, 50*time.Millisecond) // Short sleep time
+	elapsed := time.Since(start)
+
+	assert.NoError(t, err)
+	assert.True(t, elapsed.Milliseconds() >= int64(50), "Sleep duration should be at least 50ms") // Ensure it slept approximately 50ms
+}