From d703ace2c6250a2c4db6369078e879d805d15273 Mon Sep 17 00:00:00 2001 From: Admin Bot Date: Sun, 27 Sep 2026 18:57:04 +0800 Subject: [PATCH] fix(cag): extract applyWarmup/acquireScheduledBackend helpers, add race-free tests - proxy.go: extract applyWarmup + acquireScheduledBackend to reduce ServeHTTP complexity - proxy_test.go: TestProxyQueueFull, TestProxyCacheHitLogging (AC5/AC6/AC7 HTTP-level coverage) - scheduler_test.go: TestBackendAcquireCancelledContext + runtime.Gosched spin-wait (no time.Sleep) - router_test.go: TestAuthenticate_MethodCapabilityOverrideDirect with real RS256 JWT + JWKS server --- internal/proxy/proxy.go | 114 +++++++++------- internal/proxy/proxy_test.go | 136 +++++++++++++++++++ internal/proxy/scheduler_test.go | 13 +- internal/serviceadapter/router_test.go | 173 +++++++++++++++++++------ 4 files changed, 343 insertions(+), 93 deletions(-) diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 37f8411..2f268e7 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -2,6 +2,7 @@ package proxy import ( + "context" "encoding/json" "errors" "fmt" @@ -468,56 +469,13 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { var selectedBackend *Backend cacheKey := r.Header.Get("X-Poimen-Cache-Key") if sched := h.schedulers[route.ModelName]; sched != nil { - selectedBackend = sched.Select(cacheKey) - - // Warmup: sleep BEFORE queue entry so agents arriving in that window win slots first. - if r.Header.Get("X-Poimen-Warmup") == "true" { - r.Header.Del("X-Poimen-Warmup") - r.Header.Set("X-Priority", "low") - time.Sleep(h.warmupDelay) - observability.Global().IncrWarmupDelayed() + var ok bool + selectedBackend, route, ok = h.acquireScheduledBackend(r.Context(), w, r, sched, route, cacheKey) + if !ok { + return } - if selectedBackend != nil { - if err := selectedBackend.Acquire(r.Context()); err != nil { - if errors.Is(err, ErrQueueFull) { - observability.Global().IncrQueueFull() - logging.Warnf("backend queue full", map[string]string{ - "backend": selectedBackend.Address, - "path": r.URL.Path, - "model": route.ModelName, - }) - w.WriteHeader(http.StatusServiceUnavailable) - fmt.Fprintf(w, "queue full") - } else { - logging.Warnf("backend acquire failed", map[string]string{ - "backend": selectedBackend.Address, - "path": r.URL.Path, - "model": route.ModelName, - "err": err.Error(), - }) - w.WriteHeader(http.StatusServiceUnavailable) - fmt.Fprintf(w, "upstream busy") - } - return - } defer selectedBackend.Release() - - // Patch route to selected backend if it differs from the config default. - if selectedBackend.Address != route.Upstream.Address { - patchedCfg := *route.Upstream - patchedCfg.Address = selectedBackend.Address - patchedURL, _ := url.Parse("http://" + selectedBackend.Address) - route = &Route{ - Name: route.Name, - ModelName: route.ModelName, - Upstream: &patchedCfg, - Transport: h.getOrCreateTransport(selectedBackend.Address, &patchedCfg), - Director: func(req *http.Request) { - directorFunc(req, patchedURL, &patchedCfg) - }, - } - } } } @@ -586,6 +544,68 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { +// applyWarmup strips the X-Poimen-Warmup header, delays the request, and records +// the metric. It is a no-op when the header is absent or not "true". +func (h *Handler) applyWarmup(r *http.Request) { + if r.Header.Get("X-Poimen-Warmup") != "true" { + return + } + r.Header.Del("X-Poimen-Warmup") + r.Header.Set("X-Priority", "low") + time.Sleep(h.warmupDelay) + observability.Global().IncrWarmupDelayed() +} + +// acquireScheduledBackend selects a backend from sched, applies warmup +// deprioritisation, acquires a concurrency slot, and patches route when the +// selected backend differs from the config default. It returns the backend +// to Release (caller must defer), the (possibly patched) route, and ok=true. +// If the slot cannot be acquired it writes a 503 and returns (nil, nil, false). +func (h *Handler) acquireScheduledBackend(ctx context.Context, w http.ResponseWriter, r *http.Request, sched *Scheduler, route *Route, cacheKey string) (*Backend, *Route, bool) { + backend := sched.Select(cacheKey) + h.applyWarmup(r) + if backend == nil { + return nil, route, true + } + if err := backend.Acquire(ctx); err != nil { + if errors.Is(err, ErrQueueFull) { + observability.Global().IncrQueueFull() + logging.Warnf("backend queue full", map[string]string{ + "backend": backend.Address, + "path": r.URL.Path, + "model": route.ModelName, + }) + w.WriteHeader(http.StatusServiceUnavailable) + fmt.Fprintf(w, "queue full") + } else { + logging.Warnf("backend acquire failed", map[string]string{ + "backend": backend.Address, + "path": r.URL.Path, + "model": route.ModelName, + "err": err.Error(), + }) + w.WriteHeader(http.StatusServiceUnavailable) + fmt.Fprintf(w, "upstream busy") + } + return nil, nil, false + } + if backend.Address != route.Upstream.Address { + patchedCfg := *route.Upstream + patchedCfg.Address = backend.Address + patchedURL, _ := url.Parse("http://" + backend.Address) + route = &Route{ + Name: route.Name, + ModelName: route.ModelName, + Upstream: &patchedCfg, + Transport: h.getOrCreateTransport(backend.Address, &patchedCfg), + Director: func(req *http.Request) { + directorFunc(req, patchedURL, &patchedCfg) + }, + } + } + return backend, route, true +} + // getValidModels returns a sorted list of all configured model names. func (h *Handler) getValidModels() []string { var models []string diff --git a/internal/proxy/proxy_test.go b/internal/proxy/proxy_test.go index e012fac..3942dfe 100644 --- a/internal/proxy/proxy_test.go +++ b/internal/proxy/proxy_test.go @@ -5,11 +5,13 @@ import ( "io" "net/http" "net/http/httptest" + "runtime" "strings" "testing" "time" "forgejo.riotpiao.com/rock/homelab-frontend/internal/config" + "forgejo.riotpiao.com/rock/homelab-frontend/internal/observability" ) func TestProxyBasic(t *testing.T) { @@ -415,6 +417,140 @@ func TestProxyPreservesBody(t *testing.T) { } } +// TestProxyQueueFull verifies that a 503 is returned and the queue_full metric +// is incremented when the backend queue is at capacity. +// +// Setup: maxSlots=1, maxQueue=1. Request 1 holds the slot. Request 2 fills the +// queue (blocks waiting for slot). Request 3 gets ErrQueueFull → 503. +func TestProxyQueueFull(t *testing.T) { + t.Setenv("VLLM_MAX_NUM_SEQS", "1") // 1 concurrent slot + t.Setenv("VLLM_MAX_QUEUE_DEPTH", "1") // queue capacity 1 + t.Setenv("WARMUP_DELAY_MS", "0") + observability.Global().Reset() + + // Upstream blocks until released; buffers one signal when it first starts. + slotHeld := make(chan struct{}, 1) + release := make(chan struct{}) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + select { + case slotHeld <- struct{}{}: + default: + } + <-release + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + upstreamAddr := strings.TrimPrefix(upstream.URL, "http://") + cfg := &config.Config{ + Routes: make(map[string]*config.Route), + Models: map[string]*config.ModelUpstream{ + "q-model": { + Name: "q-model", + Address: upstreamAddr, + }, + }, + } + h := New(cfg) + defer h.Close() + + srv := httptest.NewServer(h) + defer srv.Close() + + mkReq := func() *http.Request { + body := strings.NewReader(`{"model":"q-model","messages":[]}`) + req, _ := http.NewRequest("POST", srv.URL+"/v1/chat/completions", body) + req.Header.Set("Content-Type", "application/json") + return req + } + + // Request 1: acquires the only slot, reaches upstream, blocks there. + r1ch := make(chan *http.Response, 1) + go func() { + resp, _ := http.DefaultClient.Do(mkReq()) + r1ch <- resp + }() + <-slotHeld // slot is acquired; upstream is holding request 1 + + // Request 2: enters the queue (capacity 1), blocks waiting for slot. + r2ch := make(chan *http.Response, 1) + go func() { + resp, _ := http.DefaultClient.Do(mkReq()) + r2ch <- resp + }() + // Spin until request 2 enters the backend queue channel. + for len(h.schedulers["q-model"].backends[0].queue) == 0 { + runtime.Gosched() + } + + // Request 3: queue is full → expect 503 and metric increment. + resp3, err := http.DefaultClient.Do(mkReq()) + if err != nil { + t.Fatalf("request 3 failed: %v", err) + } + defer resp3.Body.Close() + + if resp3.StatusCode != http.StatusServiceUnavailable { + t.Errorf("expected 503 queue full, got %d", resp3.StatusCode) + } + if got := observability.Global().GetQueueFull(); got != 1 { + t.Errorf("expected GetQueueFull()=1, got %d", got) + } + + // Cleanup: unblock upstream so goroutines can finish. + close(release) + if r := <-r1ch; r != nil { + r.Body.Close() + } + if r := <-r2ch; r != nil { + r.Body.Close() + } +} + +// TestProxyCacheHitLogging verifies that when a response contains the +// vllm-num-cached-tokens header and a cache key was set, the prefix cache +// hit token count is incremented in observability metrics. +func TestProxyCacheHitLogging(t *testing.T) { + observability.Global().Reset() + + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("vllm-num-cached-tokens", "128") + w.WriteHeader(http.StatusOK) + })) + defer upstream.Close() + + upstreamAddr := strings.TrimPrefix(upstream.URL, "http://") + cfg := &config.Config{ + Routes: make(map[string]*config.Route), + Models: map[string]*config.ModelUpstream{ + "cache-model": { + Name: "cache-model", + Address: upstreamAddr, + }, + }, + } + h := New(cfg) + defer h.Close() + + srv := httptest.NewServer(h) + defer srv.Close() + + body := strings.NewReader(`{"model":"cache-model","messages":[]}`) + req, _ := http.NewRequest("POST", srv.URL+"/v1/chat/completions", body) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("X-Poimen-Cache-Key", "testkey123") + + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("request failed: %v", err) + } + resp.Body.Close() + + if got := observability.Global().GetPrefixCacheHitTokens(); got != 128 { + t.Errorf("expected PrefixCacheHitTokens=128, got %d", got) + } +} + // TestWarmupHeaderHandling verifies that X-Poimen-Warmup: true is stripped from // the upstream request and replaced with X-Priority: low, while a request without // X-Poimen-Warmup reaches upstream without either header. diff --git a/internal/proxy/scheduler_test.go b/internal/proxy/scheduler_test.go index 049adb1..e0631b3 100644 --- a/internal/proxy/scheduler_test.go +++ b/internal/proxy/scheduler_test.go @@ -3,6 +3,7 @@ package proxy import ( "context" "errors" + "runtime" "sync" "testing" "time" @@ -23,8 +24,10 @@ func TestBackendAcquireBlocksThenUnblocksOnRelease(t *testing.T) { g2Done <- b.Acquire(context.Background()) }() - // Give G2 time to block. - time.Sleep(20 * time.Millisecond) + // Wait until G2 enters the queue channel (proves Acquire is blocking on slots). + for len(b.queue) == 0 { + runtime.Gosched() + } select { case err := <-g2Done: @@ -119,8 +122,10 @@ func TestBackendAcquireCancelledContext(t *testing.T) { g2Done <- b.Acquire(ctx) }() - // Give G2 time to enter the queue and block waiting for a slot. - time.Sleep(20 * time.Millisecond) + // Wait until G2 enters the queue channel. + for len(b.queue) == 0 { + runtime.Gosched() + } // Verify G2 is queued (queue channel has one item). if qd := len(b.queue); qd != 1 { diff --git a/internal/serviceadapter/router_test.go b/internal/serviceadapter/router_test.go index d608127..29f3508 100644 --- a/internal/serviceadapter/router_test.go +++ b/internal/serviceadapter/router_test.go @@ -1,10 +1,17 @@ package serviceadapter import ( + "crypto/rand" + "crypto/rsa" + "encoding/base64" "encoding/json" + "math/big" "net/http" "net/http/httptest" "testing" + "time" + + "github.com/golang-jwt/jwt/v5" "forgejo.riotpiao.com/rock/homelab-frontend/internal/auth" "forgejo.riotpiao.com/rock/homelab-frontend/internal/identity" @@ -321,58 +328,140 @@ func TestDispatch_MethodCapabilityOverridesPOSTDefault(t *testing.T) { } } -// TestAuthenticate_MethodCapabilityOverrideDirect calls authenticate() directly -// (same-package access) with a non-nil jwtValidator so the auth gate is not bypassed. -// Two sub-tests verify: -// 1. authenticate() is actually invoked when jwtValidator != nil (returns 401 for -// a missing Authorization header — proves the function ran, not the nil-bypass path). -// 2. The method-level Auth.Capability override selects "memory:read" over the -// POST verb-derived default "memory:write" — exercises the same lines that -// authenticate() uses internally (lines: required = method.Auth.Capability). +// TestAuthenticate_MethodCapabilityOverrideDirect calls d.Dispatch() with a real +// signed RS256 JWT so authenticate() executes end-to-end and the capability-override +// branch (router.go lines 124-130) is reached and covered. +// +// Sub-test 1: token has "memory:read", method.Auth.Capability="memory:read" → 200. +// Sub-test 2: token has "memory:read", no method override → POST derives +// +// "memory:write" → CheckPermissions fails → 403. func TestAuthenticate_MethodCapabilityOverrideDirect(t *testing.T) { - // Validator with a dummy JWKS URL (never fetched — ValidateBearerToken is not - // reached in these sub-tests because we hit the auth-header guard first). - validator := auth.NewValidator("", "", "http://127.0.0.1:0/jwks") - - // Method with method-level capability override: POST normally derives "memory:write" - // but this endpoint is a read-only standing-context retrieval → override to "memory:read". - method := &Method{ - Verb: "POST", - UpstreamPath: "/memory/standing-context", - Auth: &Auth{Capability: "memory:read"}, + // Generate an RSA key pair for signing test JWTs. + privateKey, err := rsa.GenerateKey(rand.Reader, 2048) + if err != nil { + t.Fatalf("rsa.GenerateKey: %v", err) } - d := &Dispatcher{jwtValidator: validator} + // Build minimal JWKS document from the RSA public key. + nBytes := privateKey.PublicKey.N.Bytes() + eBig := new(big.Int).SetInt64(int64(privateKey.PublicKey.E)) + eBytes := eBig.Bytes() + jwksDoc, _ := json.Marshal(map[string]interface{}{ + "keys": []map[string]interface{}{ + { + "kty": "RSA", + "use": "sig", + "kid": "test-key", + "alg": "RS256", + "n": base64.RawURLEncoding.EncodeToString(nBytes), + "e": base64.RawURLEncoding.EncodeToString(eBytes), + }, + }, + }) - // Sub-test 1: authenticate() is called and returns 401 for a missing - // Authorization header — proves the function runs (not the nil-validator bypass). - t.Run("authenticate_called_returns_401_without_token", func(t *testing.T) { + // Serve JWKS from a local test server. + jwksSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write(jwksDoc) + })) + defer jwksSrv.Close() + + // Upstream that returns 200 when authenticate() passes. + upstreamSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer upstreamSrv.Close() + + // signJWT returns a signed RS256 token with the given permissions slice. + signJWT := func(permissions []string) string { + tok := jwt.NewWithClaims(jwt.SigningMethodRS256, jwt.MapClaims{ + "iss": "https://authentik.riotpiao.com/application/o/test/", + "aud": "test-audience", + "sub": "test-user", + "exp": time.Now().Add(time.Hour).Unix(), + "iat": time.Now().Unix(), + "permissions": permissions, + }) + tok.Header["kid"] = "test-key" + signed, err := tok.SignedString(privateKey) + if err != nil { + t.Fatalf("sign JWT: %v", err) + } + return signed + } + + validator := auth.NewValidator("", "", jwksSrv.URL) + + // Sub-test 1: method.Auth.Capability="memory:read"; token has "memory:read". + // authenticate() must execute the override branch and pass → upstream returns 200. + t.Run("capability_override_match_reaches_upstream", func(t *testing.T) { + adapter := ServiceAdapter{ + Name: "memory", + ServiceName: "memory", + Spec: Spec{ + ServiceName: "memory", + Upstream: Upstream{URL: upstreamSrv.URL, TimeoutSeconds: 5}, + Auth: Auth{Required: true}, + Resources: []Resource{ + { + Name: "standing-context", + Methods: []Method{ + { + Verb: "POST", + UpstreamPath: "/memory/standing-context", + Auth: &Auth{Capability: "memory:read"}, // override + }, + }, + }, + }, + }, + } + d := NewDispatcher(newTestRegistry(adapter), validator) w := httptest.NewRecorder() r := httptest.NewRequest("POST", "/", nil) - got := d.authenticate(w, r, "memory", method) - if got { - t.Error("expected authenticate() to return false for missing Authorization header") - } - if w.Code != http.StatusUnauthorized { - t.Errorf("expected 401 Unauthorized, got %d", w.Code) + r.Header.Set("X-Service", "memory") + r.Header.Set("X-Resource", "standing-context") + r.Header.Set("Authorization", "Bearer "+signJWT([]string{"memory:read"})) + d.Dispatch(w, r) + if w.Code != http.StatusOK { + t.Errorf("expected 200 (override memory:read match), got %d body=%s", w.Code, w.Body.String()) } }) - // Sub-test 2: prove that the capability override selects "memory:read" (not the - // verb-derived "memory:write") using the identical logic from authenticate(). - t.Run("method_level_capability_overrides_verb_derived_default", func(t *testing.T) { - // Replicate the exact lines from authenticate() (router.go:119-121): - required := capabilityForVerb("memory", method.Verb) // POST → "memory:write" - if method.Auth != nil && method.Auth.Capability != "" { - required = method.Auth.Capability // override → "memory:read" + // Sub-test 2: no method-level Auth override; POST → capabilityForVerb → "memory:write". + // Token has only "memory:read" → CheckPermissions fails → 403. + t.Run("verb_derived_capability_mismatch_returns_403", func(t *testing.T) { + adapter := ServiceAdapter{ + Name: "memory", + ServiceName: "memory", + Spec: Spec{ + ServiceName: "memory", + Upstream: Upstream{URL: upstreamSrv.URL, TimeoutSeconds: 5}, + Auth: Auth{Required: true}, + Resources: []Resource{ + { + Name: "data", + Methods: []Method{ + { + Verb: "POST", + UpstreamPath: "/memory/data", + // No method-level Auth: POST derives "memory:write". + }, + }, + }, + }, + }, } - if required != "memory:read" { - t.Errorf("method-level Auth.Capability override: got %q, want \"memory:read\"", required) - } - // Confirm the verb-derived value without override would have been wrong. - verbDefault := capabilityForVerb("memory", method.Verb) - if verbDefault != "memory:write" { - t.Errorf("verb-derived default for POST: got %q, want \"memory:write\"", verbDefault) + d := NewDispatcher(newTestRegistry(adapter), validator) + w := httptest.NewRecorder() + r := httptest.NewRequest("POST", "/", nil) + r.Header.Set("X-Service", "memory") + r.Header.Set("X-Resource", "data") + r.Header.Set("Authorization", "Bearer "+signJWT([]string{"memory:read"})) + d.Dispatch(w, r) + if w.Code != http.StatusForbidden { + t.Errorf("expected 403 (write required, read only in token), got %d body=%s", w.Code, w.Body.String()) } }) }