Extend retry backoff; rework logging (LOGLEVEL, request lines, colors)
Backoff now covers all dial-phase errors (refused, timeout, DNS), TLS handshake failures, and 5xx responses with replayable bodies; streamed POSTs are never replayed to avoid duplicate work. First retry of an episode logs at WARN, subsequent attempts at INFO. LOGLEVEL (LOG_LEVEL kept as alias) now defaults to warn: startup logs version plus every setting; INFO adds one line per incoming request and per response with status/duration, ANSI-colored in text mode (bypasses slog's escaping so colors render in docker compose logs); NO_COLOR or LOG_FORMAT=json disables colors.
This commit is contained in:
mram
+175
-28
@@ -4,16 +4,22 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httputil"
|
||||
"net/url"
|
||||
"syscall"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gpu-turnstile/internal/comfy"
|
||||
@@ -38,6 +44,10 @@ type Config struct {
|
||||
Metrics *metrics.Metrics
|
||||
Log *slog.Logger
|
||||
|
||||
// LogColor enables ANSI colors in per-request log lines. Ignored when
|
||||
// the log level is above INFO (request lines are not emitted at all).
|
||||
LogColor bool
|
||||
|
||||
LLMWaitTimeout time.Duration
|
||||
UnloadTimeout time.Duration
|
||||
JobTimeout time.Duration
|
||||
@@ -136,12 +146,13 @@ func New(cfg Config) (*Server, error) {
|
||||
}, nil
|
||||
}
|
||||
|
||||
// retryTransport retries requests that fail with "connection refused" —
|
||||
// typically the upstream service simply not being up yet — using an
|
||||
// exponential backoff: initial, doubling each attempt, capped at max. The
|
||||
// retry loop runs until the request succeeds, fails with a different
|
||||
// error, or the client's context is cancelled. Retrying a refused
|
||||
// connection is always safe: no request bytes were ever sent.
|
||||
// retryTransport retries requests whose failure means the upstream never
|
||||
// saw them — any dial-phase error (connection refused, dial timeout, DNS
|
||||
// failure), TLS handshake errors — plus 5xx responses when the request
|
||||
// body can be replayed (GETs and requests with GetBody set). The wait
|
||||
// doubles from initial up to max between attempts. The loop runs until
|
||||
// the request succeeds, fails in a non-retryable way, or the client's
|
||||
// context is cancelled.
|
||||
type retryTransport struct {
|
||||
base http.RoundTripper
|
||||
initial time.Duration
|
||||
@@ -149,37 +160,89 @@ type retryTransport struct {
|
||||
log *slog.Logger
|
||||
}
|
||||
|
||||
// wsaECONNREFUSED is the Windows error code for a refused connection.
|
||||
// syscall.ECONNREFUSED is the POSIX value and never matches on Windows,
|
||||
// where dial errors carry the WSA code.
|
||||
const wsaECONNREFUSED syscall.Errno = 10061
|
||||
|
||||
func isConnRefused(err error) bool {
|
||||
if errors.Is(err, syscall.ECONNREFUSED) {
|
||||
// shouldRetry reports whether a RoundTrip error means the request never
|
||||
// reached the upstream application and is therefore safe to send again.
|
||||
func shouldRetry(err error) bool {
|
||||
// Dial-phase failures: refused, timeout, unreachable, DNS (wrapped).
|
||||
var opErr *net.OpError
|
||||
if errors.As(err, &opErr) && opErr.Op == "dial" {
|
||||
return true
|
||||
}
|
||||
var errno syscall.Errno
|
||||
return errors.As(err, &errno) && errno == wsaECONNREFUSED
|
||||
var dnsErr *net.DNSError
|
||||
if errors.As(err, &dnsErr) {
|
||||
return true
|
||||
}
|
||||
// TLS handshake failures: the HTTP request was never written.
|
||||
var recordErr tls.RecordHeaderError
|
||||
if errors.As(err, &recordErr) {
|
||||
return true
|
||||
}
|
||||
var certErr *tls.CertificateVerificationError
|
||||
if errors.As(err, &certErr) {
|
||||
return true
|
||||
}
|
||||
var alertErr tls.AlertError
|
||||
if errors.As(err, &alertErr) {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (t *retryTransport) logWarn(msg string, args ...any) {
|
||||
if t.log != nil {
|
||||
t.log.Warn(msg, args...)
|
||||
// replayable reports whether the request body can be sent again. Bodies
|
||||
// streamed from the client (GetBody == nil) cannot, so 5xx responses to
|
||||
// POSTs are not retried: the upstream may have partially processed them,
|
||||
// and re-sending could duplicate work (e.g. a second ComfyUI prompt).
|
||||
func replayable(req *http.Request) bool {
|
||||
return req.Body == nil || req.Body == http.NoBody || req.GetBody != nil
|
||||
}
|
||||
|
||||
func (t *retryTransport) logAt(level slog.Level, msg string, args ...any) {
|
||||
if t.log != nil && t.log.Enabled(context.Background(), level) {
|
||||
t.log.Log(context.Background(), level, msg, args...)
|
||||
}
|
||||
}
|
||||
|
||||
func (t *retryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
wait := t.initial
|
||||
attempt := 0
|
||||
for {
|
||||
resp, err := t.base.RoundTrip(req)
|
||||
if err == nil || !isConnRefused(err) {
|
||||
return resp, err
|
||||
switch {
|
||||
case err != nil && !shouldRetry(err):
|
||||
return nil, err
|
||||
case err == nil && (resp.StatusCode < 500 || !replayable(req)):
|
||||
return resp, nil
|
||||
}
|
||||
t.logWarn("upstream connection refused; backing off",
|
||||
"path", req.URL.Path, "retry_in", wait)
|
||||
|
||||
// Retryable failure: a transport error, or a 5xx response.
|
||||
var reason string
|
||||
if err != nil {
|
||||
reason = err.Error()
|
||||
} else {
|
||||
reason = resp.Status
|
||||
io.Copy(io.Discard, resp.Body)
|
||||
resp.Body.Close()
|
||||
if req.GetBody != nil {
|
||||
if body, berr := req.GetBody(); berr == nil {
|
||||
req.Body = body
|
||||
}
|
||||
}
|
||||
}
|
||||
attempt++
|
||||
// One WARN per outage episode; subsequent attempts at INFO.
|
||||
level := slog.LevelInfo
|
||||
if attempt == 1 {
|
||||
level = slog.LevelWarn
|
||||
}
|
||||
t.logAt(level, "upstream unavailable; retrying with backoff",
|
||||
"path", req.URL.Path, "reason", reason, "attempt", attempt, "retry_in", wait)
|
||||
|
||||
select {
|
||||
case <-req.Context().Done():
|
||||
return nil, err
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return nil, req.Context().Err()
|
||||
case <-time.After(wait):
|
||||
}
|
||||
wait *= 2
|
||||
@@ -222,6 +285,90 @@ func (s *Server) writeMetrics(w http.ResponseWriter) {
|
||||
s.cfg.Metrics.Render(w, string(state), n, pending)
|
||||
}
|
||||
|
||||
// ANSI colors for per-request log lines.
|
||||
const (
|
||||
ansiReset = "\x1b[0m"
|
||||
ansiCyan = "\x1b[36m"
|
||||
ansiGreen = "\x1b[32m"
|
||||
ansiYellow = "\x1b[33m"
|
||||
ansiRed = "\x1b[31m"
|
||||
)
|
||||
|
||||
// statusRecorder remembers the response status while passing everything
|
||||
// through, including streaming flushes and websocket hijacks.
|
||||
type statusRecorder struct {
|
||||
http.ResponseWriter
|
||||
status int
|
||||
}
|
||||
|
||||
func (r *statusRecorder) WriteHeader(code int) {
|
||||
r.status = code
|
||||
r.ResponseWriter.WriteHeader(code)
|
||||
}
|
||||
|
||||
func (r *statusRecorder) Flush() {
|
||||
if f, ok := r.ResponseWriter.(http.Flusher); ok {
|
||||
f.Flush()
|
||||
}
|
||||
}
|
||||
|
||||
func (r *statusRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
|
||||
h, ok := r.ResponseWriter.(http.Hijacker)
|
||||
if !ok {
|
||||
return nil, nil, errors.New("response writer does not support hijacking")
|
||||
}
|
||||
return h.Hijack()
|
||||
}
|
||||
|
||||
func (r *statusRecorder) Unwrap() http.ResponseWriter { return r.ResponseWriter }
|
||||
|
||||
// reqLine emits one request log line. With color enabled slog cannot be
|
||||
// used: its text handler escapes the ANSI sequences, so the line is
|
||||
// written to stderr directly in the same key=value shape. Without color
|
||||
// it is a plain slog INFO line.
|
||||
func (s *Server) reqLine(log *slog.Logger, code, line string, attrs ...any) {
|
||||
if !s.cfg.LogColor {
|
||||
log.Info(line, attrs...)
|
||||
return
|
||||
}
|
||||
var sb strings.Builder
|
||||
sb.WriteString("time=" + time.Now().Format("2006-01-02T15:04:05.000Z07:00") + " level=INFO ")
|
||||
sb.WriteString(code + line + ansiReset)
|
||||
for i := 0; i+1 < len(attrs); i += 2 {
|
||||
fmt.Fprintf(&sb, " %v=%v", attrs[i], attrs[i+1])
|
||||
}
|
||||
sb.WriteByte('\n')
|
||||
os.Stderr.WriteString(sb.String())
|
||||
}
|
||||
|
||||
// logRequests logs one line per incoming request and one per completed
|
||||
// response at INFO level, colored when enabled: cyan "-->" for incoming,
|
||||
// green/yellow/red "<--" for responses by status class. At log levels
|
||||
// above INFO it is a pass-through.
|
||||
func (s *Server) logRequests(listener string, next http.Handler) http.Handler {
|
||||
log := s.log.With("listener", listener)
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if !log.Enabled(r.Context(), slog.LevelInfo) {
|
||||
next.ServeHTTP(w, r)
|
||||
return
|
||||
}
|
||||
start := time.Now()
|
||||
s.reqLine(log, ansiCyan, "--> "+r.Method+" "+r.URL.RequestURI(),
|
||||
"listener", listener, "remote", r.RemoteAddr)
|
||||
rec := &statusRecorder{ResponseWriter: w, status: http.StatusOK}
|
||||
next.ServeHTTP(rec, r)
|
||||
code := ansiGreen
|
||||
switch {
|
||||
case rec.status >= 500:
|
||||
code = ansiRed
|
||||
case rec.status >= 400:
|
||||
code = ansiYellow
|
||||
}
|
||||
s.reqLine(log, code, "<-- "+strconv.Itoa(rec.status)+" "+r.Method+" "+r.URL.RequestURI(),
|
||||
"listener", listener, "ms", time.Since(start).Milliseconds())
|
||||
})
|
||||
}
|
||||
|
||||
// llmPaths are the Ollama endpoints that load models into VRAM and therefore
|
||||
// take the LLM lock. Everything else passes through unlocked.
|
||||
var llmPaths = map[string]bool{
|
||||
@@ -240,7 +387,7 @@ func isLLMRequest(r *http.Request) bool {
|
||||
|
||||
// OllamaHandler serves the Ollama-facing listener.
|
||||
func (s *Server) OllamaHandler() http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
return s.logRequests("ollama", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/healthz":
|
||||
s.writeHealthz(w)
|
||||
@@ -267,12 +414,12 @@ func (s *Server) OllamaHandler() http.Handler {
|
||||
}
|
||||
defer s.cfg.Lock.ReleaseLLM()
|
||||
s.ollamaProxy.ServeHTTP(w, r)
|
||||
})
|
||||
}))
|
||||
}
|
||||
|
||||
// ComfyHandler serves the ComfyUI-facing listener.
|
||||
func (s *Server) ComfyHandler() http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
return s.logRequests("comfy", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/healthz" {
|
||||
s.writeHealthz(w)
|
||||
return
|
||||
@@ -282,7 +429,7 @@ func (s *Server) ComfyHandler() http.Handler {
|
||||
return
|
||||
}
|
||||
s.comfyProxy.ServeHTTP(w, r)
|
||||
})
|
||||
}))
|
||||
}
|
||||
|
||||
// captureWriter passes the response through unchanged while recording the
|
||||
|
||||
Reference in New Issue
Block a user