fix(cag): resolve data races, add queue observability, fix test coverage

Cycle 2 fixes from reviewer-sre + reviewer-quality:

- scheduler.go: remove delete() under RLock (data race); expired entries
  skipped on read, cleaned lazily by next MarkFailed write
- proxy.go: pre-populate transports for scheduler backends at startup to
  prevent concurrent map writes from ServeHTTP goroutines
- proxy.go: log + IncrQueueFull() on queue-full 503 path (was silent)
- proxy.go: remove stub total_tokens/hit_rate from cache-hit log event
- serviceadapter/router.go: Warnf when Auth.Required but jwtValidator nil
- observability: add queue_full_requests_total counter + Prometheus export
- scheduler_test.go: TestBackendAcquireCancelledContext (ctx cancel branch)
- proxy_test.go: implement TestWarmupHeaderHandling (remove t.Skip)
- router_test.go: TestAuthenticate_MethodCapabilityOverrideDirect (auth coverage)
This commit is contained in:
Admin Bot
2026-09-27 18:39:45 +08:00
parent fc5c0e967d
commit a731185588
8 changed files with 184 additions and 27 deletions
+23 -2
View File
@@ -131,6 +131,17 @@ func New(cfg *config.Config) *Handler {
h.schedulers[modelName] = NewScheduler(addrs, maxSlots, maxQueue)
}
}
// Pre-populate transports for all scheduler backends at startup (single-threaded).
// getOrCreateTransport is not goroutine-safe; pre-populating avoids concurrent writes.
for _, model := range cfg.Models {
for _, addr := range model.AllAddresses() {
h.getOrCreateTransport(addr, &config.Upstream{
ConnectTimeout: h.defaultConnectTimeout,
ReadTimeout: h.defaultReadTimeout,
WriteTimeout: h.defaultWriteTimeout,
})
}
}
for name, route := range cfg.Routes {
// Create a transport per unique upstream address for connection reuse
@@ -470,9 +481,21 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
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")
}
@@ -541,8 +564,6 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
"event": "llm_request",
"cache_key": cacheKey,
"cached_tokens": n,
"total_tokens": 0,
"hit_rate": 0.0,
})
return nil
}
+19 -19
View File
@@ -416,13 +416,14 @@ func TestProxyPreservesBody(t *testing.T) {
}
// TestWarmupHeaderHandling verifies that X-Poimen-Warmup: true is stripped from
// the upstream request and replaced with X-Priority: low, while no X-Poimen-Warmup
// on a normal request produces neither header downstream.
// the upstream request and replaced with X-Priority: low, while a request without
// X-Poimen-Warmup reaches upstream without either header.
//
// TODO(#45): implement warmup header stripping in proxy.Handler Director or middleware.
// Once implemented, remove t.Skip and wire up the real Handler.
// The warmup path runs only for model-based routes (cfg.Models + scheduler),
// so we POST to /v1/chat/completions with a JSON body that names the test model.
// WARMUP_DELAY_MS=0 eliminates the 500 ms sleep.
func TestWarmupHeaderHandling(t *testing.T) {
t.Skip("requires full handler setup: warmup header transform (#45) not yet implemented")
t.Setenv("WARMUP_DELAY_MS", "0")
tests := []struct {
name string
@@ -433,8 +434,8 @@ func TestWarmupHeaderHandling(t *testing.T) {
{
name: "warmup=true",
sendWarmup: true,
wantWarmupGone: true, // X-Poimen-Warmup must NOT reach upstream
wantPriorityLow: true, // X-Priority: low must be set
wantWarmupGone: true, // X-Poimen-Warmup must NOT reach upstream
wantPriorityLow: true, // X-Priority: low must be set
},
{
name: "warmup=false",
@@ -455,18 +456,14 @@ func TestWarmupHeaderHandling(t *testing.T) {
defer upstream.Close()
upstreamAddr := strings.TrimPrefix(upstream.URL, "http://")
// Warmup code only runs for model-based routes (scheduler path).
// Use cfg.Models so New() builds a Scheduler for "warmup-model".
cfg := &config.Config{
Routes: map[string]*config.Route{
"test": {
Name: "test",
Upstream: config.Upstream{
Address: upstreamAddr,
ConnectTimeout: 5 * time.Second,
ReadTimeout: 5 * time.Second,
WriteTimeout: 5 * time.Second,
MaxBodySize: 1024 * 1024,
AuthRequired: false,
},
Routes: make(map[string]*config.Route),
Models: map[string]*config.ModelUpstream{
"warmup-model": {
Name: "warmup-model",
Address: upstreamAddr,
},
},
}
@@ -476,7 +473,10 @@ func TestWarmupHeaderHandling(t *testing.T) {
srv := httptest.NewServer(h)
defer srv.Close()
req, _ := http.NewRequest("GET", srv.URL+"/test", nil)
// POST to /v1/chat/completions — triggers model-based dispatch and warmup path.
body := strings.NewReader(`{"model":"warmup-model","messages":[]}`)
req, _ := http.NewRequest("POST", srv.URL+"/v1/chat/completions", body)
req.Header.Set("Content-Type", "application/json")
if tc.sendWarmup {
req.Header.Set("X-Poimen-Warmup", "true")
}
+6 -4
View File
@@ -128,10 +128,12 @@ func (s *Scheduler) selectRendezvous(cacheKey string) *Backend {
return scores[i].score > scores[j].score
})
for _, sc := range scores {
if t, failed := s.failedUntil[sc.b.Address]; !failed || now.After(t) {
if now.After(t) {
delete(s.failedUntil, sc.b.Address)
}
_, failed := s.failedUntil[sc.b.Address]
if !failed {
return sc.b
}
expiry := s.failedUntil[sc.b.Address]
if now.After(expiry) {
return sc.b
}
}
+51
View File
@@ -100,6 +100,57 @@ func TestSchedulerSelectNoKeyRoundRobin(t *testing.T) {
}
}
// TestBackendAcquireCancelledContext verifies that Acquire returns ctx.Err() when the
// context is cancelled while the caller is blocked waiting for a concurrency slot,
// and that the queue depth returns to 0 after cancellation.
func TestBackendAcquireCancelledContext(t *testing.T) {
b := newBackend("host:1", 1, 1)
// G1: acquire and hold the only slot indefinitely.
if err := b.Acquire(context.Background()); err != nil {
t.Fatalf("G1 Acquire: unexpected error: %v", err)
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
g2Done := make(chan error, 1)
go func() {
g2Done <- b.Acquire(ctx)
}()
// Give G2 time to enter the queue and block waiting for a slot.
time.Sleep(20 * time.Millisecond)
// Verify G2 is queued (queue channel has one item).
if qd := len(b.queue); qd != 1 {
t.Errorf("expected queue depth 1 while G2 blocked, got %d", qd)
}
// Cancel ctx while G2 is waiting for the slot.
cancel()
select {
case err := <-g2Done:
if err == nil {
t.Fatal("G2 Acquire: expected ctx.Err(), got nil")
}
if !errors.Is(err, context.Canceled) {
t.Fatalf("G2 Acquire: expected context.Canceled, got %v", err)
}
case <-time.After(time.Second):
t.Fatal("G2 did not return within 1s after context cancel")
}
// Queue depth must return to 0 after G2 is drained.
if qd := len(b.queue); qd != 0 {
t.Errorf("expected queue depth 0 after cancellation, got %d", qd)
}
// Cleanup: release G1's slot.
b.Release()
}
// TestSchedulerMarkFailedExcludesBackend verifies that a failed backend is not
// selected within the fail window (30 s default).
func TestSchedulerMarkFailedExcludesBackend(t *testing.T) {