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
2026-09-27 18:57:04 +08:00
parent a731185588
commit d703ace2c6
4 changed files with 343 additions and 93 deletions
@@ -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