diff --git a/README.md b/README.md index 3f5b6ced..9dd85522 100644 --- a/README.md +++ b/README.md @@ -59,6 +59,43 @@ if err != nil { Operations like `Add`, `DeleteByID`, `GetByID`, and `Update` are supported by most services that are exposed through the client if not exposed in the package. These operations are configured at runtime since the Octopus REST API is hypermedia-driven. +### Rate limiting + +Octopus Server can [limit the rate of requests](https://octopus.com/docs/administration/managing-infrastructure/rate-limiting) from each client. When a client sends too many requests, the server returns `429 Too Many Requests` and does not process the request. + +By default, the client does not handle rate limits. To make the client obey the rate limits, create an `http.Client` with `ratelimit.NewHttpClient` and give it to the client: + +```go +import ( + "github.com/OctopusDeploy/go-octopusdeploy/v2/pkg/client" + "github.com/OctopusDeploy/go-octopusdeploy/v2/pkg/client/ratelimit" +) + +rateLimitedClient := ratelimit.NewHttpClient(nil, ratelimit.DefaultOptions()) + +octopusClient, err := client.NewClient(rateLimitedClient, apiURL, apiKey, spaceID) +if err != nil { + _ = fmt.Errorf("error creating API client: %v", err) + return +} +``` + +With rate limiting on, the client does these things: + +- When the server returns `429 Too Many Requests`, the client waits and sends the request again. It waits for the time in the `Retry-After` header, or for `DefaultDelay` if the server does not send this header. The delay doubles with each retry, and has random jitter. +- The client paces its requests with the `Octopus-RateLimit-Policy` and `Octopus-RateLimit` advisory headers, so that it stays in the rate limit. This prevents most `429` responses. The client uses GCRA (the Generic Cell Rate Algorithm), the same as the Octopus C# and TypeScript clients. + +You can change these settings in `ratelimit.Options`: + +| Setting | Default | Description | +| --------------------- | ------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `RetryCount` | `5` | The maximum number of times to send a request again after a `429` response. Set to `0` to turn off retries. | +| `DefaultDelay` | `2s` | The time to wait before a retry when the `429` response has no `Retry-After` header. If `DefaultDelay` is `0` or less, the default of `2s` is used. | +| `MaxDelay` | `30s` | The maximum time to wait before a retry. If `Retry-After` is longer than this time, the client does not retry and returns the `429` response. If `MaxDelay` is `0` or less, the default of `30s` is used. | +| `UseRateLimitHeaders` | `true` | Pace the requests with the advisory rate limit headers. | + +The client retries a request only if it can send the request body again. For example, the client does not retry a package upload. + Numerous code samples that showcase the API and this client are available in the [examples](/examples) directory. There are also many [integration](/test) and unit tests available to examine that demonstrate the capabilities of this API client. ## Testing diff --git a/examples/rate_limiting/create_rate_limited_client.go b/examples/rate_limiting/create_rate_limited_client.go new file mode 100644 index 00000000..ce10e7f5 --- /dev/null +++ b/examples/rate_limiting/create_rate_limited_client.go @@ -0,0 +1,54 @@ +package examples + +import ( + "fmt" + "net/url" + "time" + + "github.com/OctopusDeploy/go-octopusdeploy/v2/pkg/client" + "github.com/OctopusDeploy/go-octopusdeploy/v2/pkg/client/ratelimit" + "github.com/OctopusDeploy/go-octopusdeploy/v2/pkg/spaces" +) + +// CreateRateLimitedClientExample provides an example of how to create a Go API client +// that obeys the rate limits of Octopus Server. +func CreateRateLimitedClientExample() { + var ( + apiKey string = "API-YOUR_API_KEY" + octopusURL string = "https://your_octopus_url" + spaceID string = "space-id" + ) + + apiURL, err := url.Parse(octopusURL) + if err != nil { + _ = fmt.Errorf("error parsing URL for Octopus API: %v", err) + return + } + + // start from the recommended settings: 5 retries, a default delay of 2 seconds, + // a maximum delay of 30 seconds, and pacing with the rate limit headers + rateLimitOptions := ratelimit.DefaultOptions() + + // change the settings if necessary + rateLimitOptions.RetryCount = 10 + rateLimitOptions.MaxDelay = 60 * time.Second + + // the first parameter for ratelimit.NewHttpClient can accept a http.Client if you wish to + // override the default; ratelimit.NewHttpClient returns a copy and does not change it + httpClient := ratelimit.NewHttpClient(nil, rateLimitOptions) + + octopusClient, err := client.NewClient(httpClient, apiURL, apiKey, spaceID) + if err != nil { + _ = fmt.Errorf("error creating API client: %v", err) + return + } + + // all requests from this client now obey the rate limits + space, err := spaces.GetByID(octopusClient, spaceID) + if err != nil { + _ = fmt.Errorf("error getting space: %v", err) + return + } + + fmt.Printf("space: (%s)\n", space.GetID()) +} diff --git a/pkg/client/rate_limit_test.go b/pkg/client/rate_limit_test.go new file mode 100644 index 00000000..122aad39 --- /dev/null +++ b/pkg/client/rate_limit_test.go @@ -0,0 +1,73 @@ +package client + +import ( + "io" + "net/http" + "net/url" + "strings" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/OctopusDeploy/go-octopusdeploy/v2/pkg/client/ratelimit" + "github.com/OctopusDeploy/go-octopusdeploy/v2/pkg/constants" +) + +const testRateLimitAPIKey = "API-ABCDEFGHIJKLMNOPQRSTUVWXYZ" + +// newRateLimitedHttpClient returns an http.Client that returns 429 for the first request to the root resource. +// After that, it returns the root resource for /api or /api/ and an empty list for all other paths. +func newRateLimitedHttpClient() (*http.Client, func() int) { + var mutex sync.Mutex + rootRequestCount := 0 + + httpClient := &http.Client{ + Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) { + if strings.TrimSuffix(request.URL.Path, "/") == "/api" { + mutex.Lock() + rootRequestCount++ + count := rootRequestCount + mutex.Unlock() + + if count == 1 { + return &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{"Retry-After": []string{"0"}}, + Body: io.NopCloser(strings.NewReader(`{"ErrorMessage":"Too many requests"}`)), + }, nil + } + + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(constants.ApiReplyRoot)), + }, nil + } + + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(`{"Items":[]}`)), + }, nil + }), + } + + return httpClient, func() int { + mutex.Lock() + defer mutex.Unlock() + return rootRequestCount + } +} + +func TestNewClientWithRateLimitedHttpClientRetriesTooManyRequests(t *testing.T) { + httpClient, rootRequestCount := newRateLimitedHttpClient() + apiURL, err := url.Parse("https://octopus.example") + require.NoError(t, err) + rateLimitedHttpClient := ratelimit.NewHttpClient(httpClient, ratelimit.Options{RetryCount: 3, MaxDelay: time.Second}) + + octopusClient, err := NewClient(rateLimitedHttpClient, apiURL, testRateLimitAPIKey, "") + + require.NoError(t, err) + require.NotNil(t, octopusClient) + require.Equal(t, 2, rootRequestCount()) +} diff --git a/pkg/client/ratelimit/http_client.go b/pkg/client/ratelimit/http_client.go new file mode 100644 index 00000000..70bd7db6 --- /dev/null +++ b/pkg/client/ratelimit/http_client.go @@ -0,0 +1,19 @@ +package ratelimit + +import "net/http" + +// NewHttpClient returns a copy of base that obeys the rate limits of Octopus Server. +// The copy sends requests through a Transport that wraps the transport of base. +// If base is nil, NewHttpClient uses a new http.Client. NewHttpClient does not change base. +// +// The Timeout of base includes the time that the Transport waits before a retry or a paced request. +// If Timeout is less than MaxDelay, a request can fail with a timeout while the Transport waits. +func NewHttpClient(base *http.Client, options Options) *http.Client { + if base == nil { + base = &http.Client{} + } + + client := *base + client.Transport = NewTransport(base.Transport, options) + return &client +} diff --git a/pkg/client/ratelimit/http_client_test.go b/pkg/client/ratelimit/http_client_test.go new file mode 100644 index 00000000..c4f60bb9 --- /dev/null +++ b/pkg/client/ratelimit/http_client_test.go @@ -0,0 +1,36 @@ +package ratelimit + +import ( + "net/http" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestNewHttpClientWithNilBaseUsesDefaultTransport(t *testing.T) { + client := NewHttpClient(nil, testOptions()) + + require.NotNil(t, client) + require.IsType(t, &Transport{}, client.Transport) + require.Equal(t, http.DefaultTransport, client.Transport.(*Transport).Base()) +} + +func TestNewHttpClientWrapsTheTransportOfBase(t *testing.T) { + server := &fakeServer{responses: []func() *http.Response{newResponse(http.StatusOK, nil)}} + base := &http.Client{Transport: server} + + client := NewHttpClient(base, testOptions()) + + require.IsType(t, &Transport{}, client.Transport) + require.Equal(t, server, client.Transport.(*Transport).Base()) +} + +func TestNewHttpClientDoesNotChangeBase(t *testing.T) { + server := &fakeServer{responses: []func() *http.Response{newResponse(http.StatusOK, nil)}} + base := &http.Client{Transport: server} + + client := NewHttpClient(base, testOptions()) + + require.NotSame(t, base, client) + require.Equal(t, server, base.Transport) +} diff --git a/pkg/client/ratelimit/pacer.go b/pkg/client/ratelimit/pacer.go new file mode 100644 index 00000000..47461f77 --- /dev/null +++ b/pkg/client/ratelimit/pacer.go @@ -0,0 +1,187 @@ +package ratelimit + +import ( + "regexp" + "strconv" + "strings" + "sync" + "time" +) + +// reservedRequests is the number of requests of the allowance that the pacer does not spend. The server tells +// the remaining requests as a whole number and its bucket refills on a timer, thus the arithmetic of the server +// and the pacer can be different by a fraction of a request. When the pacer stops one request before the limit, +// that difference causes a small wait and not a 429 response. +const reservedRequests = 1 + +// numberPattern matches the numbers that the advisory headers can contain. strconv.ParseFloat also accepts +// values such as "0x10", "Inf" and "NaN", which the server does not send. +var numberPattern = regexp.MustCompile(`^[-+]?(\d+(\.\d*)?|\.\d+)([eE][-+]?\d+)?$`) + +// pacer keeps requests in the rate limit that the server advertises in the policyHeader and rateLimitHeader +// headers. Thus the client spreads its requests out and does not send them until the server rejects them. +// +// The pacer uses GCRA (the Generic Cell Rate Algorithm), which has the same behaviour as the token bucket of the +// server. The pacer keeps one value, the theoretical arrival time (TAT). The TAT is the time when the requests +// that the pacer already let go are paid for. A request can go when the TAT is in the burst tolerance from now, +// and each request moves the TAT one emission interval later. Thus the client can send a burst of requests and +// then continues at the steady rate of the policy. +// +// Each response sets the TAT again from the remaining requests that the server tells. Thus the pacer corrects +// itself when something else spends the same allowance. The pacer only moves the TAT later, never earlier, +// thus an old response that arrives after a newer one cannot make the client go faster. +// +// The pacer is the same as the RateLimitPacer of the C# and TypeScript clients. +// +// A pacer is safe for concurrent use by multiple goroutines. +type pacer struct { + // now returns a monotonic time. Tests replace it. + now func() time.Duration + + mutex sync.Mutex + havePolicy bool + burstLimit int + requestsPerMinute float64 + emissionInterval time.Duration + burstTolerance time.Duration + theoreticalArrivalTime time.Duration +} + +// newPacer returns a pacer that uses the monotonic clock of the time package. +func newPacer() *pacer { + start := time.Now() + return newPacerWithClock(func() time.Duration { + return time.Since(start) + }) +} + +// newPacerWithClock returns a pacer that uses the clock now. +func newPacerWithClock(now func() time.Duration) *pacer { + return &pacer{now: now} +} + +// reserveSlot takes the next slot in the schedule and returns the time to wait before the request is sent. +// The slot is taken also when the caller does not wait. Thus requests that start at the same time get +// consecutive slots, and they do not all go at the same time. +func (p *pacer) reserveSlot() time.Duration { + p.mutex.Lock() + defer p.mutex.Unlock() + + // Until the server tells the policy, there is no limit to pace against. + if !p.havePolicy { + return 0 + } + + now := p.now() + earliestSendTime := p.theoreticalArrivalTime - p.burstTolerance + + p.theoreticalArrivalTime = max(p.theoreticalArrivalTime, now) + p.emissionInterval + + return max(earliestSendTime-now, 0) +} + +// observeResponse reads the advisory headers of a response and sets the schedule again from them. +// Both values can be empty: an older server, or an endpoint that the server does not rate limit. +// Then the pacer keeps the policy that it has, and gives back the slot of the request (see refundSlot). +// +// Call observeResponse for each response, also for error responses. The server writes these headers before +// the endpoint runs, thus a 429 or a 500 response from a rate limited endpoint also tells the remaining requests. +func (p *pacer) observeResponse(policyHeaderValue string, rateLimitHeaderValue string) { + p.mutex.Lock() + defer p.mutex.Unlock() + + limit, haveLimit := getParameter(policyHeaderValue, "l") + rate, haveRate := getParameter(policyHeaderValue, "rpm") + + // Use only a policy that is valid, and only when it changes. The policy header is the same on almost all + // responses, and if the pacer applies it again, it loses the schedule that it keeps. + if haveLimit && haveRate && limit >= 1 && rate > 0 && + (!p.havePolicy || limit != float64(p.burstLimit) || rate != p.requestsPerMinute) { + p.burstLimit = int(limit) + p.requestsPerMinute = rate + p.emissionInterval = time.Duration(float64(time.Minute) / rate) + + // One request for the request that is sent, and one request that the pacer keeps back. + // The remaining requests of the bucket are the burst. + allowedBurst := p.burstLimit - 1 - reservedRequests + p.burstTolerance = 0 + if allowedBurst > 0 { + p.burstTolerance = p.emissionInterval * time.Duration(allowedBurst) + } + + // The old schedule uses the units of the old policy, thus it is not correct now. + // The remaining requests in this response set the schedule again below. + p.theoreticalArrivalTime = p.now() + + p.havePolicy = true + } + + if !p.havePolicy { + return + } + + remaining, haveRemaining := getParameter(rateLimitHeaderValue, "r") + if !haveRemaining { + p.refundSlot() + return + } + + remaining = min(max(remaining, 0), float64(p.burstLimit)) + + // The allowance that the server tells is spent, as a time in the schedule. The pacer only moves the + // schedule later, never earlier: responses can arrive out of order, and an older response that arrives + // last tells a bucket that is fuller than it is. + spent := time.Duration(float64(p.emissionInterval) * (float64(p.burstLimit) - remaining)) + p.theoreticalArrivalTime = max(p.theoreticalArrivalTime, p.now()+spent) +} + +// refundSlot gives back the slot that a request took, because the response tells that the request did not cost +// the limiter anything. A response without a remaining count is from an endpoint that the server does not rate +// limit, or the request did not get to the limiter. +// +// Without a refund, the pacer spends the schedule on requests that the server did not count. Then a client that +// sends requests to limited and unlimited endpoints paces all of its requests against a limit that only applies +// to some of them. +// +// A refund that is too large is safe: the next response that has a remaining count moves the schedule back. +// A schedule that is too small has no correction, because only time moves it earlier. +// +// The caller must hold the mutex. +func (p *pacer) refundSlot() { + // Only give back time that the schedule holds. A schedule that is already behind the clock does not pace + // anything, and if the pacer moves it further into the past, it gives credit for a burst that is not earned. + now := p.now() + if p.theoreticalArrivalTime > now { + p.theoreticalArrivalTime = max(p.theoreticalArrivalTime-p.emissionInterval, now) + } +} + +// getParameter reads one name=value parameter from a header such as "l=200;rpm=600". The name is not case-sensitive. +// It returns false if the header does not have the parameter, or if the value is not a number. +// It ignores parameters that it does not know, thus the server can add parameters without a change to the client. +func getParameter(headerValue string, name string) (float64, bool) { + for parameter := range strings.SplitSeq(headerValue, ";") { + parameterName, rawValue, found := strings.Cut(parameter, "=") + if !found { + continue + } + + if !strings.EqualFold(strings.TrimSpace(parameterName), name) { + continue + } + + rawValue = strings.TrimSpace(rawValue) + if !numberPattern.MatchString(rawValue) { + return 0, false + } + + value, err := strconv.ParseFloat(rawValue, 64) + if err != nil { + return 0, false + } + + return value, true + } + + return 0, false +} diff --git a/pkg/client/ratelimit/pacer_test.go b/pkg/client/ratelimit/pacer_test.go new file mode 100644 index 00000000..dc7c32ba --- /dev/null +++ b/pkg/client/ratelimit/pacer_test.go @@ -0,0 +1,382 @@ +package ratelimit + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +// A policy of 60 requests per minute makes one request cost exactly 1 second, which keeps the arithmetic in these +// tests easy to read. A burst limit of 10 gives 8 requests of burst room: one for the request that is sent, and one +// that the pacer keeps back. +const ( + testPolicy = "l=10;rpm=60" + testEmissionInterval = time.Second + testBurstRoom = 8 +) + +// newTestPacer returns a pacer with a clock that only moves when the test calls advance. Thus the arithmetic is exact. +func newTestPacer() (*pacer, func(time.Duration)) { + var now time.Duration + p := newPacerWithClock(func() time.Duration { + return now + }) + + return p, func(duration time.Duration) { + now += duration + } +} + +func TestPacerDoesNotPaceUntilTheServerTellsTheLimit(t *testing.T) { + p, _ := newTestPacer() + + // An older server or a policy without a limit gets here, and neither must make the client slower. + require.Equal(t, time.Duration(0), p.reserveSlot()) + require.Equal(t, time.Duration(0), p.reserveSlot()) +} + +func TestPacerDoesNotPaceWhenAResponseHasNoRateLimitHeaders(t *testing.T) { + p, _ := newTestPacer() + + // A server that has no policy for this caller sends no headers, and the pacer must not make up a limit. + p.observeResponse("", "") + + require.Equal(t, time.Duration(0), p.reserveSlot()) +} + +func TestPacerDoesNotWaitWhileRequestsRemain(t *testing.T) { + testCases := []struct { + name string + rateLimit string + }{ + {"FullBucket", "r=10;t=0"}, + {"HalfBucket", "r=5;t=5"}, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + p, _ := newTestPacer() + + p.observeResponse(testPolicy, testCase.rateLimit) + + require.Equal(t, time.Duration(0), p.reserveSlot()) + }) + } +} + +func TestPacerWaitsWhenOnlyTheReservedRequestRemains(t *testing.T) { + p, _ := newTestPacer() + + // If the client spends the last request, the server has nothing left for the rounding difference between + // its bucket and the pacer, thus the pacer waits for the bucket to refill one request first. + p.observeResponse(testPolicy, "r=1;t=9") + + require.Equal(t, testEmissionInterval, p.reserveSlot()) +} + +func TestPacerWaitsForTheBucketToRefillWhenNoRequestsRemain(t *testing.T) { + p, _ := newTestPacer() + + p.observeResponse(testPolicy, "r=0;t=10") + + // The pacer also waits for the request that it keeps back. + require.Equal(t, 2*testEmissionInterval, p.reserveSlot()) +} + +func TestPacerSendsTheBurstThenPacesAtTheSteadyRate(t *testing.T) { + p, _ := newTestPacer() + p.observeResponse(testPolicy, "r=10;t=0") + + // All of the requests in the burst go immediately, at the same time. + for i := range testBurstRoom + 1 { + require.Equal(t, time.Duration(0), p.reserveSlot(), "request %d is in the burst", i+1) + } + + // After the burst, each request waits for the bucket to refill for it. The requests wait one emission + // interval after each other, and they do not all get the same wait. + require.Equal(t, testEmissionInterval, p.reserveSlot()) + require.Equal(t, 2*testEmissionInterval, p.reserveSlot()) + require.Equal(t, 3*testEmissionInterval, p.reserveSlot()) +} + +func TestPacerCountsTimeThatPassedAgainstTheWait(t *testing.T) { + p, advance := newTestPacer() + p.observeResponse(testPolicy, "r=0;t=10") + + require.Equal(t, 2*testEmissionInterval, p.reserveSlot()) + + // The caller waited, thus the bucket refilled for the next request. + advance(2 * testEmissionInterval) + require.Equal(t, testEmissionInterval, p.reserveSlot()) + advance(3 * testEmissionInterval) + require.Equal(t, time.Duration(0), p.reserveSlot()) +} + +func TestPacerSlowsDownWhenTheServerTellsFewerRequestsThanExpected(t *testing.T) { + testCases := []struct { + name string + rateLimit string + expectedDelay time.Duration + }{ + {"OneRemains", "r=1;t=9", testEmissionInterval}, + {"NoneRemain", "r=0;t=10", 2 * testEmissionInterval}, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + p, _ := newTestPacer() + p.observeResponse(testPolicy, "r=10;t=0") + require.Equal(t, time.Duration(0), p.reserveSlot()) + + // Something else spends the same allowance, for example the web portal or a script of the same user. + p.observeResponse(testPolicy, testCase.rateLimit) + + require.Equal(t, testCase.expectedDelay, p.reserveSlot()) + }) + } +} + +func TestPacerDoesNotSpeedUpForAResponseThatArrivedOutOfOrder(t *testing.T) { + testCases := []struct { + name string + newer string + older string + expectedDelay time.Duration + }{ + {"NoneRemain", "r=0;t=10", "r=10;t=0", 2 * testEmissionInterval}, + {"OneRemains", "r=1;t=9", "r=9;t=1", testEmissionInterval}, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + p, _ := newTestPacer() + p.observeResponse(testPolicy, testCase.newer) + + // A response that left the server before the newer one, but arrived after it, tells a bucket that is + // fuller than it is. If the pacer used it, the client would spend requests that are already spent. + p.observeResponse(testPolicy, testCase.older) + + require.Equal(t, testCase.expectedDelay, p.reserveSlot()) + }) + } +} + +func TestPacerHandlesAPolicyWithNoRoomToBurst(t *testing.T) { + p, _ := newTestPacer() + + // A burst limit of 1 gives a negative burst room after the request that is sent and the request that the + // pacer keeps back. The burst room must be 0, so that the pacer does not pace into the past. + p.observeResponse("l=1;rpm=60", "r=1;t=0") + + require.Equal(t, time.Duration(0), p.reserveSlot()) + require.Equal(t, testEmissionInterval, p.reserveSlot()) +} + +func TestPacerWaitsWithABurstOfOneWhenNoRequestsRemain(t *testing.T) { + p, _ := newTestPacer() + + p.observeResponse("l=1;rpm=60", "r=1;t=0") + require.Equal(t, time.Duration(0), p.reserveSlot()) + + p.observeResponse("l=1;rpm=60", "r=0;t=1") + require.Equal(t, testEmissionInterval, p.reserveSlot()) +} + +func TestPacerStartsAgainWhenThePolicyChanges(t *testing.T) { + testCases := []struct { + name string + oldPolicy string + oldRateLimit string + expectedDelay time.Duration + newPolicy string + newRateLimit string + }{ + {"MoreRequests", testPolicy, "r=0;t=10", 2 * testEmissionInterval, "l=100;rpm=600", "r=100;t=0"}, + {"FasterRate", "l=10;rpm=6", "r=1;t=90", 10 * time.Second, testPolicy, "r=10;t=0"}, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + p, _ := newTestPacer() + p.observeResponse(testCase.oldPolicy, testCase.oldRateLimit) + require.Equal(t, testCase.expectedDelay, p.reserveSlot()) + + // An administrator changes the policy. The old schedule uses the units of the old policy, thus it is + // not correct now. The remaining requests in this response set the schedule again. + p.observeResponse(testCase.newPolicy, testCase.newRateLimit) + + require.Equal(t, time.Duration(0), p.reserveSlot()) + }) + } +} + +func TestPacerDoesNotPaceWhenThePolicyIsNotValid(t *testing.T) { + p, _ := newTestPacer() + + p.observeResponse("l=nonsense;rpm=0", "r=0;t=10") + + // The pacer has no limit to pace against. It is better to send the request and let the limiter decide + // than to make up a limit and keep back requests that the server would accept. + require.Equal(t, time.Duration(0), p.reserveSlot()) +} + +func TestPacerKeepsThePolicyWhenALaterPolicyHeaderIsNotValid(t *testing.T) { + p, _ := newTestPacer() + p.observeResponse(testPolicy, "r=10;t=0") + + p.observeResponse("", "r=0;t=10") + + require.Equal(t, 2*testEmissionInterval, p.reserveSlot()) +} + +func TestPacerGivesTheSlotBackWhenAResponseHasNoRemainingCount(t *testing.T) { + p, _ := newTestPacer() + p.observeResponse(testPolicy, "r=10;t=0") + + // Spend the burst, so that the next request must wait for the schedule. + for range testBurstRoom + 2 { + p.reserveSlot() + } + require.Equal(t, 2*testEmissionInterval, p.reserveSlot()) + + // Two of those requests went to endpoints that the server does not rate limit. They did not cost + // anything, thus the pacer gives back the schedule that they took. + p.observeResponse("", "") + p.observeResponse(testPolicy, "") + + require.Equal(t, testEmissionInterval, p.reserveSlot()) +} + +func TestPacerDoesNotGiveCreditForAFutureBurstFromRefunds(t *testing.T) { + testCases := []struct { + name string + advance time.Duration + }{ + // The schedule is level with the clock. + {"ScheduleLevelWithTheClock", 0}, + // The schedule fell behind the clock. + {"ScheduleBehindTheClock", 10 * testEmissionInterval}, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + p, advance := newTestPacer() + p.observeResponse(testPolicy, "r=10;t=0") + p.reserveSlot() + + // When the schedule is not ahead of the clock, there is nothing to give back, however many unlimited + // responses arrive. If not, a quiet period would give the client a burst that it did not earn. + advance(testCase.advance) + for range 20 { + p.observeResponse("", "") + } + + // The burst is the burst that the policy permits, and no longer. + for i := range testBurstRoom + 1 { + require.Equal(t, time.Duration(0), p.reserveSlot(), "request %d is in the burst", i+1) + } + require.Equal(t, testEmissionInterval, p.reserveSlot()) + }) + } +} + +func TestPacerMovesTheScheduleBackAfterARefundThatIsTooLarge(t *testing.T) { + p, _ := newTestPacer() + p.observeResponse(testPolicy, "r=0;t=10") + + // A request that the pacer refunds can be a request that got to the limiter, for example a cancelled request. + // A refund that is too large is safe, because the next response with a remaining count sets the schedule + // again. If these three refunds stayed, the request below would go immediately. + p.observeResponse("", "") + p.observeResponse("", "") + p.observeResponse("", "") + p.observeResponse(testPolicy, "r=0;t=10") + + require.Equal(t, 2*testEmissionInterval, p.reserveSlot()) +} + +func TestPacerStopsPacingWhenTheServerStopsTheRateLimit(t *testing.T) { + p, _ := newTestPacer() + p.observeResponse(testPolicy, "r=0;t=10") + require.Equal(t, 2*testEmissionInterval, p.reserveSlot()) + + // An administrator turns off the policy, thus the responses have no headers. The refunds remove the schedule, + // so that the client does not pace against a limit that is gone. + for range 11 { + p.observeResponse("", "") + } + + require.Equal(t, time.Duration(0), p.reserveSlot()) +} + +func TestPacerPacesAgainstAFractionalRequestRate(t *testing.T) { + p, _ := newTestPacer() + + // Policies use requests per minute, thus a slow policy has a fraction. + p.observeResponse("l=2;rpm=0.5", "r=0;t=240") + + require.Equal(t, 2*2*time.Minute, p.reserveSlot()) +} + +func TestGetParameterReadsTheValue(t *testing.T) { + testCases := []struct { + name string + header string + param string + expected float64 + }{ + {"BurstLimit", "l=200;rpm=600", "l", 200}, + {"RequestsPerMinute", "l=200;rpm=600", "rpm", 600}, + {"SecondsToRefill", "r=7;t=18", "t", 18}, + {"Fraction", "l=10;rpm=0.5", "rpm", 0.5}, + {"SpaceAfterSeparator", "l=200; rpm=600", "rpm", 600}, + {"SpacesAroundParameters", " l = 200 ; rpm = 600 ", "rpm", 600}, + {"NameIsNotCaseSensitive", "L=200", "l", 200}, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + value, found := getParameter(testCase.header, testCase.param) + + require.True(t, found) + require.Equal(t, testCase.expected, value) + }) + } +} + +func TestGetParameterReturnsFalseWhenTheValueIsNotThere(t *testing.T) { + testCases := []struct { + name string + header string + param string + }{ + {"EmptyHeader", "", "l"}, + {"NoSuchParameter", "l=200;rpm=600", "r"}, + {"EmptyName", "l=200", ""}, + {"NoSeparator", "nonsense", "l"}, + {"NameOnly", "l", "l"}, + {"EmptyValue", "l=", "l"}, + {"NotANumber", "l=nonsense", "l"}, + // strconv.ParseFloat reads these, but they are not what the server means. + {"Hexadecimal", "l=0x10", "l"}, + {"Infinity", "l=Inf", "l"}, + {"NotANumberLiteral", "l=NaN", "l"}, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + _, found := getParameter(testCase.header, testCase.param) + + require.False(t, found) + }) + } +} + +func TestGetParameterIgnoresParametersThatItDoesNotKnow(t *testing.T) { + // So that the server can add parameters to these headers without a change to the client. + value, found := getParameter("l=200;something=new;rpm=600", "rpm") + + require.True(t, found) + require.Equal(t, float64(600), value) +} diff --git a/pkg/client/ratelimit/retry.go b/pkg/client/ratelimit/retry.go new file mode 100644 index 00000000..0ae9f02a --- /dev/null +++ b/pkg/client/ratelimit/retry.go @@ -0,0 +1,78 @@ +package ratelimit + +import ( + "math" + "regexp" + "strconv" + "strings" + "time" +) + +// jitterFraction is the maximum random jitter that is added to a retry delay, as a fraction of the delay. +// The limiter tells all of the requests that it rejects in a burst to wait the same time. Without jitter, they +// all come back at the same time, and the requests that it rejects again stay together for all of their retries. +const jitterFraction = 0.25 + +// retryAfterPattern matches a Retry-After value that is a whole number of seconds. +var retryAfterPattern = regexp.MustCompile(`^\s*-?\d+\s*$`) + +// retryDelay returns the time to wait before a retry, from the Retry-After header value that the server sent and +// the number of retries that were already made. It returns false if the client must not retry, because the server +// tells it to wait longer than options.MaxDelay. +// +// attempt is 0 for the first retry. The delay doubles with each attempt. Retry-After tells the time until the server +// can accept one more request, not the time until it can accept all of the requests that it rejected. When the +// server rejects a burst of requests, a retry of all of them after the same delay lets only some of them through. +// The backoff gives the limiter time to refill for all of the requests. +// +// random is a value in the range [0, 1) that sets the jitter. +// +// The retry delay is the same as the getRateLimitRetryDelay function of the TypeScript client. +func retryDelay(retryAfterHeaderValue string, options Options, attempt int, random float64) (time.Duration, bool) { + baseDelay, found := parseRetryAfterHeader(retryAfterHeaderValue) + if !found { + baseDelay = options.DefaultDelay + } + + // A Retry-After of 0 or less tells the client that it can retry immediately. + baseDelay = max(baseDelay, 0) + + // The server tells the client to wait longer than it is permitted to wait, thus do not retry. + if baseDelay > options.MaxDelay { + return 0, false + } + + delay := math.Min(float64(baseDelay)*math.Pow(2, float64(attempt)), float64(options.MaxDelay)) + + // The jitter only makes the delay longer, never shorter: the server told the earliest time to retry. + jitter := math.Min(delay*jitterFraction, float64(options.MaxDelay)-delay) + return time.Duration(delay + random*jitter), true +} + +// parseRetryAfterHeader returns the delay in a Retry-After header value. It returns false if the value is empty or +// is not a whole number of seconds. Then the caller uses the default delay. +// +// Retry-After can also be an HTTP date, but parseRetryAfterHeader does not read dates, the same as the C# and +// TypeScript clients. A delay from a date uses the clock of the client, which can be wrong by minutes. If the clock +// is wrong in one direction, the client retries without a delay. If it is wrong in the other direction, the delay +// is longer than MaxDelay and the client does not retry. Octopus Server always sends a number of seconds. +func parseRetryAfterHeader(headerValue string) (time.Duration, bool) { + if !retryAfterPattern.MatchString(headerValue) { + return 0, false + } + + // The pattern only matches whole numbers, thus ParseInt can only fail when the value is out of range. + // Then ParseInt returns the largest or smallest int64, which the checks below handle. + seconds, _ := strconv.ParseInt(strings.TrimSpace(headerValue), 10, 64) + + // Prevent an overflow. A delay that is too long is longer than MaxDelay, and a negative delay is 0. + maxSeconds := int64(math.MaxInt64 / time.Second) + if seconds > maxSeconds { + return time.Duration(math.MaxInt64), true + } + if seconds < -maxSeconds { + return 0, true + } + + return time.Duration(seconds) * time.Second, true +} diff --git a/pkg/client/ratelimit/retry_test.go b/pkg/client/ratelimit/retry_test.go new file mode 100644 index 00000000..6baf7c7b --- /dev/null +++ b/pkg/client/ratelimit/retry_test.go @@ -0,0 +1,176 @@ +package ratelimit + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +// noJitter makes retryDelay add no jitter, so that the arithmetic is easy to read. Jitter has its own tests. +const noJitter = 0 + +// maxJitter makes retryDelay add almost all of the jitter that it can add. +const maxJitter = 0.999 + +func retryOptions() Options { + return Options{ + RetryCount: 5, + DefaultDelay: 2 * time.Second, + MaxDelay: 30 * time.Second, + } +} + +func TestRetryDelayWaitsForTheSecondsInRetryAfter(t *testing.T) { + delay, ok := retryDelay("5", retryOptions(), 0, noJitter) + + require.True(t, ok) + require.Equal(t, 5*time.Second, delay) +} + +func TestRetryDelayWaitsTheDefaultDelay(t *testing.T) { + testCases := []struct { + name string + retryAfter string + }{ + {"MissingRetryAfter", ""}, + {"RetryAfterIsNotANumber", "soon please"}, + // A malformed delay waits the default delay, and does not retry immediately. + {"Fraction", "5.5"}, + {"PlusSign", "+5"}, + {"FractionWithoutWholeNumber", ".5"}, + // retryDelay does not read Retry-After as a date. A client clock that is wrong by some minutes would + // give a delay of 0 and send too many requests, or a delay longer than MaxDelay and stop the retries. + {"FutureDate", "Sun, 30 Aug 2026 10:00:07 GMT"}, + {"PastDate", "Sun, 30 Aug 2026 09:59:50 GMT"}, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + delay, ok := retryDelay(testCase.retryAfter, retryOptions(), 0, noJitter) + + require.True(t, ok) + require.Equal(t, 2*time.Second, delay) + }) + } +} + +func TestRetryDelayRetriesImmediatelyWhenRetryAfterIsZeroOrNegative(t *testing.T) { + testCases := []struct { + name string + retryAfter string + }{ + {"Zero", "0"}, + {"Negative", "-3"}, + {"VeryNegative", "-99999999999999999999"}, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + delay, ok := retryDelay(testCase.retryAfter, retryOptions(), 0, noJitter) + + require.True(t, ok) + require.Equal(t, time.Duration(0), delay) + }) + } +} + +func TestRetryDelayGivesUpWhenTheServerTellsTheClientToWaitLongerThanMaxDelay(t *testing.T) { + _, ok := retryDelay("31", retryOptions(), 0, noJitter) + require.False(t, ok) + + delay, ok := retryDelay("30", retryOptions(), 0, noJitter) + require.True(t, ok) + require.Equal(t, 30*time.Second, delay) +} + +func TestRetryDelayGivesUpWhenRetryAfterIsTooLargeForADuration(t *testing.T) { + testCases := []struct { + name string + retryAfter string + }{ + {"LargerThanADuration", "9999999999999"}, + {"LargerThanAnInt64", "99999999999999999999"}, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + _, ok := retryDelay(testCase.retryAfter, retryOptions(), 0, noJitter) + + require.False(t, ok) + }) + } +} + +func TestRetryDelayDoublesTheWaitForEachRetry(t *testing.T) { + // The server tells all of the requests in a burst that it rejects to wait the same time. A retry after that + // time only lets through the requests that the limiter refills in that time, thus the client backs off to + // give the limiter time to refill for all of the requests that it rejected. + expectedDelays := []time.Duration{2 * time.Second, 4 * time.Second, 8 * time.Second, 16 * time.Second} + + for attempt, expectedDelay := range expectedDelays { + delay, ok := retryDelay("2", retryOptions(), attempt, noJitter) + + require.True(t, ok) + require.Equal(t, expectedDelay, delay, "attempt %d", attempt) + } +} + +func TestRetryDelayDoesNotBackOffPastMaxDelay(t *testing.T) { + delay, ok := retryDelay("2", retryOptions(), 10, noJitter) + + require.True(t, ok) + require.Equal(t, 30*time.Second, delay) +} + +func TestRetryDelayDoesNotOverflowForLargeAttempts(t *testing.T) { + delay, ok := retryDelay("2", retryOptions(), 1000, noJitter) + + require.True(t, ok) + require.Equal(t, 30*time.Second, delay) +} + +func TestRetryDelayBacksOffFromTheDefaultDelayWhenRetryAfterIsMissing(t *testing.T) { + delay, ok := retryDelay("", retryOptions(), 2, noJitter) + + require.True(t, ok) + require.Equal(t, 8*time.Second, delay) +} + +func TestRetryDelayAddsJitterAfterTheWaitNeverBeforeIt(t *testing.T) { + // The server told the earliest time to retry, thus jitter only makes the delay longer. + delay, ok := retryDelay("2", retryOptions(), 0, maxJitter) + + require.True(t, ok) + require.Greater(t, delay, 2*time.Second) + require.LessOrEqual(t, delay, 2500*time.Millisecond) +} + +func TestRetryDelayKeepsTheJitterInMaxDelay(t *testing.T) { + delay, ok := retryDelay("2", retryOptions(), 10, maxJitter) + require.True(t, ok) + require.Equal(t, 30*time.Second, delay) + + delay, ok = retryDelay("28", retryOptions(), 0, maxJitter) + require.True(t, ok) + require.LessOrEqual(t, delay, 30*time.Second) +} + +func TestWaitReturnsWhenTheDelayIsOver(t *testing.T) { + require.NoError(t, wait(context.Background(), 0)) +} + +func TestWaitReturnsTheErrorWhenTheContextIsAlreadyCancelled(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + require.ErrorIs(t, wait(ctx, 0), context.Canceled) +} + +func TestWaitReturnsTheErrorWhenTheContextIsCancelledDuringTheWait(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + time.AfterFunc(10*time.Millisecond, cancel) + + require.ErrorIs(t, wait(ctx, time.Minute), context.Canceled) +} diff --git a/pkg/client/ratelimit/transport.go b/pkg/client/ratelimit/transport.go new file mode 100644 index 00000000..6e4a4e5a --- /dev/null +++ b/pkg/client/ratelimit/transport.go @@ -0,0 +1,213 @@ +// Package ratelimit contains an http.RoundTripper that obeys the rate limits of Octopus Server. +// +// See https://octopus.com/docs/administration/managing-infrastructure/rate-limiting +package ratelimit + +import ( + "context" + "io" + "math/rand/v2" + "net/http" + "time" +) + +const ( + // policyHeader is the advisory header that tells the burst limit and the requests per minute. + // For example: "l=200;rpm=600". + policyHeader = "Octopus-RateLimit-Policy" + + // rateLimitHeader is the advisory header that tells the remaining requests and the seconds until all of the + // requests are available again. For example: "r=7;t=18". + rateLimitHeader = "Octopus-RateLimit" + + retryAfterHeader = "Retry-After" +) + +// Options are the settings for a Transport. +type Options struct { + // RetryCount is the maximum number of times that the Transport sends a request again after a 429 response. + // If RetryCount is 0 or less, the Transport does not retry, and the caller gets the 429 response. + RetryCount int + + // DefaultDelay is the time to wait before a retry when the 429 response has no Retry-After header. + // If DefaultDelay is 0 or less, the Transport uses the default of 2 seconds from DefaultOptions. + DefaultDelay time.Duration + + // MaxDelay is the maximum time to wait before a retry. If the server tells the Transport to wait longer, + // the Transport does not retry, and the caller gets the 429 response. + // If MaxDelay is 0 or less, the Transport uses the default of 30 seconds from DefaultOptions. + MaxDelay time.Duration + + // UseRateLimitHeaders makes the Transport pace its requests with the Octopus-RateLimit-Policy and + // Octopus-RateLimit headers that the server sends. Then the Transport stays in the rate limit and + // prevents most 429 responses. + UseRateLimitHeaders bool +} + +// DefaultOptions returns the recommended settings for a Transport. +func DefaultOptions() Options { + return Options{ + RetryCount: 5, + DefaultDelay: 2 * time.Second, + MaxDelay: 30 * time.Second, + UseRateLimitHeaders: true, + } +} + +// Transport is an http.RoundTripper that obeys the rate limits of Octopus Server. +// +// When the server returns 429 (Too Many Requests), the Transport waits and sends the request again. +// The server does not process a request that it rejects, thus a retry is safe for all HTTP methods. +// The Transport retries a request only if it can read the request body again (the body is empty, or +// http.Request.GetBody is set). The delay doubles with each retry, and has random jitter. +// +// When UseRateLimitHeaders is true, the Transport also paces each request, retries included, with the advisory +// headers that the server sends. The pacer uses GCRA (the Generic Cell Rate Algorithm), the same as the C# and +// TypeScript clients. +// +// Use NewTransport to create a Transport. A Transport is safe for concurrent use by multiple goroutines. +// Do not copy a Transport after its first use. +type Transport struct { + base http.RoundTripper + options Options + + // pacer is nil when UseRateLimitHeaders is false. + pacer *pacer +} + +var _ http.RoundTripper = (*Transport)(nil) + +// NewTransport returns a Transport that sends requests with base. If base is nil, the Transport uses http.DefaultTransport. +// If options.DefaultDelay or options.MaxDelay is 0 or less, the Transport uses the value from DefaultOptions. +func NewTransport(base http.RoundTripper, options Options) *Transport { + if base == nil { + base = http.DefaultTransport + } + + if options.DefaultDelay <= 0 { + options.DefaultDelay = DefaultOptions().DefaultDelay + } + + if options.MaxDelay <= 0 { + options.MaxDelay = DefaultOptions().MaxDelay + } + + transport := &Transport{ + base: base, + options: options, + } + + if options.UseRateLimitHeaders { + transport.pacer = newPacer() + } + + return transport +} + +// RoundTrip implements http.RoundTripper. +func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) { + for attempt := 0; ; attempt++ { + // The pacer also paces retries: a request that the limiter rejected must not go back to it immediately. + if t.pacer != nil { + if err := wait(req.Context(), t.pacer.reserveSlot()); err != nil { + // The base transport did not get the request, so the body must be closed here. + if attempt == 0 && req.Body != nil { + _ = req.Body.Close() + } + return nil, err + } + } + + attemptRequest, err := requestForAttempt(req, attempt) + if err != nil { + return nil, err + } + + resp, err := t.base.RoundTrip(attemptRequest) + if err != nil { + return nil, err + } + + if t.pacer != nil { + t.pacer.observeResponse(resp.Header.Get(policyHeader), resp.Header.Get(rateLimitHeader)) + } + + if resp.StatusCode != http.StatusTooManyRequests || attempt >= t.options.RetryCount || !canRetry(req) { + return resp, nil + } + + delay, ok := retryDelay(resp.Header.Get(retryAfterHeader), t.options, attempt, rand.Float64()) + if !ok { + return resp, nil + } + + closeResponse(resp) + if err := wait(req.Context(), delay); err != nil { + return nil, err + } + } +} + +// Base returns the http.RoundTripper that the Transport uses to send requests. +func (t *Transport) Base() http.RoundTripper { + return t.base +} + +// CloseIdleConnections closes the idle connections of the base transport, if the base transport supports it. +// http.Client.CloseIdleConnections calls this method. +func (t *Transport) CloseIdleConnections() { + if closer, ok := t.base.(interface{ CloseIdleConnections() }); ok { + closer.CloseIdleConnections() + } +} + +// canRetry returns true if the body of req can be read again. +func canRetry(req *http.Request) bool { + return req.Body == nil || req.Body == http.NoBody || req.GetBody != nil +} + +// requestForAttempt returns req for the first attempt. For a retry, it returns a copy of req with a new body. +func requestForAttempt(req *http.Request, attempt int) (*http.Request, error) { + if attempt == 0 { + return req, nil + } + + retryRequest := req.Clone(req.Context()) + if req.Body != nil && req.Body != http.NoBody { + body, err := req.GetBody() + if err != nil { + return nil, err + } + retryRequest.Body = body + } + + return retryRequest, nil +} + +// wait stops for the time in delay, or until ctx is cancelled. It returns the error of ctx if ctx is already +// cancelled, also when delay is 0, thus the Transport does not send a request that nobody waits for. +func wait(ctx context.Context, delay time.Duration) error { + if err := ctx.Err(); err != nil { + return err + } + + if delay <= 0 { + return nil + } + + timer := time.NewTimer(delay) + defer timer.Stop() + + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +// closeResponse reads all of the response body and closes it, so that the connection can be used again. +func closeResponse(resp *http.Response) { + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() +} diff --git a/pkg/client/ratelimit/transport_test.go b/pkg/client/ratelimit/transport_test.go new file mode 100644 index 00000000..b010a800 --- /dev/null +++ b/pkg/client/ratelimit/transport_test.go @@ -0,0 +1,479 @@ +package ratelimit + +import ( + "bytes" + "context" + "io" + "net/http" + "strings" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +const testURL = "https://octopus.example/api" + +type roundTripFunc func(req *http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} + +// fakeServer returns the responses in order and records each request that it gets. +// After the last response, it returns the last response again. +type fakeServer struct { + mutex sync.Mutex + responses []func() *http.Response + requests []*http.Request + bodies []string +} + +func (s *fakeServer) RoundTrip(req *http.Request) (*http.Response, error) { + s.mutex.Lock() + defer s.mutex.Unlock() + + body := "" + if req.Body != nil { + content, err := io.ReadAll(req.Body) + if err != nil { + return nil, err + } + body = string(content) + } + + s.requests = append(s.requests, req) + s.bodies = append(s.bodies, body) + + index := min(len(s.requests)-1, len(s.responses)-1) + return s.responses[index](), nil +} + +func (s *fakeServer) requestCount() int { + s.mutex.Lock() + defer s.mutex.Unlock() + return len(s.requests) +} + +func newResponse(statusCode int, headers map[string]string) func() *http.Response { + return func() *http.Response { + header := http.Header{} + for name, value := range headers { + header.Set(name, value) + } + + return &http.Response{ + StatusCode: statusCode, + Header: header, + Body: io.NopCloser(strings.NewReader("")), + } + } +} + +// tooManyRequests returns a 429 response with the Retry-After value. If retryAfter is empty, the response has no Retry-After header. +func tooManyRequests(retryAfter string) func() *http.Response { + if retryAfter == "" { + return newResponse(http.StatusTooManyRequests, nil) + } + + return newResponse(http.StatusTooManyRequests, map[string]string{retryAfterHeader: retryAfter}) +} + +// testOptions are the settings for the transport tests. The delays are short, so that the tests are fast. +func testOptions() Options { + return Options{ + RetryCount: 3, + DefaultDelay: 50 * time.Millisecond, + MaxDelay: time.Second, + UseRateLimitHeaders: true, + } +} + +func newGetRequest(t *testing.T) *http.Request { + req, err := http.NewRequest(http.MethodGet, testURL, nil) + require.NoError(t, err) + return req +} + +func TestNewTransportWithNilBaseUsesDefaultTransport(t *testing.T) { + transport := NewTransport(nil, testOptions()) + + require.Equal(t, http.DefaultTransport, transport.Base()) +} + +func TestBaseReturnsTheBaseTransport(t *testing.T) { + server := &fakeServer{responses: []func() *http.Response{newResponse(http.StatusOK, nil)}} + + transport := NewTransport(server, testOptions()) + + require.Same(t, server, transport.Base()) +} + +func TestNewTransportUsesTheDefaultDelaysWhenTheDelaysAreNotPositive(t *testing.T) { + testCases := []struct { + name string + delay time.Duration + }{ + {"Zero", 0}, + {"Negative", -time.Second}, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + transport := NewTransport(nil, Options{RetryCount: 3, DefaultDelay: testCase.delay, MaxDelay: testCase.delay}) + + require.Equal(t, DefaultOptions().DefaultDelay, transport.options.DefaultDelay) + require.Equal(t, DefaultOptions().MaxDelay, transport.options.MaxDelay) + }) + } +} + +func TestNewTransportKeepsPositiveDelays(t *testing.T) { + transport := NewTransport(nil, Options{RetryCount: 3, DefaultDelay: time.Second, MaxDelay: 5 * time.Second}) + + require.Equal(t, time.Second, transport.options.DefaultDelay) + require.Equal(t, 5*time.Second, transport.options.MaxDelay) +} + +func TestRoundTripReturnsSuccessfulResponse(t *testing.T) { + server := &fakeServer{responses: []func() *http.Response{newResponse(http.StatusOK, nil)}} + transport := NewTransport(server, testOptions()) + + resp, err := transport.RoundTrip(newGetRequest(t)) + + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Equal(t, 1, server.requestCount()) +} + +func TestRoundTripDoesNotRetryOtherErrorResponses(t *testing.T) { + server := &fakeServer{responses: []func() *http.Response{newResponse(http.StatusServiceUnavailable, nil)}} + transport := NewTransport(server, testOptions()) + + resp, err := transport.RoundTrip(newGetRequest(t)) + + require.NoError(t, err) + require.Equal(t, http.StatusServiceUnavailable, resp.StatusCode) + require.Equal(t, 1, server.requestCount()) +} + +func TestRoundTripRetriesUntilTheServerAcceptsTheRequest(t *testing.T) { + // The server rejects the first two requests, then accepts the request. + server := &fakeServer{responses: []func() *http.Response{ + tooManyRequests("0"), + tooManyRequests("0"), + newResponse(http.StatusOK, nil), + }} + transport := NewTransport(server, testOptions()) + + resp, err := transport.RoundTrip(newGetRequest(t)) + + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Equal(t, 3, server.requestCount()) +} + +func TestRoundTripWaitsTheDefaultDelayWhenThereIsNoRetryAfterHeader(t *testing.T) { + server := &fakeServer{responses: []func() *http.Response{ + tooManyRequests(""), + newResponse(http.StatusOK, nil), + }} + transport := NewTransport(server, testOptions()) + + start := time.Now() + resp, err := transport.RoundTrip(newGetRequest(t)) + + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Equal(t, 2, server.requestCount()) + require.GreaterOrEqual(t, time.Since(start), testOptions().DefaultDelay) +} + +func TestRoundTripGivesUpAfterTheConfiguredNumberOfRetries(t *testing.T) { + server := &fakeServer{responses: []func() *http.Response{tooManyRequests("0")}} + transport := NewTransport(server, testOptions()) + + resp, err := transport.RoundTrip(newGetRequest(t)) + + require.NoError(t, err) + require.Equal(t, http.StatusTooManyRequests, resp.StatusCode) + // The first request, and RetryCount (3) retries. + require.Equal(t, 4, server.requestCount()) +} + +func TestRoundTripDoesNotRetryWhenRetryCountIsZeroOrLess(t *testing.T) { + testCases := []struct { + name string + retryCount int + }{ + {"Zero", 0}, + {"Negative", -1}, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + server := &fakeServer{responses: []func() *http.Response{ + tooManyRequests("0"), + newResponse(http.StatusOK, nil), + }} + options := testOptions() + options.RetryCount = testCase.retryCount + transport := NewTransport(server, options) + + resp, err := transport.RoundTrip(newGetRequest(t)) + + require.NoError(t, err) + require.Equal(t, http.StatusTooManyRequests, resp.StatusCode) + require.Equal(t, 1, server.requestCount()) + }) + } +} + +func TestRoundTripDoesNotRetryWhenRetryAfterIsLongerThanMaxDelay(t *testing.T) { + server := &fakeServer{responses: []func() *http.Response{ + tooManyRequests("120"), + newResponse(http.StatusOK, nil), + }} + transport := NewTransport(server, testOptions()) + + resp, err := transport.RoundTrip(newGetRequest(t)) + + require.NoError(t, err) + require.Equal(t, http.StatusTooManyRequests, resp.StatusCode) + require.Equal(t, 1, server.requestCount()) +} + +func TestRoundTripSendsTheSameBodyOnRetry(t *testing.T) { + server := &fakeServer{responses: []func() *http.Response{ + tooManyRequests("0"), + newResponse(http.StatusOK, nil), + }} + transport := NewTransport(server, testOptions()) + + req, err := http.NewRequest(http.MethodPost, testURL, bytes.NewReader([]byte(`{"Name":"test"}`))) + require.NoError(t, err) + + resp, err := transport.RoundTrip(req) + + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Equal(t, []string{`{"Name":"test"}`, `{"Name":"test"}`}, server.bodies) + require.Equal(t, http.MethodPost, server.requests[1].Method) +} + +func TestRoundTripDoesNotRetryWhenBodyCannotBeReadAgain(t *testing.T) { + server := &fakeServer{responses: []func() *http.Response{ + tooManyRequests("0"), + newResponse(http.StatusOK, nil), + }} + transport := NewTransport(server, testOptions()) + + req, err := http.NewRequest(http.MethodPost, testURL, io.NopCloser(strings.NewReader("package content"))) + require.NoError(t, err) + require.Nil(t, req.GetBody) + + resp, err := transport.RoundTrip(req) + + require.NoError(t, err) + require.Equal(t, http.StatusTooManyRequests, resp.StatusCode) + require.Equal(t, 1, server.requestCount()) +} + +func TestRoundTripReturnsErrorWhenContextIsCancelledDuringRetryWait(t *testing.T) { + server := &fakeServer{responses: []func() *http.Response{ + tooManyRequests("30"), + newResponse(http.StatusOK, nil), + }} + options := testOptions() + options.MaxDelay = time.Minute + transport := NewTransport(server, options) + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + req, err := http.NewRequestWithContext(ctx, http.MethodGet, testURL, nil) + require.NoError(t, err) + + resp, err := transport.RoundTrip(req) + + require.ErrorIs(t, err, context.DeadlineExceeded) + require.Nil(t, resp) + require.Equal(t, 1, server.requestCount()) +} + +func TestRoundTripReturnsErrorFromBaseTransport(t *testing.T) { + baseErr := io.ErrUnexpectedEOF + transport := NewTransport(roundTripFunc(func(req *http.Request) (*http.Response, error) { + return nil, baseErr + }), testOptions()) + + resp, err := transport.RoundTrip(newGetRequest(t)) + + require.ErrorIs(t, err, baseErr) + require.Nil(t, resp) +} + +func TestRoundTripGivesTheRateLimitHeadersToThePacer(t *testing.T) { + testCases := []struct { + name string + statusCode int + }{ + {"SuccessfulResponse", http.StatusOK}, + // The server writes these headers before the endpoint runs, thus error responses also have them. + {"ErrorResponse", http.StatusInternalServerError}, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + server := &fakeServer{responses: []func() *http.Response{ + newResponse(testCase.statusCode, map[string]string{policyHeader: testPolicy, rateLimitHeader: "r=0;t=10"}), + }} + transport := NewTransport(server, testOptions()) + testPacer, _ := newTestPacer() + transport.pacer = testPacer + + _, err := transport.RoundTrip(newGetRequest(t)) + + require.NoError(t, err) + require.Equal(t, 2*testEmissionInterval, testPacer.reserveSlot()) + }) + } +} + +func TestRoundTripWaitsForThePacer(t *testing.T) { + // 600 requests per minute gives one request each 100 milliseconds. No requests remain, thus the + // next request waits for two requests: the request that it sends, and the request that the pacer keeps back. + server := &fakeServer{responses: []func() *http.Response{ + newResponse(http.StatusOK, map[string]string{policyHeader: "l=10;rpm=600", rateLimitHeader: "r=0;t=1"}), + }} + transport := NewTransport(server, testOptions()) + + _, err := transport.RoundTrip(newGetRequest(t)) + require.NoError(t, err) + + start := time.Now() + _, err = transport.RoundTrip(newGetRequest(t)) + + require.NoError(t, err) + require.GreaterOrEqual(t, time.Since(start), 150*time.Millisecond) +} + +func TestRoundTripDoesNotPaceWhenUseRateLimitHeadersIsFalse(t *testing.T) { + // 1 request per minute and no requests remain. With pacing, the next request must wait for two minutes. + server := &fakeServer{responses: []func() *http.Response{ + newResponse(http.StatusOK, map[string]string{policyHeader: "l=10;rpm=1", rateLimitHeader: "r=0;t=600"}), + }} + options := testOptions() + options.UseRateLimitHeaders = false + transport := NewTransport(server, options) + + _, err := transport.RoundTrip(newGetRequest(t)) + require.NoError(t, err) + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + req, err := http.NewRequestWithContext(ctx, http.MethodGet, testURL, nil) + require.NoError(t, err) + + resp, err := transport.RoundTrip(req) + + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Equal(t, 2, server.requestCount()) +} + +func TestRoundTripPacesRetries(t *testing.T) { + // The 429 response tells the client to retry immediately, but no requests remain, thus the pacer makes the + // retry wait. A request that the limiter rejected must not go back to it immediately. + server := &fakeServer{responses: []func() *http.Response{ + newResponse(http.StatusTooManyRequests, map[string]string{retryAfterHeader: "0", policyHeader: "l=10;rpm=600", rateLimitHeader: "r=0;t=1"}), + newResponse(http.StatusOK, nil), + }} + transport := NewTransport(server, testOptions()) + + start := time.Now() + resp, err := transport.RoundTrip(newGetRequest(t)) + + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Equal(t, 2, server.requestCount()) + require.GreaterOrEqual(t, time.Since(start), 150*time.Millisecond) +} + +func TestRoundTripReturnsErrorWhenContextIsCancelledDuringPacerWait(t *testing.T) { + // 1 request per minute and no requests remain, thus the next request must wait for two minutes. + server := &fakeServer{responses: []func() *http.Response{ + newResponse(http.StatusOK, map[string]string{policyHeader: "l=10;rpm=1", rateLimitHeader: "r=0;t=600"}), + }} + transport := NewTransport(server, testOptions()) + + _, err := transport.RoundTrip(newGetRequest(t)) + require.NoError(t, err) + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + req, err := http.NewRequestWithContext(ctx, http.MethodGet, testURL, nil) + require.NoError(t, err) + + resp, err := transport.RoundTrip(req) + + require.ErrorIs(t, err, context.DeadlineExceeded) + require.Nil(t, resp) + require.Equal(t, 1, server.requestCount()) +} + +type closeRecorder struct { + io.Reader + closed bool +} + +func (c *closeRecorder) Close() error { + c.closed = true + return nil +} + +func TestRoundTripClosesBodyWhenContextIsCancelledBeforeFirstAttempt(t *testing.T) { + server := &fakeServer{responses: []func() *http.Response{newResponse(http.StatusOK, nil)}} + transport := NewTransport(server, testOptions()) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + body := &closeRecorder{Reader: strings.NewReader(`{"Name":"test"}`)} + req, err := http.NewRequestWithContext(ctx, http.MethodPost, testURL, body) + require.NoError(t, err) + + _, err = transport.RoundTrip(req) + + require.ErrorIs(t, err, context.Canceled) + require.True(t, body.closed) + require.Equal(t, 0, server.requestCount()) +} + +// idleConnectionCloser is a base transport that records the calls to CloseIdleConnections. +type idleConnectionCloser struct { + closeCount int +} + +func (c *idleConnectionCloser) RoundTrip(req *http.Request) (*http.Response, error) { + return newResponse(http.StatusOK, nil)(), nil +} + +func (c *idleConnectionCloser) CloseIdleConnections() { + c.closeCount++ +} + +func TestCloseIdleConnectionsCallsTheBaseTransport(t *testing.T) { + base := &idleConnectionCloser{} + transport := NewTransport(base, testOptions()) + + transport.CloseIdleConnections() + + require.Equal(t, 1, base.closeCount) +} + +func TestCloseIdleConnectionsDoesNothingWhenTheBaseTransportDoesNotSupportIt(t *testing.T) { + server := &fakeServer{responses: []func() *http.Response{newResponse(http.StatusOK, nil)}} + transport := NewTransport(server, testOptions()) + + require.NotPanics(t, transport.CloseIdleConnections) +} diff --git a/pkg/newclient/httpsession.go b/pkg/newclient/httpsession.go index b1e42819..36c47482 100644 --- a/pkg/newclient/httpsession.go +++ b/pkg/newclient/httpsession.go @@ -73,6 +73,9 @@ func (h *HttpSession) DoRawJsonRequest(req *http.Request, requestBody any, outpu req.Header.Set("Content-Type", "application/json") } req.Body = io.NopCloser(bytes.NewReader(body)) + req.GetBody = func() (io.ReadCloser, error) { + return io.NopCloser(bytes.NewReader(body)), nil + } } if req.Header.Get("Accept") == "" { diff --git a/pkg/newclient/httpsession_test.go b/pkg/newclient/httpsession_test.go new file mode 100644 index 00000000..bc637e7a --- /dev/null +++ b/pkg/newclient/httpsession_test.go @@ -0,0 +1,48 @@ +package newclient + +import ( + "io" + "net/http" + "net/url" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +type roundTripFunc func(req *http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} + +func TestDoRawJsonRequestSetsGetBody(t *testing.T) { + var sentRequest *http.Request + httpSession := &HttpSession{ + HttpClient: &http.Client{ + Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + sentRequest = req + return &http.Response{ + StatusCode: http.StatusNoContent, + Body: io.NopCloser(strings.NewReader("")), + }, nil + }), + }, + } + + requestURL, err := url.Parse("https://octopus.example/api/projects") + require.NoError(t, err) + req := &http.Request{Method: http.MethodPost, URL: requestURL} + + _, err = httpSession.DoRawJsonRequest(req, map[string]string{"Name": "test"}, nil, nil) + require.NoError(t, err) + + require.NotNil(t, sentRequest.GetBody) + for range 2 { + body, err := sentRequest.GetBody() + require.NoError(t, err) + content, err := io.ReadAll(body) + require.NoError(t, err) + require.JSONEq(t, `{"Name":"test"}`, string(content)) + } +}