3 Commits
Author SHA1 Message Date
mram 5f6a22c2cf Pin compose example to v0.1.2
ci / test (push) Successful in 47s
ci / docker (push) Successful in 1m5s
2026-09-20 20:07:38 +02:00
mram 42e1386811 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
mram 20d5439b5f Retry upstream connection-refused with exponential backoff
A refused dial (service down/restarting) is retried with a wait that
doubles from BACKOFF_INITIAL (1s) up to BACKOFF_MAX (60s) until the
upstream answers or the client disconnects. Handles the Windows WSA
errno (10061) as well as POSIX ECONNREFUSED. compose.yaml.example now
uses host.docker.internal like the working local deployment.
2026-09-20 19:45:48 +02:00
6 changed files with 524 additions and 29 deletions
+9 -1
View File
@@ -41,7 +41,7 @@ startup.
| `JOB_TIMEOUT` | `15m` | Wait for a ComfyUI job to finish |
| `LLM_WAIT_TIMEOUT` | `10m` | Max lock wait for an LLM request before 503 |
| `WARM_MODEL` | _(empty)_ | Model to reload after an image job (off by default) |
| `LOG_LEVEL` | `info` | `debug` logs every lock transition |
| `LOGLEVEL` | `warn` | `info` logs every request (colored arrows in text mode), `debug` adds lock transitions. `LOG_LEVEL` works as an alias |
| `LOG_FORMAT` | `text` | `json` for structured JSON logs |
| `UNLOAD_POLL_INTERVAL` | `500ms` | `/api/ps` poll interval while unloading |
| `HISTORY_POLL_INTERVAL` | `1s` | `/history/<id>` poll interval while a job runs |
@@ -49,6 +49,8 @@ startup.
| `FREE_TIMEOUT` | `30s` | `POST /free` call after an image job |
| `WARM_TIMEOUT` | `2m` | Warm-model reload after an image job |
| `SHUTDOWN_TIMEOUT` | `10s` | Graceful shutdown on SIGINT/SIGTERM |
| `BACKOFF_INITIAL` | `1s` | First retry wait when an upstream refuses a connection |
| `BACKOFF_MAX` | `60s` | Cap for the exponential retry backoff |
| `PROMPT_CAPTURE_LIMIT` | `65536` | Bytes of the `/prompt` response buffered to find `prompt_id` (pass-through is unaffected) |
## Observability
@@ -58,6 +60,12 @@ startup.
`gpu_turnstile_llm_inflight`, `gpu_turnstile_image_pending`,
`gpu_turnstile_image_jobs_total`, `gpu_turnstile_lock_wait_seconds`
(histogram, `kind="llm|image"`), `gpu_turnstile_unload_seconds`.
- Logs: startup logs the version and every setting (visible even at the
default `warn` level). With `LOGLEVEL=info` or `debug`, every request
logs a `-->` incoming line and a `<--` response line with status and
duration — ANSI-colored (cyan incoming; green/yellow/red by status
class) in text mode, which renders in `docker compose logs` on Windows
Terminal. Set `NO_COLOR` to disable colors.
## Build and run
+18 -2
View File
@@ -112,13 +112,15 @@ load time. Off by default.
| `JOB_TIMEOUT` | `15m` | wait for ComfyUI job |
| `LLM_WAIT_TIMEOUT` | `10m` | max time an LLM request waits for the lock before 503 |
| `WARM_MODEL` | `` | optional model to reload after an image job |
| `LOG_LEVEL` | `info` | `debug` logs every lock transition |
| `LOGLEVEL` | `warn` | `info` logs every request (colored arrows in text mode), `debug` adds lock transitions. `LOG_LEVEL` is accepted as an alias |
| `UNLOAD_POLL_INTERVAL` | `500ms` | `/api/ps` poll interval while unloading |
| `HISTORY_POLL_INTERVAL` | `1s` | `/history/<id>` poll interval while a job runs |
| `PROBE_TIMEOUT` | `5s` | startup probe of both upstreams |
| `FREE_TIMEOUT` | `30s` | `POST /free` call after an image job |
| `WARM_TIMEOUT` | `2m` | warm-model reload after an image job |
| `SHUTDOWN_TIMEOUT` | `10s` | graceful shutdown on SIGINT/SIGTERM |
| `BACKOFF_INITIAL` | `1s` | first retry wait when an upstream refuses a connection |
| `BACKOFF_MAX` | `60s` | cap for the exponential retry backoff |
| `PROMPT_CAPTURE_LIMIT` | `65536` | bytes of the `/prompt` response buffered to find `prompt_id` (pass-through is unaffected) |
Startup fails fast on unparsable values. Both upstreams are probed once at
@@ -134,7 +136,13 @@ start (`/api/version`, `/system_stats`); failure is logged, not fatal.
`gpu_turnstile_image_jobs_total`, `gpu_turnstile_lock_wait_seconds`
(histogram, label `kind="llm|image"`), `gpu_turnstile_unload_seconds`.
- Structured logs (`log/slog`, JSON when `LOG_FORMAT=json`), one line per
state transition and per image job phase with `prompt_id`.
state transition and per image job phase with `prompt_id`. Startup logs
the version and every setting (visible even at the default `warn`
level). With `LOGLEVEL=info` or `debug`, every request logs a `-->`
incoming line and a `<--` response line with status and duration —
ANSI-colored (cyan incoming; green/yellow/red by status class) in text
mode, which renders in `docker compose logs` on Windows Terminal. Set
`NO_COLOR` to disable colors.
## Edge cases to handle
@@ -146,6 +154,14 @@ start (`/api/version`, `/system_stats`); failure is logged, not fatal.
`JOB_TIMEOUT` releases the lock; log at warn.
- Ollama unreachable during unload: continue with the image job; the whole
point is not to block users on a misbehaving neighbour.
- Upstream unreachable while proxying (connection refused, dial timeout,
DNS failure, TLS handshake error): retry with exponential backoff —
`BACKOFF_INITIAL`, doubling per attempt, capped at `BACKOFF_MAX` — until
the upstream answers or the client disconnects. These are safe to retry:
the request never reached the upstream application. 5xx responses are
retried the same way, but only when the request body can be replayed
(GETs, or bodies with `GetBody`); streamed POSTs are never replayed to
avoid duplicate work such as a double-enqueued ComfyUI prompt.
- `POST /prompt` with a body that ComfyUI rejects (400): lock released
immediately, body passed back.
- Websocket `/ws` connections are long-lived and never take the lock.
+36 -5
View File
@@ -40,6 +40,8 @@ type config struct {
freeTimeout time.Duration
warmTimeout time.Duration
shutdownTimeout time.Duration
backoffInitial time.Duration
backoffMax time.Duration
promptCaptureLimit int64
warmModel string
@@ -76,9 +78,11 @@ func loadConfig(getenv func(string) string) (config, error) {
freeTimeout: 30 * time.Second,
warmTimeout: 2 * time.Minute,
shutdownTimeout: 10 * time.Second,
backoffInitial: time.Second,
backoffMax: time.Minute,
promptCaptureLimit: 64 * 1024,
logLevel: slog.LevelInfo,
logLevel: slog.LevelWarn,
}
for _, e := range []struct {
name string
@@ -107,6 +111,8 @@ func loadConfig(getenv func(string) string) (config, error) {
{"FREE_TIMEOUT", &cfg.freeTimeout},
{"WARM_TIMEOUT", &cfg.warmTimeout},
{"SHUTDOWN_TIMEOUT", &cfg.shutdownTimeout},
{"BACKOFF_INITIAL", &cfg.backoffInitial},
{"BACKOFF_MAX", &cfg.backoffMax},
} {
if err := envDuration(getenv, e.name, e.dst); err != nil {
return cfg, err
@@ -119,10 +125,15 @@ func loadConfig(getenv func(string) string) (config, error) {
}
cfg.promptCaptureLimit = n
}
if v := getenv("LOG_LEVEL"); v != "" {
// LOGLEVEL is the canonical spelling; LOG_LEVEL is kept as an alias.
logLevelValue := getenv("LOGLEVEL")
if logLevelValue == "" {
logLevelValue = getenv("LOG_LEVEL")
}
if logLevelValue != "" {
var level slog.Level
if err := level.UnmarshalText([]byte(v)); err != nil {
return cfg, fmt.Errorf("LOG_LEVEL: %w", err)
if err := level.UnmarshalText([]byte(logLevelValue)); err != nil {
return cfg, fmt.Errorf("LOGLEVEL: %w", err)
}
cfg.logLevel = level
}
@@ -151,12 +162,29 @@ func main() {
log := slog.New(handler)
slog.SetDefault(log)
log.Info("starting gpu-turnstile",
// The startup line carries the version and every setting and is emitted
// at WARN so it is visible even with the default (quiet) log level.
log.Log(context.Background(), slog.LevelWarn, "starting gpu-turnstile",
"version", version,
"listen_ollama", cfg.listenOllama,
"listen_comfy", cfg.listenComfy,
"ollama_url", cfg.ollamaURL,
"comfy_url", cfg.comfyURL,
"unload_timeout", cfg.unloadTimeout,
"job_timeout", cfg.jobTimeout,
"llm_wait_timeout", cfg.llmWaitTimeout,
"unload_poll_interval", cfg.unloadPollInterval,
"history_poll_interval", cfg.historyPollInterval,
"probe_timeout", cfg.probeTimeout,
"free_timeout", cfg.freeTimeout,
"warm_timeout", cfg.warmTimeout,
"shutdown_timeout", cfg.shutdownTimeout,
"backoff_initial", cfg.backoffInitial,
"backoff_max", cfg.backoffMax,
"prompt_capture_limit", cfg.promptCaptureLimit,
"warm_model", cfg.warmModel,
"log_level", cfg.logLevel,
"log_format", map[bool]string{true: "json", false: "text"}[cfg.logJSON],
)
ollamaClient, err := ollama.New(cfg.ollamaURL, log)
@@ -178,6 +206,7 @@ func main() {
Comfy: comfyClient,
Metrics: metrics.New(),
Log: log,
LogColor: !cfg.logJSON && os.Getenv("NO_COLOR") == "",
LLMWaitTimeout: cfg.llmWaitTimeout,
UnloadTimeout: cfg.unloadTimeout,
JobTimeout: cfg.jobTimeout,
@@ -185,6 +214,8 @@ func main() {
HistoryPollInterval: cfg.historyPollInterval,
FreeTimeout: cfg.freeTimeout,
WarmTimeout: cfg.warmTimeout,
BackoffInitial: cfg.backoffInitial,
BackoffMax: cfg.backoffMax,
PromptCaptureLimit: cfg.promptCaptureLimit,
WarmModel: cfg.warmModel,
})
+10 -4
View File
@@ -1,14 +1,20 @@
# Example deployment for gpu-turnstile. Copy to compose.yaml and adjust.
#
# gpu-turnstile listens on the ports the services normally use; the actual
# Ollama and ComfyUI instances run one port higher on the workstation.
# Ollama and ComfyUI instances run one port higher (11435 / 8189) and must
# bind 0.0.0.0 so the container can reach them (OLLAMA_HOST=0.0.0.0:11435,
# ComfyUI --listen 0.0.0.0 --port 8189).
services:
gpu-turnstile:
image: git.rambossek.at/public/gpu-turnstile:v0.1.1
image: git.rambossek.at/public/gpu-turnstile:v0.1.2
restart: unless-stopped
environment:
OLLAMA_URL: http://<workstation-ip>:11435
COMFY_URL: http://<workstation-ip>:8189
# Services on the Docker host itself:
OLLAMA_URL: http://host.docker.internal:11435
COMFY_URL: http://host.docker.internal:8189
# Services on another machine: use its LAN IP instead, e.g.
# OLLAMA_URL: http://192.168.1.10:11435
# COMFY_URL: http://192.168.1.10:8189
# UNLOAD_TIMEOUT: 60s
# JOB_TIMEOUT: 15m
# LLM_WAIT_TIMEOUT: 10m
+234 -7
View File
@@ -4,15 +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"
"os"
"strconv"
"strings"
"time"
"gpu-turnstile/internal/comfy"
@@ -37,10 +44,21 @@ 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
// BackoffInitial and BackoffMax control the exponential retry backoff
// when an upstream refuses a connection: the wait doubles from
// BackoffInitial up to BackoffMax between attempts. Zero selects the
// defaults (1s / 60s).
BackoffInitial time.Duration
BackoffMax time.Duration
// UnloadPollInterval and HistoryPollInterval override the clients'
// /api/ps and /history poll intervals when > 0.
UnloadPollInterval time.Duration
@@ -62,6 +80,8 @@ type Server struct {
freeTimeout time.Duration
warmTimeout time.Duration
captureLimit int64
backoffInitial time.Duration
backoffMax time.Duration
ollamaProxy *httputil.ReverseProxy
comfyProxy *httputil.ReverseProxy
@@ -99,19 +119,142 @@ func New(cfg Config) (*Server, error) {
if captureLimit <= 0 {
captureLimit = defaultCaptureLimit
}
backoffInitial := cfg.BackoffInitial
if backoffInitial <= 0 {
backoffInitial = time.Second
}
backoffMax := cfg.BackoffMax
if backoffMax <= 0 {
backoffMax = time.Minute
}
retry := &retryTransport{
base: http.DefaultTransport,
initial: backoffInitial,
max: backoffMax,
log: log,
}
return &Server{
cfg: cfg,
log: log,
freeTimeout: freeTimeout,
warmTimeout: warmTimeout,
captureLimit: captureLimit,
ollamaProxy: newReverseProxy(ollamaURL, log.With("upstream", "ollama")),
comfyProxy: newReverseProxy(comfyURL, log.With("upstream", "comfy")),
backoffInitial: backoffInitial,
backoffMax: backoffMax,
ollamaProxy: newReverseProxy(ollamaURL, retry, log.With("upstream", "ollama")),
comfyProxy: newReverseProxy(comfyURL, retry, log.With("upstream", "comfy")),
}, nil
}
func newReverseProxy(target *url.URL, log *slog.Logger) *httputil.ReverseProxy {
// 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
max time.Duration
log *slog.Logger
}
// 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 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
}
// 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)
switch {
case err != nil && !shouldRetry(err):
return nil, err
case err == nil && (resp.StatusCode < 500 || !replayable(req)):
return resp, nil
}
// 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():
if err != nil {
return nil, err
}
return nil, req.Context().Err()
case <-time.After(wait):
}
wait *= 2
if wait > t.max {
wait = t.max
}
}
}
func newReverseProxy(target *url.URL, transport http.RoundTripper, log *slog.Logger) *httputil.ReverseProxy {
return &httputil.ReverseProxy{
Transport: transport,
Rewrite: func(pr *httputil.ProxyRequest) {
pr.SetURL(target)
pr.SetXForwarded()
@@ -142,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{
@@ -160,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)
@@ -187,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
@@ -202,7 +429,7 @@ func (s *Server) ComfyHandler() http.Handler {
return
}
s.comfyProxy.ServeHTTP(w, r)
})
}))
}
// captureWriter passes the response through unchanged while recording the
+207
View File
@@ -0,0 +1,207 @@
package proxy
import (
"context"
"crypto/tls"
"io"
"net"
"net/http"
"strings"
"syscall"
"testing"
"time"
)
// stubTransport fails with ECONNREFUSED for the first fails requests, then
// returns a 200 response.
type stubTransport struct {
fails int
calls int
}
func (s *stubTransport) RoundTrip(req *http.Request) (*http.Response, error) {
s.calls++
if s.calls <= s.fails {
return nil, &net.OpError{Op: "dial", Net: "tcp", Err: syscall.ECONNREFUSED}
}
return &http.Response{
StatusCode: 200,
Body: io.NopCloser(strings.NewReader("ok")),
Header: make(http.Header),
}, nil
}
func TestRetryTransportBackoff(t *testing.T) {
st := &stubTransport{fails: 3}
rt := &retryTransport{
base: st,
initial: 10 * time.Millisecond,
max: 25 * time.Millisecond,
}
start := time.Now()
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 st.calls != 4 {
t.Fatalf("calls = %d, want 4", st.calls)
}
// Waits: 10ms + 20ms + 25ms (capped) = 55ms minimum.
elapsed := time.Since(start)
if elapsed < 50*time.Millisecond {
t.Fatalf("elapsed = %v, want >= ~55ms of backoff", elapsed)
}
if elapsed > 5*time.Second {
t.Fatalf("elapsed = %v, suspiciously long", elapsed)
}
}
func TestRetryTransportNonRefusedErrorNotRetried(t *testing.T) {
rt := &retryTransport{
base: &stubTransport{fails: 0},
initial: time.Millisecond,
max: time.Millisecond,
}
req, _ := http.NewRequest(http.MethodGet, "http://upstream/", nil)
resp, err := rt.RoundTrip(req)
if err != nil || resp.StatusCode != 200 {
t.Fatalf("resp=%v err=%v", resp, err)
}
}
type alwaysRefused struct{ calls int }
func (a *alwaysRefused) RoundTrip(*http.Request) (*http.Response, error) {
a.calls++
return nil, &net.OpError{Op: "dial", Net: "tcp", Err: syscall.ECONNREFUSED}
}
func TestRetryTransportContextCancel(t *testing.T) {
ar := &alwaysRefused{}
rt := &retryTransport{base: ar, initial: time.Second, max: time.Second}
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://upstream/", nil)
start := time.Now()
_, err := rt.RoundTrip(req)
if err == nil {
t.Fatal("expected error after context cancellation")
}
if time.Since(start) > 2*time.Second {
t.Fatal("retry loop did not stop on context cancellation")
}
}
func TestRetryViaReverseProxy(t *testing.T) {
// Point the proxy at a port nothing listens on; the request should be
// retried (not instantly 502) until the client context ends.
srv, err := New(Config{
OllamaURL: "http://127.0.0.1:1",
ComfyURL: "http://127.0.0.1:1",
Lock: nil,
Metrics: nil,
BackoffInitial: 10 * time.Millisecond,
BackoffMax: 20 * time.Millisecond,
})
if err != nil {
t.Fatal(err)
}
_ = srv // construction must not panic with minimal config
rt := &retryTransport{base: http.DefaultTransport, initial: 10 * time.Millisecond, max: 20 * time.Millisecond}
ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
defer cancel()
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://127.0.0.1:1/", nil)
_, err = rt.RoundTrip(req)
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 }