perf(registry): cache the encoded /v1/models response per group

GET /v1/models cost 15.9ms, 21.3MB, 552,884 allocs per call at the
live 1,033-model catalog (registry.go's listFor plus routes_unified.go's
json.Encoder), uncached — and OpenAI SDK clients poll this endpoint per
session, so an unchanged catalog was being re-walked and re-marshaled
on every request.

Add modelRegistry.modelsJSON(grp), caching the fully encoded
{object:list,data:[...]} body per (*group, modelsGen). modelsGen
is an atomic counter bumped by a new modelRegistry.finishRefresh
wrapper around providerState.finishRefresh — the only two places that
ever mutate a provider's discovered set or discovery-captured metadata
(warmFill, refreshProvider) already funnel through providerState.
finishRefresh once per refresh cycle, so wrapping it is sufficient to
invalidate the cache on every possible change, including a
metadata-only update on an otherwise-failed refresh (captureModelMetadata
runs unconditionally alongside finishRefresh in both call sites).
Bumped unconditionally rather than only on success, since a failed
listModels fetch can still pair with a successful metadata capture.

Cached per *group pointer: authStore builds every group once, at
construction, and never rebuilds or replaces that map afterward — a
config reload constructs an entirely new Gateway (and modelRegistry)
with its own empty cache, so a *group key can never alias a different
group's catalog either within one registry's lifetime or across a
reload.

routes_unified.go's handleModels call site shrinks to a single write of
modelsJSON's result, replacing the previous inline listFor+json.Encode;
routes_unified.go is otherwise untouched (owned by a sibling change in
this same effort).
This commit is contained in:
2026-08-22 14:36:48 +01:00
parent 90eb58171e
commit bf33732123
3 changed files with 229 additions and 14 deletions
+111 -9
View File
@@ -2,11 +2,13 @@ package traefikllmgateway
import (
"context"
"encoding/json"
"errors"
"fmt"
"sort"
"strings"
"sync"
"sync/atomic"
"time"
)
@@ -199,8 +201,10 @@ func (st *providerState) snapshot() (models []string, lastRefresh time.Time, las
// discovered model ids into the gateway's model-resolution and /v1/models
// listing surfaces. adapters, providerNames, and states are all fixed at
// construction; only providerState's own fields mutate afterward (each
// guarded by its own mutex), so modelRegistry itself needs no lock beyond
// warnedMu for the collision-log dedup below.
// guarded by its own mutex). modelRegistry itself needs warnedMu for the
// collision-log dedup below, and modelsCacheMu for modelsCache (perf
// finding 3) — modelsGen is read/written only via sync/atomic and needs
// no mutex of its own.
type modelRegistry struct {
adapters map[string]providerAdapter
states map[string]*providerState
@@ -213,7 +217,9 @@ type modelRegistry struct {
// shape through a method or interface method is unaffected — this is
// specific to a field holding a func value. Verified empirically
// against yaegi v0.16.1 (see tools/yaegi-check) and does not apply to
// g.logf/g.errorf, which are methods, not fields.
// g.logf/g.errorf, which are methods, not fields. modelsJSON (perf
// finding 3) follows this same rule: m.log("%s", fmt.Sprintf(...)),
// never a second variadic argument.
log func(string, ...any)
// aliases is the validated alias->target map from Config.ModelAliases
// (spec §5, v0.2), built once by newModelRegistry via
@@ -232,11 +238,35 @@ type modelRegistry struct {
// lookup (always a miss) that adds no branching cost to the
// no-overrides case — the same nil-vs-empty convention aliases above
// already establishes.
modelMeta map[string]*ModelMetaConfig
nowFn func() time.Time
warned map[string]bool
modelMeta map[string]*ModelMetaConfig
nowFn func() time.Time
warned map[string]bool
// modelsCache holds one encoded GET /v1/models response body per
// group, tagged with the modelsGen it was built from (perf finding 3,
// 2026-08-2x audit: 15.9ms/21.3MB/552,884 allocs per uncached call at
// the live 1,033-model catalog) — see modelsJSON's own doc comment.
// Guarded by modelsCacheMu.
modelsCache map[*group]modelsCacheEntry
providerNames []string
warnedMu sync.Mutex
// modelsGen counts how many times finishRefresh (below) has recorded a
// discovery attempt for ANY provider — bumped unconditionally, success
// or failure, since a failed listModels fetch can still pair with a
// successful metadata capture (captureModelMetadata's own doc comment)
// that changes what listFor/modelsJSON would produce. Read and written
// only via sync/atomic — modelsJSON (below) reads it outside any lock,
// finishRefresh writes it from whichever goroutine (warmFill's caller,
// or maybeRefresh's background goroutine) is recording that refresh.
modelsGen int64
warnedMu sync.Mutex
// modelsCacheMu guards modelsCache.
modelsCacheMu sync.Mutex
}
// modelsCacheEntry is one cached, encoded GET /v1/models response body
// (modelRegistry.modelsCache) plus the modelsGen it was built from.
type modelsCacheEntry struct {
body []byte
gen int64
}
// newModelRegistry builds a modelRegistry from adapters and cfg's matching
@@ -253,6 +283,7 @@ func newModelRegistry(adapters map[string]providerAdapter, cfg *Config, log func
nowFn: time.Now,
warned: make(map[string]bool),
providerNames: make([]string, 0, len(adapters)),
modelsCache: make(map[*group]modelsCacheEntry),
}
for name := range adapters {
m.providerNames = append(m.providerNames, name)
@@ -386,6 +417,21 @@ func (m *modelRegistry) now() time.Time {
return m.nowFn()
}
// finishRefresh records the outcome of a refresh attempt (st.finishRefresh)
// and invalidates modelsJSON's cache by bumping modelsGen — the ONLY two
// places that ever mutate st.discovered or st.discoveredContext both
// funnel through here (warmFill, refreshProvider), so every change
// listFor's output could possibly reflect is covered. Bumped
// unconditionally, not just on a successful refresh: captureModelMetadata
// (both call sites' own doc comments) always runs immediately alongside
// this same refresh cycle regardless of whether listModels itself
// succeeded, and a metadata-only change must invalidate the cache just as
// much as a discovered-set change would.
func (m *modelRegistry) finishRefresh(st *providerState, now time.Time, ids []string, err error) {
st.finishRefresh(now, ids, err)
atomic.AddInt64(&m.modelsGen, 1)
}
// warmFill performs newGateway's synchronous first discovery fill: for
// every discovery-enabled provider, it fetches listModels once, bounded by
// warmFillTimeout, and records the result via finishRefresh. A fetch error
@@ -402,7 +448,7 @@ func (m *modelRegistry) warmFill(ctx context.Context) {
fctx, cancel := context.WithTimeout(ctx, warmFillTimeout)
ids, err := m.adapters[name].listModels(fctx)
cancel()
st.finishRefresh(m.now(), ids, err)
m.finishRefresh(st, m.now(), ids, err)
if err != nil {
m.log("%s", fmt.Sprintf("model registry: initial discovery for provider %q failed: %v", name, err))
}
@@ -541,7 +587,7 @@ func (m *modelRegistry) refreshProvider(name string, st *providerState, adapter
if rec := recover(); rec != nil {
err = fmt.Errorf("panic: %v", rec)
}
st.finishRefresh(m.now(), ids, err)
m.finishRefresh(st, m.now(), ids, err)
if err != nil {
m.log("%s", fmt.Sprintf("model registry: discovery refresh for provider %q failed: %v", name, err))
}
@@ -807,6 +853,62 @@ func (m *modelRegistry) warnCollisionOnce(id, winner string, provs []string) {
m.log("%s", fmt.Sprintf("model registry: model id %q is provided by multiple providers %v; %q wins the bare id", id, provs, winner))
}
// modelsJSON returns grp's encoded GET /v1/models response body
// ({"object":"list","data":[...]} — routes_unified.go's handleModels
// writes this straight to the response), cached per (group, modelsGen)
// (perf finding 3, 2026-08-2x audit: 15.9ms/21.3MB/552,884 allocs per
// call, uncached, at the live 1,033-model catalog — registry.go's listFor
// plus routes_unified.go's json.Encoder, called on every request an
// OpenAI SDK client's session-start model poll happens to land on).
// listFor's output only changes when a discovery refresh runs
// (finishRefresh, above, bumps modelsGen unconditionally on every
// attempt) — between refreshes, every call for the same group returns
// the identical cached bytes without re-walking every provider's model
// set or re-marshaling.
//
// Cached per *group pointer, not by name or any other derived key:
// authStore builds every group once, at construction (auth.go's
// newAuthStore), and never rebuilds or replaces that map afterward — a
// config reload constructs an entirely new Gateway (and so a new
// modelRegistry with its own empty modelsCache), so a *group pointer here
// can never alias a different group's catalog, either within one
// registry's lifetime or across a reload.
func (m *modelRegistry) modelsJSON(grp *group) []byte {
gen := atomic.LoadInt64(&m.modelsGen)
m.modelsCacheMu.Lock()
if entry, ok := m.modelsCache[grp]; ok && entry.gen == gen {
m.modelsCacheMu.Unlock()
return entry.body
}
m.modelsCacheMu.Unlock()
body, err := json.Marshal(map[string]any{
"object": "list",
"data": m.listFor(grp),
})
if err != nil {
// listFor's map[string]any values are all JSON-safe primitives
// (string, int, float64, bool, and nested maps/slices of the
// same) — Marshal cannot fail on them in practice. Not caching a
// failure keeps this equivalent to the pre-cache
// json.NewEncoder(w).Encode call site, which also had nothing
// useful to do on an encode failure beyond logging.
m.log("%s", fmt.Sprintf("model registry: encoding /v1/models response for group %q failed: %v", grp.name, err))
return nil
}
// json.Encoder.Encode (the pre-cache call site) always appends a
// trailing '\n' after the value; json.Marshal does not — appended
// here so the cached body is byte-identical to what the old
// uncached path wrote (review fix, Should-Fix 5).
body = append(body, '\n')
m.modelsCacheMu.Lock()
m.modelsCache[grp] = modelsCacheEntry{gen: gen, body: body}
m.modelsCacheMu.Unlock()
return body
}
// listFor returns grp's visible model catalog as OpenAI-compatible model
// objects ({"id","object":"model","owned_by"}), sorted by id. A bare id
// owned by only one provider is listed once, bare. A bare id owned by more
+112
View File
@@ -1,6 +1,7 @@
package traefikllmgateway
import (
"bytes"
"context"
"encoding/json"
"errors"
@@ -874,6 +875,117 @@ func TestModelRegistry_ListFor_Collision_LosingProviderStillReachableViaPrefixed
}
}
// --- modelsJSON: cached encoded /v1/models body (perf finding 3) ---
// TestModelRegistry_ModelsJSON_CachesUntilFinishRefresh proves the perf
// fix: a group's encoded body is the SAME []byte (not merely
// byte-equal — literally not re-marshaled) across repeated calls, until
// finishRefresh actually records a refresh attempt for some provider
// (modelsGen bumps), at which point the next call reflects the change.
func TestModelRegistry_ModelsJSON_CachesUntilFinishRefresh(t *testing.T) {
t.Parallel()
adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")}
cfg := &Config{Providers: map[string]*ProviderConfig{"openai": {Models: []string{"gpt-test"}}}}
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
if err != nil {
t.Fatalf("newModelRegistry: %v", err)
}
grp := &group{name: "all"}
first := reg.modelsJSON(grp)
if len(first) == 0 {
t.Fatal("modelsJSON returned an empty body")
}
second := reg.modelsJSON(grp)
if len(first) == 0 || len(second) == 0 || &first[0] != &second[0] {
t.Error("two calls between refreshes returned different backing arrays, want the identical cached []byte")
}
if !strings.Contains(string(first), "gpt-test") {
t.Errorf("body = %s, want it to contain %q", first, "gpt-test")
}
// A discovery refresh — even for a provider whose own discovery is
// disabled here, driven directly via the registry's own finishRefresh
// wrapper rather than a real background goroutine — must invalidate
// the cache.
reg.finishRefresh(reg.states["openai"], reg.now(), []string{"gpt-test", "gpt-new"}, nil)
third := reg.modelsJSON(grp)
if string(third) == string(second) {
t.Error("modelsJSON body unchanged after finishRefresh added a new model, want it to reflect the new catalog")
}
if !strings.Contains(string(third), "gpt-new") {
t.Errorf("body after refresh = %s, want it to contain the newly discovered %q", third, "gpt-new")
}
}
// TestModelRegistry_ModelsJSON_PerGroupIsolation proves the cache never
// leaks one group's catalog to another: two groups authorized for
// disjoint providers must each get their own encoded body, keyed by
// *group pointer identity (modelsJSON's own doc comment).
func TestModelRegistry_ModelsJSON_PerGroupIsolation(t *testing.T) {
t.Parallel()
adapters := map[string]providerAdapter{
"openai": newFakeAdapter("openai"),
"anthropic": newFakeAdapter("anthropic"),
}
cfg := &Config{Providers: map[string]*ProviderConfig{
"openai": {Models: []string{"gpt-test"}},
"anthropic": {Models: []string{"claude-x"}},
}}
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
if err != nil {
t.Fatalf("newModelRegistry: %v", err)
}
grpOpenAI := &group{name: "openai-only", providers: []string{"openai"}}
grpAnthropic := &group{name: "anthropic-only", providers: []string{"anthropic"}}
bodyOpenAI := reg.modelsJSON(grpOpenAI)
bodyAnthropic := reg.modelsJSON(grpAnthropic)
if strings.Contains(string(bodyOpenAI), "claude-x") {
t.Errorf("openai-only group's body leaked anthropic's model: %s", bodyOpenAI)
}
if strings.Contains(string(bodyAnthropic), "gpt-test") {
t.Errorf("anthropic-only group's body leaked openai's model: %s", bodyAnthropic)
}
if !strings.Contains(string(bodyOpenAI), "gpt-test") {
t.Errorf("openai-only group's body = %s, want it to contain gpt-test", bodyOpenAI)
}
if !strings.Contains(string(bodyAnthropic), "claude-x") {
t.Errorf("anthropic-only group's body = %s, want it to contain claude-x", bodyAnthropic)
}
}
// TestModelRegistry_ModelsJSON_MatchesUncachedListFor proves modelsJSON's
// output is byte-identical to the OLD, pre-cache call site
// (json.NewEncoder(w).Encode(...), routes_unified.go before perf finding
// 3) for the same envelope over listFor's own result — the cache changes
// WHEN the body is computed, never WHAT it computes, including
// json.Encoder.Encode's own trailing '\n' that a bare json.Marshal does
// not add (review fix, Should-Fix 5 — modelsJSON appends it explicitly).
func TestModelRegistry_ModelsJSON_MatchesUncachedListFor(t *testing.T) {
t.Parallel()
adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")}
cfg := &Config{Providers: map[string]*ProviderConfig{"openai": {Models: []string{"zeta", "alpha"}}}}
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
if err != nil {
t.Fatalf("newModelRegistry: %v", err)
}
grp := allowAllGroup()
got := reg.modelsJSON(grp)
var want bytes.Buffer
if err := json.NewEncoder(&want).Encode(map[string]any{"object": "list", "data": reg.listFor(grp)}); err != nil {
t.Fatalf("json.NewEncoder.Encode: %v", err)
}
if string(got) != want.String() {
t.Errorf("modelsJSON = %q, want %q (byte-identical to the old json.Encoder.Encode call site)", got, want.String())
}
}
// --- end-to-end: GET /v1/models through Gateway.ServeHTTP ---
func TestGateway_ServeHTTP_GetV1Models_ReturnsGroupFilteredList(t *testing.T) {
+6 -5
View File
@@ -39,7 +39,11 @@ const gatewayAliasKey = "__alias"
// handleModels implements GET /v1/models: an authenticated caller gets back
// an OpenAI-compatible {"object":"list","data":[...]} envelope of the
// models their group can see, via modelRegistry.listFor.
// models their group can see, via modelRegistry.modelsJSON (perf finding
// 3: modelsJSON caches the encoded body per (group, discovery generation)
// so a request between discovery refreshes never re-walks listFor/
// re-marshals — registry.go owns the cache; this call site just writes
// its result).
func (g *Gateway) handleModels(w http.ResponseWriter, r *http.Request) {
u, grp, ok := g.auth.identify(r)
g.logAuthEvent(ok, authEventUserName(u), r)
@@ -49,10 +53,7 @@ func (g *Gateway) handleModels(w http.ResponseWriter, r *http.Request) {
}
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(map[string]any{ // headers already committed; nothing useful to do on encode failure
"object": "list",
"data": g.registry.listFor(grp),
})
_, _ = w.Write(g.registry.modelsJSON(grp)) // headers already committed; nothing useful to do on write failure
}
// adapterCall abstracts providerAdapter.chatCompletion and .embeddings: