fix(proxy): mutex on transports, ctx-aware warmup delay, Retry-After on 503
CI / CI (pull_request) Successful in 2m11s

- transportsMu (sync.RWMutex): guard getOrCreateTransport reads/writes;
  double-checked locking pattern; safe for dynamic backends
- applyWarmup: replace time.Sleep with select{time.After, ctx.Done}
  so cancelled warmup requests release the goroutine immediately
- Retry-After: 5 header on queue-full 503 (client backoff hint)
- OTel trace_id skipped: collector in tracing ns, not monitoring ns
This commit is contained in:
Admin Bot
2026-09-27 19:13:38 +08:00
parent 6257b8c201
commit 78a5d6baa2
+28 -7
View File
@@ -15,6 +15,7 @@ import (
"sort" "sort"
"strconv" "strconv"
"strings" "strings"
"sync"
"syscall" "syscall"
"time" "time"
@@ -46,6 +47,8 @@ type Handler struct {
schedulers map[string]*Scheduler schedulers map[string]*Scheduler
// CAG: delay applied to warmup requests before queue entry (#45) // CAG: delay applied to warmup requests before queue entry (#45)
warmupDelay time.Duration warmupDelay time.Duration
// transportsMu guards h.transports for goroutine-safe reads/writes.
transportsMu sync.RWMutex
} }
// Route represents a reverse proxy route. // Route represents a reverse proxy route.
@@ -167,7 +170,16 @@ func New(cfg *config.Config) *Handler {
// getOrCreateTransport returns a shared http.Transport for the given upstream address. // getOrCreateTransport returns a shared http.Transport for the given upstream address.
// This ensures connections are pooled and reused across requests to the same upstream. // This ensures connections are pooled and reused across requests to the same upstream.
func (h *Handler) getOrCreateTransport(addr string, up *config.Upstream) *http.Transport { func (h *Handler) getOrCreateTransport(addr string, up *config.Upstream) *http.Transport {
h.transportsMu.RLock()
if t, ok := h.transports[addr]; ok { if t, ok := h.transports[addr]; ok {
h.transportsMu.RUnlock()
return t
}
h.transportsMu.RUnlock()
h.transportsMu.Lock()
defer h.transportsMu.Unlock()
if t, ok := h.transports[addr]; ok { // double-check after write lock
return t return t
} }
@@ -216,7 +228,7 @@ func (h *Handler) getOrCreateTransport(addr string, up *config.Upstream) *http.T
} }
h.transports[addr] = transport h.transports[addr] = transport
return transport return transport // written under transportsMu write lock
} }
// directorFunc modifies the request to be sent to the upstream. // directorFunc modifies the request to be sent to the upstream.
@@ -351,16 +363,22 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
// applyWarmup strips the X-Poimen-Warmup header, delays the request, and records // applyWarmup strips the X-Poimen-Warmup header, delays the request while
// the metric. It is a no-op when the header is absent or not "true". // honouring context cancellation, and records the metric.
func (h *Handler) applyWarmup(r *http.Request) { // Returns false if the context was cancelled during the delay (caller must abort).
func (h *Handler) applyWarmup(ctx context.Context, r *http.Request) bool {
if r.Header.Get("X-Poimen-Warmup") != "true" { if r.Header.Get("X-Poimen-Warmup") != "true" {
return return true
} }
r.Header.Del("X-Poimen-Warmup") r.Header.Del("X-Poimen-Warmup")
r.Header.Set("X-Priority", "low") r.Header.Set("X-Priority", "low")
time.Sleep(h.warmupDelay) select {
case <-time.After(h.warmupDelay):
case <-ctx.Done():
return false
}
observability.Global().IncrWarmupDelayed() observability.Global().IncrWarmupDelayed()
return true
} }
// acquireScheduledBackend selects a backend from sched, applies warmup // acquireScheduledBackend selects a backend from sched, applies warmup
@@ -370,7 +388,9 @@ func (h *Handler) applyWarmup(r *http.Request) {
// If the slot cannot be acquired it writes a 503 and returns (nil, nil, false). // 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) { 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) backend := sched.Select(cacheKey)
h.applyWarmup(r) if !h.applyWarmup(ctx, r) {
return nil, nil, false
}
if backend == nil { if backend == nil {
return nil, route, true return nil, route, true
} }
@@ -382,6 +402,7 @@ func (h *Handler) acquireScheduledBackend(ctx context.Context, w http.ResponseWr
"path": r.URL.Path, "path": r.URL.Path,
"model": route.ModelName, "model": route.ModelName,
}) })
w.Header().Set("Retry-After", "5")
w.WriteHeader(http.StatusServiceUnavailable) w.WriteHeader(http.StatusServiceUnavailable)
fmt.Fprintf(w, "queue full") fmt.Fprintf(w, "queue full")
} else { } else {