diff --git a/registry.go b/registry.go index 14e4328..5d4b5d6 100644 --- a/registry.go +++ b/registry.go @@ -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 diff --git a/registry_test.go b/registry_test.go index a80a194..febf747 100644 --- a/registry_test.go +++ b/registry_test.go @@ -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) { diff --git a/routes_unified.go b/routes_unified.go index d8523ed..09d9678 100644 --- a/routes_unified.go +++ b/routes_unified.go @@ -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: