diff --git a/internal/observabilityproxy/limits.go b/internal/observabilityproxy/limits.go index 588a4f29e..f0392f25b 100644 --- a/internal/observabilityproxy/limits.go +++ b/internal/observabilityproxy/limits.go @@ -58,9 +58,11 @@ func newLimiter(cfg LimiterConfig) *limiter { } } -func (l *limiter) allowRate(user string) bool { +// allowRate takes a token from the user's bucket; when it is empty, it +// returns how long until the next token. +func (l *limiter) allowRate(user string) (time.Duration, bool) { if l.cfg.RatePerUser <= 0 { - return true + return 0, true } l.mu.Lock() defer l.mu.Unlock() @@ -77,10 +79,10 @@ func (l *limiter) allowRate(user string) bool { b.tokens = min(burst, b.tokens+now.Sub(b.last).Seconds()*l.cfg.RatePerUser) b.last = now if b.tokens < 1 { - return false + return time.Duration((1 - b.tokens) / l.cfg.RatePerUser * float64(time.Second)), false } b.tokens-- - return true + return 0, true } // sweepBuckets drops buckets that have refilled completely; recreating one diff --git a/internal/observabilityproxy/limits_test.go b/internal/observabilityproxy/limits_test.go index 167598971..5c3af772d 100644 --- a/internal/observabilityproxy/limits_test.go +++ b/internal/observabilityproxy/limits_test.go @@ -76,29 +76,35 @@ func TestLimiterRate(t *testing.T) { now := time.Unix(1000, 0) l := newLimiter(LimiterConfig{RatePerUser: 2, RateBurst: 2}) l.now = func() time.Time { return now } - if !l.allowRate("alice") || !l.allowRate("alice") { + if !allowed(l, "alice") || !allowed(l, "alice") { t.Fatal("burst refused") } - if l.allowRate("alice") { - t.Fatal("request over the burst admitted") + if wait, ok := l.allowRate("alice"); ok || wait != 500*time.Millisecond { + t.Fatalf("request over the burst: admitted %v, wait %s, want refused with 500ms", ok, wait) } - if !l.allowRate("bob") { + if !allowed(l, "bob") { t.Fatal("bob limited by alice's rate") } now = now.Add(500 * time.Millisecond) - if !l.allowRate("alice") { + if !allowed(l, "alice") { t.Fatal("token not refilled after 1/rate") } - if l.allowRate("alice") { - t.Fatal("refilled more than elapsed*rate") + now = now.Add(100 * time.Millisecond) + if wait, ok := l.allowRate("alice"); ok || wait != 400*time.Millisecond { + t.Fatalf("partly refilled bucket: admitted %v, wait %s, want refused with 400ms", ok, wait) } } func TestLimiterRateDisabled(t *testing.T) { l := newLimiter(LimiterConfig{}) for range 100 { - if !l.allowRate("alice") { + if !allowed(l, "alice") { t.Fatal("rate limit applied while disabled") } } } + +func allowed(l *limiter, user string) bool { + _, ok := l.allowRate(user) + return ok +} diff --git a/internal/observabilityproxy/server.go b/internal/observabilityproxy/server.go index 4ea5cebee..fcdb46141 100644 --- a/internal/observabilityproxy/server.go +++ b/internal/observabilityproxy/server.go @@ -23,6 +23,7 @@ import ( "fmt" "io" "log/slog" + "math" "net/http" "strconv" "strings" @@ -150,8 +151,9 @@ func (s *Server) serve(w http.ResponseWriter, r *http.Request, route Route) *HTT w.Header().Set("Allow", "GET, POST") return errorf(http.StatusMethodNotAllowed, "method not allowed") } - if !s.limits.allowRate(user) { - return errorf(http.StatusUnprocessableEntity, "request rate limit exceeded") + if wait, ok := s.limits.allowRate(user); !ok { + w.Header().Set("Retry-After", strconv.FormatInt(int64(math.Ceil(wait.Seconds())), 10)) + return errorf(http.StatusTooManyRequests, "request rate limit exceeded") } endpoint, ok := strings.CutPrefix(route.Endpoint, "select/logsql/") if !ok { @@ -173,7 +175,7 @@ func (s *Server) serve(w http.ResponseWriter, r *http.Request, route Route) *HTT } release, ok := acquire(user) if !ok { - return errorf(http.StatusUnprocessableEntity, "too many concurrent requests") + return errorf(http.StatusTooManyRequests, "too many concurrent requests") } defer release() diff --git a/internal/observabilityproxy/server_test.go b/internal/observabilityproxy/server_test.go index 62d809f28..53c2aaa2d 100644 --- a/internal/observabilityproxy/server_test.go +++ b/internal/observabilityproxy/server_test.go @@ -252,8 +252,8 @@ func TestServerLimits(t *testing.T) { }) first := h.do(http.MethodGet, logsRoot+"query", "good", nil, nil) defer func() { _ = first.Body.Close() }() - if got := h.status(http.MethodGet, logsRoot+"query", "good"); got != http.StatusUnprocessableEntity { - t.Fatalf("concurrent request answered %d, want 422", got) + if got := h.status(http.MethodGet, logsRoot+"query", "good"); got != http.StatusTooManyRequests { + t.Fatalf("concurrent request answered %d, want 429", got) } once.Do(func() { close(release) }) @@ -261,8 +261,11 @@ func TestServerLimits(t *testing.T) { if got := h2.status(http.MethodGet, logsRoot+"query", "good"); got != http.StatusOK { t.Fatalf("first request answered %d", got) } - if got := h2.status(http.MethodGet, logsRoot+"query", "good"); got != http.StatusUnprocessableEntity { - t.Fatalf("rate-limited request answered %d, want 422", got) + limited := h2.do(http.MethodGet, logsRoot+"query", "good", nil, nil) + _ = limited.Body.Close() + if limited.StatusCode != http.StatusTooManyRequests || limited.Header.Get("Retry-After") != "1000" { + t.Fatalf("rate-limited request answered %d with Retry-After %q, want 429 with 1000", + limited.StatusCode, limited.Header.Get("Retry-After")) } } @@ -336,8 +339,8 @@ func TestServerTailCapsAndIdle(t *testing.T) { if a.StatusCode != http.StatusOK { t.Fatalf("first tail answered %d", a.StatusCode) } - if got := h.status(http.MethodGet, logsRoot+"tail", "good"); got != http.StatusUnprocessableEntity { - t.Fatalf("second tail of one user answered %d, want 422", got) + if got := h.status(http.MethodGet, logsRoot+"tail", "good"); got != http.StatusTooManyRequests { + t.Fatalf("second tail of one user answered %d, want 429", got) } b := h.do(http.MethodGet, logsRoot+"tail", "bob-token", nil, nil) defer func() { _ = b.Body.Close() }()