From 78a5d6baa2d0bc467066e4ad5b3f4b6203e00e1c Mon Sep 17 00:00:00 2001
From: Admin Bot <admin@riotpiao.com>
Date: Sun, 27 Sep 2026 19:13:38 +0800
Subject: [PATCH] fix(proxy): mutex on transports, ctx-aware warmup delay,
 Retry-After on 503

- 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
---
 internal/proxy/proxy.go | 35 ++++++++++++++++++++++++++++-------
 1 file changed, 28 insertions(+), 7 deletions(-)

diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go
index ccc82b6..c3d9972 100644
--- a/internal/proxy/proxy.go
+++ b/internal/proxy/proxy.go
@@ -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 {