diff --git a/common_client.go b/common_client.go index d221a25..3d13a0b 100644 --- a/common_client.go +++ b/common_client.go @@ -75,12 +75,20 @@ func sendRequest(c *http.Client, req *http.Request) (*http.Response, error) { } type httpClientOption struct { - logger *slog.Logger - retryMax int - retryWaitMax time.Duration + logger *slog.Logger + + // Options for built-in retryable HTTP client. + // Ignored if a custom retryable HTTP client is provided via WithRetryableHTTPClint. + retryMax int + retryWaitMax time.Duration + + // fields added to the transport if specified rootCAs *x509.CertPool tlsInsecureSkipVerify bool proxyFunc ProxyFunc + timeout time.Duration + + retryableHTTPClient *retryablehttp.Client } func (o *httpClientOption) defaults() { @@ -93,14 +101,28 @@ func (o *httpClientOption) defaults() { if o.retryWaitMax == 0 { o.retryWaitMax = 30 * time.Second } + if o.timeout == 0 { + o.timeout = 5 * time.Minute + } } func (o *httpClientOption) newRetryableHTTPClient() (*retryablehttp.Client, error) { - retryClient := retryablehttp.NewClient() - retryClient.Logger = o.logger - retryClient.RetryMax = o.retryMax - retryClient.RetryWaitMax = o.retryWaitMax - retryClient.HTTPClient.Timeout = 5 * time.Minute // timeout must be > 1m to accomodate long polling + var retryClient *retryablehttp.Client + if o.retryableHTTPClient != nil { + retryClient = o.retryableHTTPClient + } else { + retryClient = retryablehttp.NewClient() + retryClient.RetryMax = o.retryMax + retryClient.RetryWaitMax = o.retryWaitMax + } + + if retryClient.HTTPClient.Timeout == 0 { + retryClient.HTTPClient.Timeout = o.timeout + } + + if retryClient.Logger == nil { + retryClient.Logger = o.logger + } transport, ok := retryClient.HTTPClient.Transport.(*http.Transport) if !ok { @@ -120,7 +142,9 @@ func (o *httpClientOption) newRetryableHTTPClient() (*retryablehttp.Client, erro transport.TLSClientConfig.InsecureSkipVerify = true } - transport.Proxy = o.proxyFunc + if o.proxyFunc != nil { + transport.Proxy = o.proxyFunc + } retryClient.HTTPClient.Transport = transport @@ -145,6 +169,14 @@ func (c *commonClient) setUserAgent() { // HTTPOption defines a functional option for configuring the Client. type HTTPOption func(*httpClientOption) +// WithRetryableHTTPClint allows users to provide a custom retryable HTTP client. +// If not set, a default client will be used with the specified retry and timeout settings. +func WithRetryableHTTPClint(client *retryablehttp.Client) HTTPOption { + return func(c *httpClientOption) { + c.retryableHTTPClient = client + } +} + // WithLogger sets a custom logger for the Client. // If nil is passed, a discard logger will be used. func WithLogger(logger *slog.Logger) HTTPOption { @@ -190,3 +222,10 @@ func WithProxy(proxyFunc ProxyFunc) HTTPOption { c.proxyFunc = proxyFunc } } + +// WithTimeout sets a timeout for the Client. +func WithTimeout(duration time.Duration) HTTPOption { + return func(c *httpClientOption) { + c.timeout = duration + } +} diff --git a/common_client_test.go b/common_client_test.go index d06e1de..fb16edf 100644 --- a/common_client_test.go +++ b/common_client_test.go @@ -7,8 +7,10 @@ import ( "net/http/httptest" "net/url" "testing" + "time" "github.com/actions/scaleset/internal/testserver" + "github.com/hashicorp/go-retryablehttp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/net/http/httpproxy" @@ -153,3 +155,88 @@ func TestUserAgent(t *testing.T) { assert.Equal(t, want, got) } + +// TestWithRetryableHTTPClient verifies that a custom retryable HTTP client +// provided via WithRetryableHTTPClient is actually used instead of the built-in one +func TestWithRetryableHTTPClient(t *testing.T) { + t.Run("uses custom retryable client instead of built-in", func(t *testing.T) { + attemptCount := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + attemptCount++ + if attemptCount == 1 { + w.WriteHeader(http.StatusServiceUnavailable) + return + } + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"result": "success"}`)) + })) + defer server.Close() + + // Create a custom retryable HTTP client with specific retry configuration + customRetryClient := retryablehttp.NewClient() + customRetryClient.RetryMax = 3 + customRetryClient.RetryWaitMax = 10 * time.Millisecond + + // Create options with the custom retryable client + opts := defaultHTTPClientOption() + WithRetryableHTTPClint(customRetryClient)(&opts) + + // Verify that the custom client is set in options + assert.NotNil(t, opts.retryableHTTPClient) + assert.Equal(t, customRetryClient, opts.retryableHTTPClient) + + // Create the common client with custom retryable client + client := newCommonClient(testSystemInfo, opts) + + // Make a request that will trigger a retry + req, err := http.NewRequest("GET", server.URL, nil) + require.NoError(t, err) + + resp, err := client.do(req) + require.NoError(t, err) + + // Should succeed after retry + assert.Equal(t, http.StatusOK, resp.StatusCode) + assert.Equal(t, 2, attemptCount) + + // Verify that the client used is the custom one by checking newRetryableHTTPClient + retrievedRetryClient, err := client.newRetryableHTTPClient() + require.NoError(t, err) + assert.Equal(t, customRetryClient, retrievedRetryClient, "should return the custom retryable client") + }) + + t.Run("respects custom client's retry configuration over built-in defaults", func(t *testing.T) { + attemptCount := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + attemptCount++ + w.WriteHeader(http.StatusServiceUnavailable) + })) + defer server.Close() + + // Create custom client with limited retries + customRetryClient := retryablehttp.NewClient() + customRetryClient.RetryMax = 1 // Only 1 retry (2 total attempts) + customRetryClient.RetryWaitMax = 5 * time.Millisecond + + opts := defaultHTTPClientOption() + WithRetryableHTTPClint(customRetryClient)(&opts) + + client := newCommonClient(testSystemInfo, opts) + + req, err := http.NewRequest("GET", server.URL, nil) + require.NoError(t, err) + + resp, err := client.do(req) + // When all retries are exhausted with a retryable error, the client gives up + // and an error is returned + if err != nil { + // Expected: request failed after exhausting retries + assert.Contains(t, err.Error(), "giving up after 2 attempt(s)") + } else { + // Or the final response is returned + assert.Equal(t, http.StatusServiceUnavailable, resp.StatusCode) + } + // Should have tried 1 initial + 1 retry = 2 times total + assert.Equal(t, 2, attemptCount) + }) +}