fix(observability-proxy): answer 429 to rate and concurrency limits

Clients such as Grafana and browser retry logic back off on 429 Too
Many Requests, not on 422, which they treat as a final rejection of
the query. The per-user rate limit, the per-user concurrency limit and
the open-tail caps are about the caller's pace, not the query, so they
now answer 429. The rate limiter knows when its bucket refills, so that
answer carries Retry-After; the concurrency caps cannot know when a
slot frees and omit it. Range, step, bucket and limit violations keep
422.

Assisted-by: LLM
Signed-off-by: ZverGuy <maximbel2003@gmail.com>

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

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
+}

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()
 

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() }()