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:
@@ -41,7 +41,7 @@ startup.
|
|||||||
| `JOB_TIMEOUT` | `15m` | Wait for a ComfyUI job to finish |
|
| `JOB_TIMEOUT` | `15m` | Wait for a ComfyUI job to finish |
|
||||||
| `LLM_WAIT_TIMEOUT` | `10m` | Max lock wait for an LLM request before 503 |
|
| `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) |
|
| `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 |
|
| `LOG_FORMAT` | `text` | `json` for structured JSON logs |
|
||||||
| `UNLOAD_POLL_INTERVAL` | `500ms` | `/api/ps` poll interval while unloading |
|
| `UNLOAD_POLL_INTERVAL` | `500ms` | `/api/ps` poll interval while unloading |
|
||||||
| `HISTORY_POLL_INTERVAL` | `1s` | `/history/<id>` poll interval while a job runs |
|
| `HISTORY_POLL_INTERVAL` | `1s` | `/history/<id>` poll interval while a job runs |
|
||||||
@@ -60,6 +60,12 @@ startup.
|
|||||||
`gpu_turnstile_llm_inflight`, `gpu_turnstile_image_pending`,
|
`gpu_turnstile_llm_inflight`, `gpu_turnstile_image_pending`,
|
||||||
`gpu_turnstile_image_jobs_total`, `gpu_turnstile_lock_wait_seconds`
|
`gpu_turnstile_image_jobs_total`, `gpu_turnstile_lock_wait_seconds`
|
||||||
(histogram, `kind="llm|image"`), `gpu_turnstile_unload_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
|
## Build and run
|
||||||
|
|
||||||
|
|||||||
@@ -112,7 +112,7 @@ load time. Off by default.
|
|||||||
| `JOB_TIMEOUT` | `15m` | wait for ComfyUI job |
|
| `JOB_TIMEOUT` | `15m` | wait for ComfyUI job |
|
||||||
| `LLM_WAIT_TIMEOUT` | `10m` | max time an LLM request waits for the lock before 503 |
|
| `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 |
|
| `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 |
|
| `UNLOAD_POLL_INTERVAL` | `500ms` | `/api/ps` poll interval while unloading |
|
||||||
| `HISTORY_POLL_INTERVAL` | `1s` | `/history/<id>` poll interval while a job runs |
|
| `HISTORY_POLL_INTERVAL` | `1s` | `/history/<id>` poll interval while a job runs |
|
||||||
| `PROBE_TIMEOUT` | `5s` | startup probe of both upstreams |
|
| `PROBE_TIMEOUT` | `5s` | startup probe of both upstreams |
|
||||||
@@ -136,7 +136,13 @@ start (`/api/version`, `/system_stats`); failure is logged, not fatal.
|
|||||||
`gpu_turnstile_image_jobs_total`, `gpu_turnstile_lock_wait_seconds`
|
`gpu_turnstile_image_jobs_total`, `gpu_turnstile_lock_wait_seconds`
|
||||||
(histogram, label `kind="llm|image"`), `gpu_turnstile_unload_seconds`.
|
(histogram, label `kind="llm|image"`), `gpu_turnstile_unload_seconds`.
|
||||||
- Structured logs (`log/slog`, JSON when `LOG_FORMAT=json`), one line per
|
- 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
|
## Edge cases to handle
|
||||||
|
|
||||||
@@ -148,11 +154,14 @@ start (`/api/version`, `/system_stats`); failure is logged, not fatal.
|
|||||||
`JOB_TIMEOUT` releases the lock; log at warn.
|
`JOB_TIMEOUT` releases the lock; log at warn.
|
||||||
- Ollama unreachable during unload: continue with the image job; the whole
|
- Ollama unreachable during unload: continue with the image job; the whole
|
||||||
point is not to block users on a misbehaving neighbour.
|
point is not to block users on a misbehaving neighbour.
|
||||||
- Upstream connection refused while proxying (service down or restarting):
|
- Upstream unreachable while proxying (connection refused, dial timeout,
|
||||||
retry with exponential backoff — `BACKOFF_INITIAL`, doubling per attempt,
|
DNS failure, TLS handshake error): retry with exponential backoff —
|
||||||
capped at `BACKOFF_MAX` — until the upstream answers or the client
|
`BACKOFF_INITIAL`, doubling per attempt, capped at `BACKOFF_MAX` — until
|
||||||
disconnects. Retrying a refused connection is safe: no request bytes were
|
the upstream answers or the client disconnects. These are safe to retry:
|
||||||
sent. Other upstream errors are not retried.
|
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
|
- `POST /prompt` with a body that ComfyUI rejects (400): lock released
|
||||||
immediately, body passed back.
|
immediately, body passed back.
|
||||||
- Websocket `/ws` connections are long-lived and never take the lock.
|
- Websocket `/ws` connections are long-lived and never take the lock.
|
||||||
|
|||||||
@@ -82,7 +82,7 @@ func loadConfig(getenv func(string) string) (config, error) {
|
|||||||
backoffMax: time.Minute,
|
backoffMax: time.Minute,
|
||||||
promptCaptureLimit: 64 * 1024,
|
promptCaptureLimit: 64 * 1024,
|
||||||
|
|
||||||
logLevel: slog.LevelInfo,
|
logLevel: slog.LevelWarn,
|
||||||
}
|
}
|
||||||
for _, e := range []struct {
|
for _, e := range []struct {
|
||||||
name string
|
name string
|
||||||
@@ -125,10 +125,15 @@ func loadConfig(getenv func(string) string) (config, error) {
|
|||||||
}
|
}
|
||||||
cfg.promptCaptureLimit = n
|
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
|
var level slog.Level
|
||||||
if err := level.UnmarshalText([]byte(v)); err != nil {
|
if err := level.UnmarshalText([]byte(logLevelValue)); err != nil {
|
||||||
return cfg, fmt.Errorf("LOG_LEVEL: %w", err)
|
return cfg, fmt.Errorf("LOGLEVEL: %w", err)
|
||||||
}
|
}
|
||||||
cfg.logLevel = level
|
cfg.logLevel = level
|
||||||
}
|
}
|
||||||
@@ -157,12 +162,29 @@ func main() {
|
|||||||
log := slog.New(handler)
|
log := slog.New(handler)
|
||||||
slog.SetDefault(log)
|
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,
|
"version", version,
|
||||||
"listen_ollama", cfg.listenOllama,
|
"listen_ollama", cfg.listenOllama,
|
||||||
"listen_comfy", cfg.listenComfy,
|
"listen_comfy", cfg.listenComfy,
|
||||||
"ollama_url", cfg.ollamaURL,
|
"ollama_url", cfg.ollamaURL,
|
||||||
"comfy_url", cfg.comfyURL,
|
"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)
|
ollamaClient, err := ollama.New(cfg.ollamaURL, log)
|
||||||
@@ -184,6 +206,7 @@ func main() {
|
|||||||
Comfy: comfyClient,
|
Comfy: comfyClient,
|
||||||
Metrics: metrics.New(),
|
Metrics: metrics.New(),
|
||||||
Log: log,
|
Log: log,
|
||||||
|
LogColor: !cfg.logJSON && os.Getenv("NO_COLOR") == "",
|
||||||
LLMWaitTimeout: cfg.llmWaitTimeout,
|
LLMWaitTimeout: cfg.llmWaitTimeout,
|
||||||
UnloadTimeout: cfg.unloadTimeout,
|
UnloadTimeout: cfg.unloadTimeout,
|
||||||
JobTimeout: cfg.jobTimeout,
|
JobTimeout: cfg.jobTimeout,
|
||||||
|
|||||||
+174
-27
@@ -4,16 +4,22 @@
|
|||||||
package proxy
|
package proxy
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"bufio"
|
||||||
"bytes"
|
"bytes"
|
||||||
"context"
|
"context"
|
||||||
|
"crypto/tls"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"io"
|
||||||
"log/slog"
|
"log/slog"
|
||||||
|
"net"
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httputil"
|
"net/http/httputil"
|
||||||
"net/url"
|
"net/url"
|
||||||
"syscall"
|
"os"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"gpu-turnstile/internal/comfy"
|
"gpu-turnstile/internal/comfy"
|
||||||
@@ -38,6 +44,10 @@ type Config struct {
|
|||||||
Metrics *metrics.Metrics
|
Metrics *metrics.Metrics
|
||||||
Log *slog.Logger
|
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
|
LLMWaitTimeout time.Duration
|
||||||
UnloadTimeout time.Duration
|
UnloadTimeout time.Duration
|
||||||
JobTimeout time.Duration
|
JobTimeout time.Duration
|
||||||
@@ -136,12 +146,13 @@ func New(cfg Config) (*Server, error) {
|
|||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// retryTransport retries requests that fail with "connection refused" —
|
// retryTransport retries requests whose failure means the upstream never
|
||||||
// typically the upstream service simply not being up yet — using an
|
// saw them — any dial-phase error (connection refused, dial timeout, DNS
|
||||||
// exponential backoff: initial, doubling each attempt, capped at max. The
|
// failure), TLS handshake errors — plus 5xx responses when the request
|
||||||
// retry loop runs until the request succeeds, fails with a different
|
// body can be replayed (GETs and requests with GetBody set). The wait
|
||||||
// error, or the client's context is cancelled. Retrying a refused
|
// doubles from initial up to max between attempts. The loop runs until
|
||||||
// connection is always safe: no request bytes were ever sent.
|
// the request succeeds, fails in a non-retryable way, or the client's
|
||||||
|
// context is cancelled.
|
||||||
type retryTransport struct {
|
type retryTransport struct {
|
||||||
base http.RoundTripper
|
base http.RoundTripper
|
||||||
initial time.Duration
|
initial time.Duration
|
||||||
@@ -149,37 +160,89 @@ type retryTransport struct {
|
|||||||
log *slog.Logger
|
log *slog.Logger
|
||||||
}
|
}
|
||||||
|
|
||||||
// wsaECONNREFUSED is the Windows error code for a refused connection.
|
// shouldRetry reports whether a RoundTrip error means the request never
|
||||||
// syscall.ECONNREFUSED is the POSIX value and never matches on Windows,
|
// reached the upstream application and is therefore safe to send again.
|
||||||
// where dial errors carry the WSA code.
|
func shouldRetry(err error) bool {
|
||||||
const wsaECONNREFUSED syscall.Errno = 10061
|
// Dial-phase failures: refused, timeout, unreachable, DNS (wrapped).
|
||||||
|
var opErr *net.OpError
|
||||||
func isConnRefused(err error) bool {
|
if errors.As(err, &opErr) && opErr.Op == "dial" {
|
||||||
if errors.Is(err, syscall.ECONNREFUSED) {
|
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
var errno syscall.Errno
|
var dnsErr *net.DNSError
|
||||||
return errors.As(err, &errno) && errno == wsaECONNREFUSED
|
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) {
|
// replayable reports whether the request body can be sent again. Bodies
|
||||||
if t.log != nil {
|
// streamed from the client (GetBody == nil) cannot, so 5xx responses to
|
||||||
t.log.Warn(msg, args...)
|
// 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) {
|
func (t *retryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||||
wait := t.initial
|
wait := t.initial
|
||||||
|
attempt := 0
|
||||||
for {
|
for {
|
||||||
resp, err := t.base.RoundTrip(req)
|
resp, err := t.base.RoundTrip(req)
|
||||||
if err == nil || !isConnRefused(err) {
|
switch {
|
||||||
return resp, err
|
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 {
|
select {
|
||||||
case <-req.Context().Done():
|
case <-req.Context().Done():
|
||||||
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
|
}
|
||||||
|
return nil, req.Context().Err()
|
||||||
case <-time.After(wait):
|
case <-time.After(wait):
|
||||||
}
|
}
|
||||||
wait *= 2
|
wait *= 2
|
||||||
@@ -222,6 +285,90 @@ func (s *Server) writeMetrics(w http.ResponseWriter) {
|
|||||||
s.cfg.Metrics.Render(w, string(state), n, pending)
|
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
|
// llmPaths are the Ollama endpoints that load models into VRAM and therefore
|
||||||
// take the LLM lock. Everything else passes through unlocked.
|
// take the LLM lock. Everything else passes through unlocked.
|
||||||
var llmPaths = map[string]bool{
|
var llmPaths = map[string]bool{
|
||||||
@@ -240,7 +387,7 @@ func isLLMRequest(r *http.Request) bool {
|
|||||||
|
|
||||||
// OllamaHandler serves the Ollama-facing listener.
|
// OllamaHandler serves the Ollama-facing listener.
|
||||||
func (s *Server) OllamaHandler() http.Handler {
|
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 {
|
switch r.URL.Path {
|
||||||
case "/healthz":
|
case "/healthz":
|
||||||
s.writeHealthz(w)
|
s.writeHealthz(w)
|
||||||
@@ -267,12 +414,12 @@ func (s *Server) OllamaHandler() http.Handler {
|
|||||||
}
|
}
|
||||||
defer s.cfg.Lock.ReleaseLLM()
|
defer s.cfg.Lock.ReleaseLLM()
|
||||||
s.ollamaProxy.ServeHTTP(w, r)
|
s.ollamaProxy.ServeHTTP(w, r)
|
||||||
})
|
}))
|
||||||
}
|
}
|
||||||
|
|
||||||
// ComfyHandler serves the ComfyUI-facing listener.
|
// ComfyHandler serves the ComfyUI-facing listener.
|
||||||
func (s *Server) ComfyHandler() http.Handler {
|
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" {
|
if r.URL.Path == "/healthz" {
|
||||||
s.writeHealthz(w)
|
s.writeHealthz(w)
|
||||||
return
|
return
|
||||||
@@ -282,7 +429,7 @@ func (s *Server) ComfyHandler() http.Handler {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
s.comfyProxy.ServeHTTP(w, r)
|
s.comfyProxy.ServeHTTP(w, r)
|
||||||
})
|
}))
|
||||||
}
|
}
|
||||||
|
|
||||||
// captureWriter passes the response through unchanged while recording the
|
// captureWriter passes the response through unchanged while recording the
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ package proxy
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"crypto/tls"
|
||||||
"io"
|
"io"
|
||||||
"net"
|
"net"
|
||||||
"net/http"
|
"net/http"
|
||||||
@@ -116,7 +117,91 @@ func TestRetryViaReverseProxy(t *testing.T) {
|
|||||||
defer cancel()
|
defer cancel()
|
||||||
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://127.0.0.1:1/", nil)
|
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://127.0.0.1:1/", nil)
|
||||||
_, err = rt.RoundTrip(req)
|
_, err = rt.RoundTrip(req)
|
||||||
if err == nil || !isConnRefused(err) {
|
if err == nil || !shouldRetry(err) {
|
||||||
t.Fatalf("err = %v, want connection refused", 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 }
|
||||||
|
|||||||
Reference in New Issue
Block a user