From 56e112a3ed85263c54546a2c4dacb5a47bedd32d Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Thu, 17 Sep 2026 15:23:01 +0100 Subject: [PATCH] perf(openai): presize upstream JSON response reads from Content-Length forwardJSON read upstream response bodies with io.ReadAll, which grows its buffer many times. readAllLimited returns the same data and errors, but it uses the declared Content-Length as a capacity hint. A body of known length then usually needs one allocation. A wrong Content-Length still reads to EOF or to the limit. The hint is clamped to 1 MiB, so a peer that declares a large length and then stalls cannot make the gateway allocate more than that before data arrives. The client request path keeps io.ReadAll, because the client controls its own Content-Length. Tests compare readAllLimited with io.ReadAll for exact, short, long, unknown and oversized lengths, for a mid-body error and for the clamp. With both commits, the yaegi harness measures 729565 bytes per embeddings request instead of 925587 (-21%), 12.6% fewer allocations, and 1.15 ms instead of 1.31 ms. With 16 embeddings per request, allocation falls from 5.0 MB to 3.0 MB. Claude-Session: https://claude.ai/code/session_01KwtV5ecajCrftcf8TkfPad --- provider_openai.go | 2 +- providers.go | 29 ++++++++++++ providers_test.go | 115 +++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 145 insertions(+), 1 deletion(-) diff --git a/provider_openai.go b/provider_openai.go index a11928e..e82d518 100644 --- a/provider_openai.go +++ b/provider_openai.go @@ -214,7 +214,7 @@ type chatUsagePayload struct { // since a malformed-but-readable body is the upstream's problem to answer // for, not a reason to swallow the response. func (a *openaiAdapter) forwardJSON(w http.ResponseWriter, resp *http.Response) (usage, error) { - body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes)) + body, err := readAllLimited(resp.Body, resp.ContentLength, maxResponseBytes) if err != nil { return usage{}, fmt.Errorf("%w: read response body: %w", errUpstream, err) } diff --git a/providers.go b/providers.go index 1ab1bd6..e5f3bff 100644 --- a/providers.go +++ b/providers.go @@ -80,6 +80,35 @@ func newProviderHTTPError(resp *http.Response) *providerHTTPError { return &providerHTTPError{status: resp.StatusCode, body: body} } +// readAllHintMax caps how much readAllLimited allocates up front from a +// declared Content-Length. A peer that declares a large length and then +// stalls must not make the gateway hold more than this before any byte +// arrives; a longer body still reads correctly, it just regrows past it. +const readAllHintMax = 1 << 20 + +// readAllLimited returns the same result as io.ReadAll(io.LimitReader(r, +// limit)); contentLength is only a capacity hint — a wrong value still +// reads to EOF or limit, never trusted for correctness. Used for upstream +// response bodies (forwardJSON), where the sender is the operator- +// configured provider, not an arbitrary client. When it looks usable, it +// lets a body of known length be read into a single allocation instead of +// io.ReadAll's repeated doubling regrowth, but the hint itself is clamped +// to readAllHintMax before allocating — the added bytes.MinRead keeps +// ReadFrom's final probe read (which always tries to read past what it +// expects) from itself forcing one more regrow. +func readAllLimited(r io.Reader, contentLength, limit int64) ([]byte, error) { + if contentLength <= 0 || contentLength > limit { + return io.ReadAll(io.LimitReader(r, limit)) + } + hint := contentLength + if hint > readAllHintMax { + hint = readAllHintMax + } + buf := bytes.NewBuffer(make([]byte, 0, hint+bytes.MinRead)) + _, err := buf.ReadFrom(io.LimitReader(r, limit)) + return buf.Bytes(), err +} + // errUpstream is a sentinel wrapped into errors an adapter returns when it // cannot reach, or cannot make sense of a response from, its upstream — // distinct from providerHTTPError, which means the upstream answered, just diff --git a/providers_test.go b/providers_test.go index 77a64f6..2bab1b8 100644 --- a/providers_test.go +++ b/providers_test.go @@ -1,9 +1,13 @@ package traefikllmgateway import ( + "bytes" "context" + "errors" + "io" "net/http" "testing" + "testing/iotest" ) // TestProviderPassthroughEnabled is the security+performance audit @@ -95,3 +99,114 @@ func TestWithAttemptRecorder_NilRecorder_ReturnsSameContext(t *testing.T) { t.Errorf("attemptRecorderFromContext = %v, want nil after a nil recorder was passed in", rec) } } + +// TestReadAllLimited compares readAllLimited against the reference +// io.ReadAll(io.LimitReader(r, limit)) on an identical input for every +// shape of contentLength hint: exact, too small, too large, unknown (-1), +// zero with a non-empty body, larger than limit, and a body longer than +// limit (which must still truncate to limit) both via the io.ReadAll +// fallback (hint <= 0 or hint > limit) and via the presized bytes.Buffer +// fast path directly (hint > 0 and hint <= limit) — the fast path's own +// io.LimitReader must still do the truncating, not just size the buffer. +func TestReadAllLimited(t *testing.T) { + cases := []struct { + name string + body []byte + contentLength int64 + limit int64 + }{ + {name: "exact content-length", body: []byte("hello world"), contentLength: 11, limit: 1024}, + {name: "hint smaller than body", body: []byte("hello world"), contentLength: 3, limit: 1024}, + {name: "hint larger than body", body: []byte("hi"), contentLength: 1000, limit: 1024}, + {name: "unknown content-length (-1)", body: []byte("hello"), contentLength: -1, limit: 1024}, + {name: "zero hint with non-empty body", body: []byte("hello"), contentLength: 0, limit: 1024}, + {name: "hint larger than limit", body: []byte("hello world"), contentLength: 2000, limit: 5}, + {name: "body longer than limit, hint exceeds limit (fallback path) truncates", body: bytes.Repeat([]byte("x"), 20), contentLength: 20, limit: 10}, + {name: "body longer than limit, unknown hint (fallback path) truncates", body: bytes.Repeat([]byte("x"), 20), contentLength: -1, limit: 10}, + {name: "body longer than limit, hint equals limit (fast path) truncates", body: bytes.Repeat([]byte("x"), 20), contentLength: 10, limit: 10}, + {name: "body longer than limit, hint below limit (fast path) truncates", body: bytes.Repeat([]byte("x"), 20), contentLength: 5, limit: 10}, + {name: "empty body", body: nil, contentLength: 0, limit: 1024}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + want, wantErr := io.ReadAll(io.LimitReader(bytes.NewReader(c.body), c.limit)) + got, gotErr := readAllLimited(bytes.NewReader(c.body), c.contentLength, c.limit) + + if (gotErr == nil) != (wantErr == nil) { + t.Fatalf("err = %v, want error presence %v", gotErr, wantErr != nil) + } + if !bytes.Equal(got, want) { + t.Errorf("readAllLimited(%q, %d, %d) = %q, want %q", c.body, c.contentLength, c.limit, got, want) + } + }) + } +} + +// TestReadAllLimited_ClampsHintToReadAllHintMax asserts a large declared +// Content-Length never drives an unbounded up-front allocation before any +// byte has arrived — an authenticated peer that declares a huge length and +// then stalls must not make readAllLimited hold more than readAllHintMax +// per call — while a genuinely large body (bigger than readAllHintMax) +// still reads back identically to the io.ReadAll(io.LimitReader(...)) +// reference, just via regrowth past the initial cap. +func TestReadAllLimited_ClampsHintToReadAllHintMax(t *testing.T) { + t.Run("huge declared length, tiny body: allocation stays capped", func(t *testing.T) { + body := []byte("0123456789") + const contentLength = 32 << 20 + const limit = 32 << 20 + + got, err := readAllLimited(bytes.NewReader(body), contentLength, limit) + if err != nil { + t.Fatalf("readAllLimited: %v", err) + } + if !bytes.Equal(got, body) { + t.Errorf("readAllLimited = %q, want %q", got, body) + } + if cap(got) > readAllHintMax+bytes.MinRead { + t.Errorf("cap(got) = %d, want <= %d: the declared length must not drive an unbounded up-front allocation", + cap(got), readAllHintMax+bytes.MinRead) + } + }) + + t.Run("body larger than readAllHintMax with exact hint still reads fully", func(t *testing.T) { + body := bytes.Repeat([]byte("y"), readAllHintMax+100) + contentLength := int64(len(body)) + limit := int64(len(body)) + + want, wantErr := io.ReadAll(io.LimitReader(bytes.NewReader(body), limit)) + if wantErr != nil { + t.Fatalf("reference io.ReadAll: %v", wantErr) + } + got, err := readAllLimited(bytes.NewReader(body), contentLength, limit) + if err != nil { + t.Fatalf("readAllLimited: %v", err) + } + if !bytes.Equal(got, want) { + t.Errorf("readAllLimited returned %d bytes, want %d bytes matching reference", len(got), len(want)) + } + }) +} + +// TestReadAllLimited_PropagatesNonEOFError asserts a reader that yields +// some data and then a non-EOF error fails readAllLimited the same way it +// fails the reference io.ReadAll(io.LimitReader(r, limit)) — the injected +// error must still be reachable via errors.Is, not swallowed. +func TestReadAllLimited_PropagatesNonEOFError(t *testing.T) { + injected := errors.New("stub: read failed") + newFailingReader := func() io.Reader { + return io.MultiReader(bytes.NewReader([]byte("partial")), iotest.ErrReader(injected)) + } + + _, refErr := io.ReadAll(io.LimitReader(newFailingReader(), 1024)) + if !errors.Is(refErr, injected) { + t.Fatalf("reference io.ReadAll(io.LimitReader) err = %v, want errors.Is match for %v", refErr, injected) + } + + _, gotErr := readAllLimited(newFailingReader(), 100, 1024) + if gotErr == nil { + t.Fatal("readAllLimited: want a non-nil error") + } + if !errors.Is(gotErr, injected) { + t.Errorf("readAllLimited err = %v, want errors.Is match for %v", gotErr, injected) + } +}