fix(proxy): mutex on transports, ctx-aware warmup delay, Retry-After on 503
CI / CI (pull_request) Successful in 2m11s
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:
+28
-7
@@ -15,6 +15,7 @@ import (
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
@@ -46,6 +47,8 @@ type Handler struct {
|
||||
schedulers map[string]*Scheduler
|
||||
// CAG: delay applied to warmup requests before queue entry (#45)
|
||||
warmupDelay time.Duration
|
||||
// transportsMu guards h.transports for goroutine-safe reads/writes.
|
||||
transportsMu sync.RWMutex
|
||||
}
|
||||
|
||||
// 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.
|
||||
// This ensures connections are pooled and reused across requests to the same upstream.
|
||||
func (h *Handler) getOrCreateTransport(addr string, up *config.Upstream) *http.Transport {
|
||||
h.transportsMu.RLock()
|
||||
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
|
||||
}
|
||||
|
||||
@@ -216,7 +228,7 @@ func (h *Handler) getOrCreateTransport(addr string, up *config.Upstream) *http.T
|
||||
}
|
||||
|
||||
h.transports[addr] = transport
|
||||
return transport
|
||||
return transport // written under transportsMu write lock
|
||||
}
|
||||
|
||||
// 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
|
||||
// the metric. It is a no-op when the header is absent or not "true".
|
||||
func (h *Handler) applyWarmup(r *http.Request) {
|
||||
// applyWarmup strips the X-Poimen-Warmup header, delays the request while
|
||||
// honouring context cancellation, and records the metric.
|
||||
// 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" {
|
||||
return
|
||||
return true
|
||||
}
|
||||
r.Header.Del("X-Poimen-Warmup")
|
||||
r.Header.Set("X-Priority", "low")
|
||||
time.Sleep(h.warmupDelay)
|
||||
select {
|
||||
case <-time.After(h.warmupDelay):
|
||||
case <-ctx.Done():
|
||||
return false
|
||||
}
|
||||
observability.Global().IncrWarmupDelayed()
|
||||
return true
|
||||
}
|
||||
|
||||
// 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).
|
||||
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 !h.applyWarmup(ctx, r) {
|
||||
return nil, nil, false
|
||||
}
|
||||
if backend == nil {
|
||||
return nil, route, true
|
||||
}
|
||||
@@ -382,6 +402,7 @@ func (h *Handler) acquireScheduledBackend(ctx context.Context, w http.ResponseWr
|
||||
"path": r.URL.Path,
|
||||
"model": route.ModelName,
|
||||
})
|
||||
w.Header().Set("Retry-After", "5")
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
fmt.Fprintf(w, "queue full")
|
||||
} else {
|
||||
|
||||
Reference in New Issue
Block a user