+135
pkg/robusthttp/client.go
+135
pkg/robusthttp/client.go
···
1
+
package robusthttp
2
+
3
+
import (
4
+
"context"
5
+
"log/slog"
6
+
"net/http"
7
+
"time"
8
+
9
+
"github.com/hashicorp/go-cleanhttp"
10
+
"github.com/hashicorp/go-retryablehttp"
11
+
"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
12
+
)
13
+
14
+
type LeveledSlog struct {
15
+
inner *slog.Logger
16
+
}
17
+
18
+
// re-writes HTTP client ERROR to WARN level (because of retries)
19
+
func (l LeveledSlog) Error(msg string, keysAndValues ...any) {
20
+
l.inner.Warn(msg, keysAndValues...)
21
+
}
22
+
23
+
func (l LeveledSlog) Warn(msg string, keysAndValues ...any) {
24
+
l.inner.Warn(msg, keysAndValues...)
25
+
}
26
+
27
+
func (l LeveledSlog) Info(msg string, keysAndValues ...any) {
28
+
l.inner.Info(msg, keysAndValues...)
29
+
}
30
+
31
+
func (l LeveledSlog) Debug(msg string, keysAndValues ...any) {
32
+
l.inner.Debug(msg, keysAndValues...)
33
+
}
34
+
35
+
type Option func(*retryablehttp.Client)
36
+
37
+
// WithMaxRetries sets the maximum number of retries for the HTTP client.
38
+
func WithMaxRetries(maxRetries int) Option {
39
+
return func(client *retryablehttp.Client) {
40
+
client.RetryMax = maxRetries
41
+
}
42
+
}
43
+
44
+
// WithRetryWaitMin sets the minimum wait time between retries.
45
+
func WithRetryWaitMin(waitMin time.Duration) Option {
46
+
return func(client *retryablehttp.Client) {
47
+
client.RetryWaitMin = waitMin
48
+
}
49
+
}
50
+
51
+
// WithRetryWaitMax sets the maximum wait time between retries.
52
+
func WithRetryWaitMax(waitMax time.Duration) Option {
53
+
return func(client *retryablehttp.Client) {
54
+
client.RetryWaitMax = waitMax
55
+
}
56
+
}
57
+
58
+
// WithLogger sets a custom logger for the HTTP client.
59
+
func WithLogger(logger *slog.Logger) Option {
60
+
return func(client *retryablehttp.Client) {
61
+
client.Logger = retryablehttp.LeveledLogger(LeveledSlog{inner: logger})
62
+
}
63
+
}
64
+
65
+
// WithTransport sets a custom transport for the HTTP client.
66
+
func WithTransport(transport http.RoundTripper) Option {
67
+
return func(client *retryablehttp.Client) {
68
+
client.HTTPClient.Transport = transport
69
+
}
70
+
}
71
+
72
+
// WithRetryPolicy sets a custom retry policy for the HTTP client.
73
+
func WithRetryPolicy(policy retryablehttp.CheckRetry) Option {
74
+
return func(client *retryablehttp.Client) {
75
+
client.CheckRetry = policy
76
+
}
77
+
}
78
+
79
+
// Generates an HTTP client with decent general-purpose defaults around
80
+
// timeouts and retries. The returned client has the stdlib http.Client
81
+
// interface, but has Hashicorp retryablehttp logic internally.
82
+
//
83
+
// This client will retry on connection errors, 5xx status (except 501).
84
+
// It will log intermediate failures with WARN level. This does not start from
85
+
// http.DefaultClient.
86
+
//
87
+
// This should be usable for XRPC clients, and other general inter-service
88
+
// client needs. CLI tools might want shorter timeouts and fewer retries by
89
+
// default.
90
+
func NewClient(options ...Option) *http.Client {
91
+
logger := LeveledSlog{inner: slog.Default().With("subsystem", "RobustHTTPClient")}
92
+
retryClient := retryablehttp.NewClient()
93
+
retryClient.HTTPClient.Transport = otelhttp.NewTransport(cleanhttp.DefaultPooledTransport())
94
+
retryClient.RetryMax = 3
95
+
retryClient.RetryWaitMin = 1 * time.Second
96
+
retryClient.RetryWaitMax = 10 * time.Second
97
+
retryClient.Logger = retryablehttp.LeveledLogger(logger)
98
+
retryClient.CheckRetry = DefaultRetryPolicy
99
+
100
+
for _, option := range options {
101
+
option(retryClient)
102
+
}
103
+
104
+
client := retryClient.StandardClient()
105
+
client.Timeout = 30 * time.Second
106
+
return client
107
+
}
108
+
109
+
// For use in local integration tests. Short timeouts, no retries, etc
110
+
func TestingHTTPClient() *http.Client {
111
+
112
+
client := http.DefaultClient
113
+
client.Timeout = 1 * time.Second
114
+
return client
115
+
}
116
+
117
+
// DefaultRetryPolicy is a custom wrapper around retryablehttp.DefaultRetryPolicy.
118
+
// It treats `429 Too Many Requests` as non-retryable, so the application can decide
119
+
// how to deal with rate-limiting.
120
+
func DefaultRetryPolicy(ctx context.Context, resp *http.Response, err error) (bool, error) {
121
+
if err == nil && resp.StatusCode == http.StatusTooManyRequests {
122
+
return false, nil
123
+
}
124
+
return retryablehttp.DefaultRetryPolicy(ctx, resp, err)
125
+
}
126
+
127
+
func NoInternalServerErrorPolicy(ctx context.Context, resp *http.Response, err error) (bool, error) {
128
+
if err == nil && resp.StatusCode == http.StatusTooManyRequests {
129
+
return false, nil
130
+
}
131
+
if err == nil && resp.StatusCode == http.StatusInternalServerError {
132
+
return false, nil
133
+
}
134
+
return retryablehttp.DefaultRetryPolicy(ctx, resp, err)
135
+
}