diff --git a/internal/appctx/context.go b/internal/appctx/context.go index 4dfaa7d29..2faa4b86a 100644 --- a/internal/appctx/context.go +++ b/internal/appctx/context.go @@ -134,7 +134,7 @@ func NewApp(cfg *config.Config) *App { } sdkOptions := []basecamp.ClientOption{ basecamp.WithHooks(hooks), - basecamp.WithTransport(transport), + basecamp.WithTransport(gatingHooks.Transport(transport)), basecamp.WithUserAgent(version.UserAgent() + " " + basecamp.DefaultUserAgent), } sdkClient := basecamp.NewClient(sdkCfg, &authAdapter{mgr: authMgr}, sdkOptions...) diff --git a/internal/resilience/hooks.go b/internal/resilience/hooks.go index 05ef17fef..e8dd74ebb 100644 --- a/internal/resilience/hooks.go +++ b/internal/resilience/hooks.go @@ -93,6 +93,7 @@ func (h *GatingHooks) OnOperationGate(ctx context.Context, op basecamp.Operation if err := h.rateLimiter.waitSince(ctx, start, deadline); err != nil { return ctx, err } + ctx = withPrepaidRequest(ctx) } // Acquire bulkhead slot (PID-based, released in OnOperationEnd) diff --git a/internal/resilience/rate_limiter.go b/internal/resilience/rate_limiter.go index 7e9931d04..afd858ef5 100644 --- a/internal/resilience/rate_limiter.go +++ b/internal/resilience/rate_limiter.go @@ -118,6 +118,44 @@ func (rl *RateLimiter) take() (allowed bool, wait time.Duration, blocked bool) { return allowed, wait, blocked } +func (rl *RateLimiter) takeToken() (allowed bool, wait time.Duration) { + err := rl.store.Update(func(state *State) error { + now := rl.now() + rl.refill(&state.RateLimiter, now) + if state.RateLimiter.Tokens >= rl.config.TokensPerRequest { + state.RateLimiter.Tokens -= rl.config.TokensPerRequest + allowed = true + } else { + deficit := rl.config.TokensPerRequest - state.RateLimiter.Tokens + wait = time.Duration(deficit / rl.config.RefillRate * float64(time.Second)) + } + state.UpdatedAt = now + return nil + }) + if err != nil { + return true, 0 + } + return allowed, wait +} + +func (rl *RateLimiter) Pace(ctx context.Context) error { + for { + if err := ctx.Err(); err != nil { + return err + } + allowed, wait := rl.takeToken() //nolint:contextcheck // lock acquisition is context-independent by design + if allowed { + return nil + } + if rl.onWait != nil { + rl.onWait() + } + if err := pause(ctx, jittered(max(wait, minRefillWait))); err != nil { + return err + } + } +} + // minRefillWait floors the sleep between token attempts: a deficit of a few // microseconds is not worth a wakeup, and under contention the refill is // consumed by whichever process reaches the lock first anyway. diff --git a/internal/resilience/transport.go b/internal/resilience/transport.go new file mode 100644 index 000000000..74083ba57 --- /dev/null +++ b/internal/resilience/transport.go @@ -0,0 +1,45 @@ +package resilience + +import ( + "context" + "net/http" + "sync/atomic" +) + +type prepaidKey struct{} + +func withPrepaidRequest(ctx context.Context) context.Context { + prepaid := &atomic.Bool{} + prepaid.Store(true) + return context.WithValue(ctx, prepaidKey{}, prepaid) +} + +func spendPrepaidRequest(ctx context.Context) bool { + prepaid, ok := ctx.Value(prepaidKey{}).(*atomic.Bool) + return ok && prepaid.CompareAndSwap(true, false) +} + +type pacedTransport struct { + base http.RoundTripper + rateLimiter *RateLimiter +} + +func (t *pacedTransport) RoundTrip(req *http.Request) (*http.Response, error) { + ctx := req.Context() + if !spendPrepaidRequest(ctx) { + if err := t.rateLimiter.Pace(ctx); err != nil { + if req.Body != nil { + _ = req.Body.Close() + } + return nil, err + } + } + return t.base.RoundTrip(req) +} + +func (h *GatingHooks) Transport(base http.RoundTripper) http.RoundTripper { + if h.rateLimiter == nil { + return base + } + return &pacedTransport{base: base, rateLimiter: h.rateLimiter} +} diff --git a/internal/resilience/transport_test.go b/internal/resilience/transport_test.go new file mode 100644 index 000000000..db69337d0 --- /dev/null +++ b/internal/resilience/transport_test.go @@ -0,0 +1,213 @@ +package resilience + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/basecamp/basecamp-sdk/go/pkg/basecamp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func recordingsServer(t *testing.T, pages int, failFirst int) (*httptest.Server, *atomic.Int32) { + t.Helper() + var hits atomic.Int32 + var failures atomic.Int32 + var srv *httptest.Server + srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + if int(failures.Add(1)) <= failFirst { + w.WriteHeader(http.StatusServiceUnavailable) + return + } + page, _ := strconv.Atoi(r.URL.Query().Get("page")) + if page == 0 { + page = 1 + } + if page < pages { + q := r.URL.Query() + q.Set("page", strconv.Itoa(page+1)) + w.Header().Set("Link", fmt.Sprintf("<%s%s?%s>; rel=\"next\"", srv.URL, r.URL.Path, q.Encode())) + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode([]map[string]any{{ + "id": page, "status": "active", "type": "Todo", "title": "t", + "created_at": "2026-01-01T00:00:00Z", "updated_at": "2026-01-01T00:00:00Z", + }}) + })) + t.Cleanup(srv.Close) + return srv, &hits +} + +func pacedClient(t *testing.T, baseURL string, opts ...basecamp.ClientOption) (*basecamp.AccountClient, *RateLimiter) { + t.Helper() + cfg := DefaultConfig() + cfg.RateLimiter.MaxTokens = 100 + cfg.RateLimiter.RefillRate = 0.0001 + hooks := NewGatingHooksFromConfig(NewStore(t.TempDir()), cfg) + opts = append([]basecamp.ClientOption{ + basecamp.WithHooks(hooks), + basecamp.WithTransport(hooks.Transport(http.DefaultTransport)), + }, opts...) + client := basecamp.NewClient(&basecamp.Config{BaseURL: baseURL}, &basecamp.StaticTokenProvider{Token: "test-token"}, opts...) + return client.ForAccount("99999"), hooks.rateLimiter +} + +func tokensSpent(t *testing.T, rl *RateLimiter, run func()) float64 { + t.Helper() + before, err := rl.Tokens() + require.NoError(t, err) + run() + after, err := rl.Tokens() + require.NoError(t, err) + return before - after +} + +func TestEveryPageOfAWalkSpendsAToken(t *testing.T) { + srv, hits := recordingsServer(t, 5, 0) + account, rl := pacedClient(t, srv.URL) + + spent := tokensSpent(t, rl, func() { + result, err := account.Recordings().List(context.Background(), basecamp.RecordingTypeTodo, &basecamp.RecordingsListOptions{Limit: -1}) + require.NoError(t, err) + require.Len(t, result.Recordings, 5) + }) + + assert.Equal(t, int32(5), hits.Load()) + assert.InDelta(t, 5, spent, 0.01) +} + +func TestARawRequestSpendsAToken(t *testing.T) { + srv, hits := recordingsServer(t, 1, 0) + account, rl := pacedClient(t, srv.URL) + + spent := tokensSpent(t, rl, func() { + for range 3 { + _, err := account.Get(context.Background(), "/projects/recordings.json?type=Todo") + require.NoError(t, err) + } + }) + + assert.Equal(t, int32(3), hits.Load()) + assert.InDelta(t, 3, spent, 0.01) +} + +func TestARetrySpendsAToken(t *testing.T) { + srv, hits := recordingsServer(t, 1, 1) + account, rl := pacedClient(t, srv.URL, basecamp.WithBaseDelay(time.Millisecond)) + + spent := tokensSpent(t, rl, func() { + _, err := account.Get(context.Background(), "/projects/recordings.json?type=Todo") + require.NoError(t, err) + }) + + assert.Equal(t, int32(2), hits.Load()) + assert.InDelta(t, 2, spent, 0.01) +} + +func TestARetriedOperationSpendsATokenPerAttempt(t *testing.T) { + srv, hits := recordingsServer(t, 1, 1) + account, rl := pacedClient(t, srv.URL, basecamp.WithBaseDelay(time.Millisecond)) + + spent := tokensSpent(t, rl, func() { + _, err := account.Recordings().List(context.Background(), basecamp.RecordingTypeTodo, nil) + require.NoError(t, err) + }) + + assert.Equal(t, int32(2), hits.Load()) + assert.InDelta(t, 2, spent, 0.01) +} + +type closeRecorder struct { + io.Reader + closed atomic.Bool +} + +func (c *closeRecorder) Close() error { + c.closed.Store(true) + return nil +} + +func TestACanceledPaceClosesTheRequestBodyAndSendsNothing(t *testing.T) { + var hits atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { hits.Add(1) })) + t.Cleanup(srv.Close) + + rl := NewRateLimiter(NewStore(t.TempDir()), RateLimiterConfig{MaxTokens: 1, RefillRate: 0.0001, TokensPerRequest: 1}) + require.NoError(t, rl.Pace(context.Background())) + transport := (&GatingHooks{rateLimiter: rl}).Transport(http.DefaultTransport) + + ctx, cancel := context.WithCancel(context.Background()) + previousPause := pause + pause = func(ctx context.Context, d time.Duration) error { + cancel() + return previousPause(ctx, d) + } + t.Cleanup(func() { pause = previousPause }) + + body := &closeRecorder{Reader: strings.NewReader("{}")} + req, err := http.NewRequestWithContext(ctx, http.MethodPost, srv.URL, body) + require.NoError(t, err) + + resp, err := transport.RoundTrip(req) + if resp != nil { + _ = resp.Body.Close() + } + + assert.ErrorIs(t, err, context.Canceled) + assert.True(t, body.closed.Load()) + assert.Equal(t, int32(0), hits.Load()) +} + +func TestASinglePageOperationSpendsOneToken(t *testing.T) { + srv, hits := recordingsServer(t, 1, 0) + account, rl := pacedClient(t, srv.URL) + + spent := tokensSpent(t, rl, func() { + _, err := account.Recordings().List(context.Background(), basecamp.RecordingTypeTodo, nil) + require.NoError(t, err) + }) + + assert.Equal(t, int32(1), hits.Load()) + assert.InDelta(t, 1, spent, 0.01) +} + +func TestPaceWaitsForARefillInsteadOfFailing(t *testing.T) { + clock := holdClock(t, 0) + rl := NewRateLimiter(NewStore(t.TempDir()), RateLimiterConfig{MaxTokens: 1, RefillRate: 2, TokensPerRequest: 1}) + rl.clock = clock.Now + start := clock.Now() + waits := 0 + rl.onWait = func() { waits++ } + + require.NoError(t, rl.Pace(context.Background())) + require.NoError(t, rl.Pace(context.Background())) + require.NoError(t, rl.Pace(context.Background())) + + assert.Equal(t, 2, waits) + assert.Equal(t, start.Add(time.Second), clock.Now()) +} + +func TestPaceStopsWhenTheCallerCancels(t *testing.T) { + rl := NewRateLimiter(NewStore(t.TempDir()), RateLimiterConfig{MaxTokens: 1, RefillRate: 0.0001, TokensPerRequest: 1}) + require.NoError(t, rl.Pace(context.Background())) + + ctx, cancel := context.WithCancel(context.Background()) + previousPause := pause + pause = func(ctx context.Context, d time.Duration) error { + cancel() + return previousPause(ctx, d) + } + t.Cleanup(func() { pause = previousPause }) + + assert.ErrorIs(t, rl.Pace(ctx), context.Canceled) +}