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.
2026-09-20 20:00:57 +02:00
parent 20d5439b5f
commit 42e1386811
5 changed files with 312 additions and 42 deletions
@@ -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
@@ -2,6 +2,7 @@ package proxy

 import (
 	"context"
+	"crypto/tls"
 	"io"
 	"net"
 	"net/http"
@@ -116,7 +117,91 @@ func TestRetryViaReverseProxy(t *testing.T) {
 	defer cancel()
 	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://127.0.0.1:1/", nil)
 	_, err = rt.RoundTrip(req)
-	if err == nil || !isConnRefused(err) {
+	if err == nil || !shouldRetry(err) {
 		t.Fatalf("err = %v, want connection refused", err)
 	}
 }
+
+// flakyStatus returns 500 for the first fails requests, then 200.
+type flakyStatus struct {
+	fails int
+	calls int
+}
+
+func (s *flakyStatus) RoundTrip(req *http.Request) (*http.Response, error) {
+	s.calls++
+	code := 200
+	if s.calls <= s.fails {
+		code = 500
+	}
+	return &http.Response{
+		StatusCode: code,
+		Status:     http.StatusText(code),
+		Body:       io.NopCloser(strings.NewReader("")),
+		Header:     make(http.Header),
+	}, nil
+}
+
+func TestRetryTransport5xxGet(t *testing.T) {
+	fs := &flakyStatus{fails: 2}
+	rt := &retryTransport{base: fs, initial: time.Millisecond, max: 2 * time.Millisecond}
+
+	req, _ := http.NewRequest(http.MethodGet, "http://upstream/api/version", nil)
+	resp, err := rt.RoundTrip(req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	resp.Body.Close()
+	if resp.StatusCode != 200 {
+		t.Fatalf("status = %d, want 200", resp.StatusCode)
+	}
+	if fs.calls != 3 {
+		t.Fatalf("calls = %d, want 3", fs.calls)
+	}
+}
+
+func TestRetryTransport5xxPostNotRetried(t *testing.T) {
+	fs := &flakyStatus{fails: 10}
+	rt := &retryTransport{base: fs, initial: time.Millisecond, max: time.Millisecond}
+
+	// A streamed body (no GetBody) must not be replayed after a 500.
+	req, _ := http.NewRequest(http.MethodPost, "http://upstream/prompt", io.NopCloser(strings.NewReader("{}")))
+	resp, err := rt.RoundTrip(req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	resp.Body.Close()
+	if resp.StatusCode != 500 {
+		t.Fatalf("status = %d, want 500", resp.StatusCode)
+	}
+	if fs.calls != 1 {
+		t.Fatalf("calls = %d, want 1 (no retry for streamed POST)", fs.calls)
+	}
+}
+
+func TestShouldRetryClassification(t *testing.T) {
+	cases := []struct {
+		name string
+		err  error
+		want bool
+	}{
+		{"dial refused", &net.OpError{Op: "dial", Net: "tcp", Err: syscall.ECONNREFUSED}, true},
+		{"dial timeout", &net.OpError{Op: "dial", Net: "tcp", Err: timeoutErr{}}, true},
+		{"dns", &net.DNSError{Err: "no such host", IsNotFound: true}, true},
+		{"tls record", tls.RecordHeaderError{Msg: "bad"}, true},
+		{"tls alert", tls.AlertError(42), true},
+		{"read error mid-request", &net.OpError{Op: "read", Net: "tcp", Err: syscall.ECONNRESET}, false},
+		{"plain error", io.EOF, false},
+	}
+	for _, c := range cases {
+		if got := shouldRetry(c.err); got != c.want {
+			t.Errorf("%s: shouldRetry = %v, want %v", c.name, got, c.want)
+		}
+	}
+}
+
+type timeoutErr struct{}
+
+func (timeoutErr) Error() string   { return "i/o timeout" }
+func (timeoutErr) Timeout() bool   { return true }
+func (timeoutErr) Temporary() bool { return true }