Skip to content

Commit 896f866

Browse files
authored
Adjust http client to support rate limits (#775)
* Adjust http client to support rate limits Changed the way that the CLI implements rate limits and the back off. We've removed the dependency on `roko`, which shouldn't be required with the availability of the `RateLimit-` headers that we can get from client calls. * wrap api limit output in a conditional for debug flag use
1 parent e5401fb commit 896f866

9 files changed

Lines changed: 464 additions & 446 deletions

File tree

cmd/api/api.go

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"context"
66
"encoding/json"
77
"fmt"
8+
"net/http"
89
"os"
910
"strings"
1011
"time"
@@ -114,17 +115,19 @@ func (c *ApiCmd) Run(kongCtx *kong.Context, globals cli.GlobalFlags) error {
114115

115116
fullEndpoint := buildFullEndpoint(c.Endpoint, f.Config.OrganizationSlug(), c.Analytics)
116117

117-
// Create an HTTP client with appropriate configuration
118+
// Create an HTTP client with rate-limit retry via the shared transport.
119+
rl := httpClient.NewRateLimitTransport(nil)
120+
rl.MaxRetryDelay = 60 * time.Second
121+
rl.OnRateLimit = func(attempt int, delay time.Duration) {
122+
if c.Verbose {
123+
fmt.Fprintf(os.Stderr, "WARNING: Rate limit exceeded, retrying in %v @ %q (attempt %d)\n", delay, time.Now().Add(delay).Format(time.RFC3339), attempt)
124+
}
125+
}
126+
118127
client := httpClient.NewClient(
119128
f.Config.APIToken(),
120129
httpClient.WithBaseURL(f.RestAPIClient.BaseURL.String()),
121-
httpClient.WithMaxRetries(3),
122-
httpClient.WithMaxRetryDelay(60*time.Second),
123-
httpClient.WithOnRetry(func(attempt int, delay time.Duration) {
124-
if c.Verbose {
125-
fmt.Fprintf(os.Stderr, "WARNING: Rate limit exceeded, retrying in %v @ %q (attempt %d)\n", delay, time.Now().Add(delay).Format(time.RFC3339), attempt)
126-
}
127-
}),
130+
httpClient.WithHTTPClient(&http.Client{Transport: rl}),
128131
)
129132

130133
// Process custom headers

cmd/preflight/preflight.go

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import (
1818
"github.com/buildkite/cli/v3/internal/build/watch"
1919
"github.com/buildkite/cli/v3/internal/cli"
2020
bkErrors "github.com/buildkite/cli/v3/internal/errors"
21+
bkhttp "github.com/buildkite/cli/v3/internal/http"
2122
"github.com/buildkite/cli/v3/internal/pipeline/resolver"
2223
"github.com/buildkite/cli/v3/internal/preflight"
2324
"github.com/buildkite/cli/v3/pkg/cmd/factory"
@@ -43,7 +44,8 @@ func (c *PreflightCmd) Help() string {
4344
}
4445

4546
func (c *PreflightCmd) Run(kongCtx *kong.Context, globals cli.GlobalFlags) error {
46-
f, err := newFactory(factory.WithDebug(globals.EnableDebug()))
47+
rlTransport := bkhttp.NewRateLimitTransport(http.DefaultTransport)
48+
f, err := newFactory(factory.WithDebug(globals.EnableDebug()), factory.WithTransport(rlTransport))
4749
if err != nil {
4850
return bkErrors.NewInternalError(err, "failed to initialize CLI", "This is likely a bug", "Report to Buildkite")
4951
}
@@ -90,6 +92,17 @@ func (c *PreflightCmd) Run(kongCtx *kong.Context, globals cli.GlobalFlags) error
9092

9193
renderer := newRenderer(os.Stdout, c.JSON, c.Text, stop)
9294

95+
rlTransport.OnRateLimit = func(attempt int, delay time.Duration) {
96+
if globals.EnableDebug() {
97+
_ = renderer.Render(Event{
98+
Type: EventOperation,
99+
Time: time.Now(),
100+
PreflightID: preflightID.String(),
101+
Title: fmt.Sprintf("Rate limited by API, waiting %s before retrying (attempt %d/%d)...", delay.Truncate(time.Second), attempt+1, rlTransport.MaxRetries),
102+
})
103+
}
104+
}
105+
93106
_ = renderer.Render(Event{Type: EventOperation, Time: time.Now(), PreflightID: preflightID.String(), Title: "Pushing snapshot of working tree..."})
94107

95108
var opts []preflight.SnapshotOption

go.mod

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@ require (
66
github.com/alecthomas/kong v1.15.0
77
github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be
88
github.com/buildkite/go-buildkite/v4 v4.19.0
9-
github.com/buildkite/roko v1.4.0
109
github.com/buildkite/termoji v0.0.0-20260330080310-c0aa4ebee0d1
1110
github.com/charmbracelet/bubbles v1.0.0
1211
github.com/charmbracelet/bubbletea v1.3.10

go.sum

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -35,8 +35,6 @@ github.com/bradleyjkemp/cupaloy/v2 v2.6.0 h1:knToPYa2xtfg42U3I6punFEjaGFKWQRXJwj
3535
github.com/bradleyjkemp/cupaloy/v2 v2.6.0/go.mod h1:bm7JXdkRd4BHJk9HpwqAI8BoAY1lps46Enkdqw6aRX0=
3636
github.com/buildkite/go-buildkite/v4 v4.19.0 h1:HPc6+V/Ky6v8eDWHxyvzHtxuhruCLUyuRrChOKvJE1I=
3737
github.com/buildkite/go-buildkite/v4 v4.19.0/go.mod h1:8+7GiWBKwEPAWoZnRU/kpNCt46j1iVH8kFMMbD4YDfc=
38-
github.com/buildkite/roko v1.4.0 h1:DxixoCdpNqxu4/1lXrXbfsKbJSd7r1qoxtef/TT2J80=
39-
github.com/buildkite/roko v1.4.0/go.mod h1:0vbODqUFEcVf4v2xVXRfZZRsqJVsCCHTG/TBRByGK4E=
4038
github.com/buildkite/termoji v0.0.0-20260330080310-c0aa4ebee0d1 h1:aaEl0QZURcwC+KOfFTzSp66xknw5eTmFZ1NgB87s2xk=
4139
github.com/buildkite/termoji v0.0.0-20260330080310-c0aa4ebee0d1/go.mod h1:ZTEvQlMN3+qzjROvjRb1p0X+xDQxxKpkMFhMSnaTrpw=
4240
github.com/cenkalti/backoff v2.2.1+incompatible h1:tNowT99t7UNflLxfYYSlKYsBpXdEet03Pg2g16Swow4=
@@ -225,5 +223,3 @@ gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
225223
gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
226224
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
227225
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
228-
gotest.tools/v3 v3.5.2 h1:7koQfIKdy+I8UTetycgUqXWSDwpgv193Ka+qRsmBY8Q=
229-
gotest.tools/v3 v3.5.2/go.mod h1:LtdLGcnqToBH83WByAAi/wiwSFCArdFIUV/xxN4pcjA=

internal/http/client.go

Lines changed: 1 addition & 91 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,7 @@ import (
88
"io"
99
"net/http"
1010
"net/url"
11-
"strconv"
1211
"strings"
13-
"time"
14-
15-
"github.com/buildkite/roko"
1612
)
1713

1814
// ErrorResponse represents an error response from the API
@@ -44,10 +40,6 @@ type Client struct {
4440
token string
4541
userAgent string
4642
client *http.Client
47-
48-
maxRetries int
49-
maxRetryDelay time.Duration
50-
onRetry OnRetryFunc
5143
}
5244

5345
// ClientOption is a function that modifies a Client
@@ -74,30 +66,6 @@ func WithHTTPClient(client *http.Client) ClientOption {
7466
}
7567
}
7668

77-
// WithMaxRetries sets the maximum number of retries for rate-limited requests.
78-
func WithMaxRetries(n int) ClientOption {
79-
return func(c *Client) {
80-
c.maxRetries = n
81-
}
82-
}
83-
84-
// WithMaxRetryDelay sets the maximum delay between retries
85-
func WithMaxRetryDelay(d time.Duration) ClientOption {
86-
return func(c *Client) {
87-
c.maxRetryDelay = d
88-
}
89-
}
90-
91-
// WithOnRetry sets a callback that is invoked before each retry sleep.
92-
func WithOnRetry(f OnRetryFunc) ClientOption {
93-
return func(c *Client) {
94-
c.onRetry = f
95-
}
96-
}
97-
98-
// OnRetryFunc is called before each retry sleep with the attempt number and delay duration.
99-
type OnRetryFunc func(attempt int, delay time.Duration)
100-
10169
// NewClient creates a new HTTP client with the given token and options
10270
func NewClient(token string, opts ...ClientOption) *Client {
10371
c := &Client{
@@ -164,23 +132,6 @@ func (e *ErrorResponse) IsTooManyRequests() bool {
164132
return e.StatusCode == http.StatusTooManyRequests
165133
}
166134

167-
// RetryAfter returns the duration to wait before retrying, based on the RateLimit-Reset header.
168-
// Returns 0 if the header is missing or invalid.
169-
func (e *ErrorResponse) RetryAfter() time.Duration {
170-
if e.Headers == nil {
171-
return 0
172-
}
173-
resetStr := e.Headers.Get("RateLimit-Reset")
174-
if resetStr == "" {
175-
return 0
176-
}
177-
seconds, err := strconv.Atoi(resetStr)
178-
if err != nil || seconds < 0 {
179-
return 0
180-
}
181-
return time.Duration(seconds) * time.Second
182-
}
183-
184135
// Do performs an HTTP request with the given method, endpoint, and body.
185136
func (c *Client) Do(ctx context.Context, method, endpoint string, body interface{}, v interface{}) error {
186137
// Ensure endpoint starts with "/"
@@ -215,21 +166,7 @@ func (c *Client) Do(ctx context.Context, method, endpoint string, body interface
215166
}
216167
}
217168

218-
r := roko.NewRetrier(
219-
roko.WithMaxAttempts(c.maxRetries+1),
220-
roko.WithStrategy(roko.Constant(0)),
221-
)
222-
223-
respBody, err := roko.DoFunc(ctx, r, func(r *roko.Retrier) ([]byte, error) {
224-
resp, err := c.send(ctx, method, reqURL, bodyBytes)
225-
if err != nil {
226-
if !c.handleRetry(r, err) {
227-
r.Break()
228-
}
229-
return nil, err
230-
}
231-
return resp, nil
232-
})
169+
respBody, err := c.send(ctx, method, reqURL, bodyBytes)
233170
if err != nil {
234171
return err
235172
}
@@ -284,30 +221,3 @@ func (c *Client) send(ctx context.Context, method, reqURL string, body []byte) (
284221

285222
return respBody, nil
286223
}
287-
288-
// handleRetry checks if an error is retryable and configures the retrier accordingly.
289-
// Returns true if the request should be retried, false otherwise.
290-
func (c *Client) handleRetry(r *roko.Retrier, err error) bool {
291-
errResp, ok := err.(*ErrorResponse)
292-
if !ok || !errResp.IsTooManyRequests() {
293-
return false
294-
}
295-
296-
attempt := r.AttemptCount()
297-
delay := errResp.RetryAfter()
298-
if attempt > 0 {
299-
// Got rate-limited again means contention - back off exponentially
300-
delay *= time.Duration(1 << attempt)
301-
}
302-
303-
if c.maxRetryDelay > 0 {
304-
delay = min(delay, c.maxRetryDelay)
305-
}
306-
307-
if c.onRetry != nil {
308-
c.onRetry(attempt, delay)
309-
}
310-
311-
r.SetNextInterval(delay)
312-
return true
313-
}

0 commit comments

Comments
 (0)