mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-21 05:57:53 +00:00
fix(routes): release body-admission slot after decode, not whole handler
Fix an availability regression the round-2 body-admission semaphore shipped (MUST-1, round 3, 2026-08-22 review): `defer release()` spanned the WHOLE handler in runUnified and all three media handlers, including the upstream round trip and response streaming. The stated goal was bounding decoded-body memory during the read+json.Unmarshal step; the shipped code instead bounded total in-flight requests, since an LLM completion is I/O-bound for seconds to minutes. Measured, zero-config, 400 concurrent chat requests: 220 succeeded, 180 got 503, on a cap a 2-vCPU pod would size at 16. Narrowed the held window to read+decode only, via new helpers that do the acquire/decode/release internally and return before the caller ever resolves a model or calls the adapter: - routes_unified.go: readAndDecodeUnifiedBody (runUnified no longer acquires/defers at all). - routes_media.go: decodeMediaJSONRequest now claims/releases its own slot (images.generations, audio.speech); a new readAdmittedCapped helper does the same around audio.transcriptions' readCapped call — narrower still, since that route never json.Unmarshal's into a map[string]any at all. Every release still goes through `defer` inside these narrow helpers, not a bare post-call statement, so a slot is released exactly once on every path including a panic unwinding through the read/decode itself. Since the held window shrank from "whole request" (seconds-minutes) to "one decode" (low milliseconds), raised the self-tuned default: bodyAdmissionCapPerCPU 8->24, floor 16->32, ceiling 128->256 — justified in llmgateway.go's doc comment against the measured adversarial-shape decode peak (~51.5MB LIVE heap, ~189MB TotalAlloc for a 4MiB adversarial body), not the old, no-longer-accurate ~170MB whole-pipeline estimate. Also clamped an explicit Config.MaxInFlightBodyRequests override to 10,000 (previously accepted verbatim, so e.g. a fat-fingered MaxInt would have silently disabled the bound). Corrected maxUnifiedRequestBytes' own DoS-math doc comment (SHOULD-3): the prior "~11x general case" framing was wrong — a realistic single inline base64 image measures ~1.0x (no amplification at all), while the adversarial many-small-values shape is what actually amplifies (12.3x LIVE / ~45x TotalAlloc, both now stated with real measured numbers instead of a single misleading multiplier). Documented two real, previously-silent behavior changes from finding 1a's admission-before-decode reordering (SHOULD-4): an over-budget caller whose request would otherwise have 400'd/404'd now gets 429 instead, and a malformed/unroutable request now consumes rate-limit budget it previously did not. Both are arguably improvements but were not visible from the diff alone. Replaced the sequential zero-config admission test, which could not actually exercise contention (MUST-2: the reviewer reran its exact shape with MaxInFlightBodyRequests=1 and it still passed), with a genuinely concurrent one: n = defaultBodyAdmissionCap()+20 real goroutines against a gated upstream, asserting (a) zero 503s across all n and (b) the semaphore holds zero slots while all n are parked in the upstream -- the regression's own regression test.
This commit is contained in:
+97
-13
@@ -2083,14 +2083,51 @@ func TestDefaultBodyAdmissionCap_ClampedToBounds(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestHandleChat_ZeroConfig_BodyAdmissionDoesNotThrottleOrdinaryTraffic is
|
||||
// the DEFAULT-PRESERVING gate's own explicit requirement: a config that
|
||||
// never sets MaxInFlightBodyRequests must behave exactly as before this
|
||||
// finding for ordinary, non-bursty traffic — a modest run of sequential
|
||||
// requests must never see a 503 from the new semaphore.
|
||||
func TestHandleChat_ZeroConfig_BodyAdmissionDoesNotThrottleOrdinaryTraffic(t *testing.T) {
|
||||
// TestHandleChat_ZeroConfig_ConcurrentAboveDefaultCap_NoThrottling_SlotsReleasedBeforeUpstream
|
||||
// is the DEFAULT-PRESERVING gate's own explicit requirement, replacing a
|
||||
// SEQUENTIAL version that could not actually exercise it (MUST-2, round
|
||||
// 3, 2026-08-22 review: the reviewer re-ran the sequential test's exact
|
||||
// shape with MaxInFlightBodyRequests=1 and it still passed — a
|
||||
// capacity-one semaphore is never contended by requests that arrive one
|
||||
// at a time, so it was not evidence the semaphore leaves ordinary
|
||||
// traffic unthrottled).
|
||||
//
|
||||
// This drives n = defaultBodyAdmissionCap()+20 requests genuinely
|
||||
// CONCURRENTLY (real goroutines, real network round trips to an
|
||||
// httptest.Server upstream — n is deliberately larger than the
|
||||
// self-tuned default cap) against an upstream gated to block until every
|
||||
// one of them has arrived, and asserts two things:
|
||||
//
|
||||
// 1. While all n requests are parked waiting on the gated upstream —
|
||||
// more in-flight requests than admission slots exist — the
|
||||
// semaphore (gw.bodyAdmission) holds ZERO slots. This is the
|
||||
// regression's own regression test (MUST-1, round 3): before that
|
||||
// fix, a slot was held for the WHOLE handler including this exact
|
||||
// upstream wait, so n > cap concurrent parked requests would have
|
||||
// left the semaphore fully occupied (and, with a large enough n
|
||||
// relative to cap, some later arrivals would have found it
|
||||
// exhausted). With the fix, a slot is released the instant
|
||||
// read+decode finishes, strictly before the upstream call ever
|
||||
// starts — by the time a request reaches the gated upstream, it is
|
||||
// provably holding no slot at all.
|
||||
// 2. Once every parked request completes, all n get 200 — none was
|
||||
// ever refused with 503 merely because n other requests happened to
|
||||
// be concurrently in flight, which is exactly what "zero-config
|
||||
// ordinary traffic is not throttled" has to mean on an I/O-bound
|
||||
// gateway where many requests are routinely in flight for seconds
|
||||
// at a time.
|
||||
func TestHandleChat_ZeroConfig_ConcurrentAboveDefaultCap_NoThrottling_SlotsReleasedBeforeUpstream(t *testing.T) {
|
||||
const respBody = `{"id":"c1","choices":[{"message":{"role":"assistant","content":"hi"}}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`
|
||||
|
||||
// gate blocks every upstream response until the test explicitly
|
||||
// closes it — simulating an LLM completion parked on a slow
|
||||
// upstream for a while, the exact scenario the MUST-1 regression
|
||||
// measured (a real completion is I/O-bound for seconds to minutes).
|
||||
gate := make(chan struct{})
|
||||
var arrived sync.WaitGroup
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
arrived.Done() // this request's admission slot is already released by now — see below
|
||||
<-gate
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte(respBody))
|
||||
@@ -2110,14 +2147,61 @@ func TestHandleChat_ZeroConfig_BodyAdmissionDoesNotThrottleOrdinaryTraffic(t *te
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
gw, ok := h.(*Gateway)
|
||||
if !ok {
|
||||
t.Fatal("handler is not *Gateway")
|
||||
}
|
||||
|
||||
for i := 0; i < 20; i++ {
|
||||
body := map[string]any{"model": "gpt-test", "messages": []any{}}
|
||||
req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body)
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("request %d: status = %d, want 200, body=%s", i, rec.Code, rec.Body.String())
|
||||
n := cap(gw.bodyAdmission) + 20
|
||||
arrived.Add(n)
|
||||
|
||||
statuses := make([]int, n)
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(n)
|
||||
for i := 0; i < n; i++ {
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
body := map[string]any{"model": "gpt-test", "messages": []any{}}
|
||||
req := newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body)
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
statuses[i] = rec.Code
|
||||
}(i)
|
||||
}
|
||||
|
||||
// Blocks until every one of the n requests has reached the gated
|
||||
// upstream — which, by this call chain's own synchronous ordering
|
||||
// (runUnified -> readAndDecodeUnifiedBody's deferred release ->
|
||||
// registry.resolve -> the adapter's own client.Do), can only happen
|
||||
// AFTER that request's admission slot has already been released.
|
||||
// Bounded by a timeout, not a bare Wait(): if the regression this
|
||||
// test guards against reappeared, fewer than n requests would ever
|
||||
// reach the upstream at all (the rest would be refused with an
|
||||
// immediate 503 — acquireBodyAdmission never blocks, see its own doc
|
||||
// comment — and never call arrived.Done()), which would hang a bare
|
||||
// Wait() forever instead of failing the test.
|
||||
arrivedCh := make(chan struct{})
|
||||
go func() {
|
||||
arrived.Wait()
|
||||
close(arrivedCh)
|
||||
}()
|
||||
select {
|
||||
case <-arrivedCh:
|
||||
case <-time.After(10 * time.Second):
|
||||
close(gate) // unstick any goroutine that did make it through, so the test can still finish reporting statuses
|
||||
t.Fatal("timed out waiting for all n requests to reach the gated upstream — some are likely being refused with 503 before ever reaching it, which is the regression this test guards against")
|
||||
}
|
||||
|
||||
if held := len(gw.bodyAdmission); held != 0 {
|
||||
t.Errorf("body-admission slots held while %d requests (cap=%d) are parked in the upstream = %d, want 0 — a request parked past decode must never still hold a slot", n, cap(gw.bodyAdmission), held)
|
||||
}
|
||||
|
||||
close(gate)
|
||||
wg.Wait()
|
||||
|
||||
for i, status := range statuses {
|
||||
if status != http.StatusOK {
|
||||
t.Errorf("request %d: status = %d, want 200 (cap=%d, n=%d concurrent) — ordinary traffic must never see 503 from the body-admission semaphore merely for being concurrent", i, status, cap(gw.bodyAdmission), n)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user