Spend a rate-limit token on every request, not every operation by robzolkos · Pull Request #909 · basecamp/basecamp-cli · GitHub
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion internal/appctx/context.go
1 change: 1 addition & 0 deletions internal/resilience/hooks.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
38 changes: 38 additions & 0 deletions internal/resilience/rate_limiter.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
45 changes: 45 additions & 0 deletions internal/resilience/transport.go
Original file line number Diff line number Diff line change
@@ -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
}
Comment thread
Copilot marked this conversation as resolved.
}
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}
}
213 changes: 213 additions & 0 deletions internal/resilience/transport_test.go
Loading