Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5f6a22c2cf | ||
|
|
42e1386811 | ||
|
|
20d5439b5f | ||
|
|
3aaa5d80a9 | ||
|
|
cef845c2e3 | ||
|
|
cf9be55a6c | ||
|
|
0f950f3134 | ||
|
|
73fb8ac3ad |
@@ -37,13 +37,19 @@ jobs:
|
||||
exit 1
|
||||
fi
|
||||
|
||||
- name: Compute lowercase image name
|
||||
id: meta
|
||||
run: |
|
||||
REPO=git.rambossek.at/$(echo "${{ gitea.repository }}" | tr '[:upper:]' '[:lower:]')
|
||||
echo "image=$REPO" >> "$GITHUB_OUTPUT"
|
||||
|
||||
- uses: docker/setup-buildx-action@v3
|
||||
|
||||
- uses: docker/login-action@v3
|
||||
with:
|
||||
registry: git.rambossek.at
|
||||
username: ${{ gitea.actor }}
|
||||
password: ${{ secrets.GITEA_TOKEN }}
|
||||
password: ${{ secrets.REGISTRY_TOKEN }}
|
||||
|
||||
- uses: docker/build-push-action@v6
|
||||
with:
|
||||
@@ -52,5 +58,5 @@ jobs:
|
||||
build-args: |
|
||||
VERSION=${{ gitea.ref_name }}
|
||||
tags: |
|
||||
git.rambossek.at/${{ gitea.repository }}:${{ gitea.ref_name }}
|
||||
git.rambossek.at/${{ gitea.repository }}:latest
|
||||
${{ steps.meta.outputs.image }}:${{ gitea.ref_name }}
|
||||
${{ steps.meta.outputs.image }}:latest
|
||||
|
||||
@@ -1,2 +1,4 @@
|
||||
/gpu-turnstile
|
||||
/gpu-turnstile.exe
|
||||
/compose.yaml
|
||||
/compose.yml
|
||||
|
||||
@@ -41,8 +41,17 @@ 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 |
|
||||
| `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) |
|
||||
|
||||
## Observability
|
||||
|
||||
@@ -51,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
|
||||
|
||||
@@ -73,6 +88,10 @@ runs `go vet` and `go test -race`, and pushing a semantic-version tag
|
||||
`git.rambossek.at/<owner>/gpu-turnstile:vX.Y.Z` (and updates `:latest`).
|
||||
No images are built from branches.
|
||||
|
||||
The registry login needs one repository secret (Settings → Actions →
|
||||
Secrets): `REGISTRY_TOKEN` — an access token with `write:package` scope.
|
||||
The automatic `GITEA_TOKEN` cannot push packages.
|
||||
|
||||
## Development
|
||||
|
||||
```sh
|
||||
|
||||
@@ -81,12 +81,14 @@ ComfyUI listener (`:8188` → `COMFY_URL`):
|
||||
2. Unload Ollama: `GET /api/ps`; for each model `POST /api/generate
|
||||
{"model":M,"keep_alive":0}`; if that returns non-2xx (embedding-only
|
||||
models), `POST /api/embed {"model":M,"input":"x","keep_alive":0}`. Poll
|
||||
`/api/ps` every 500 ms until empty or `UNLOAD_TIMEOUT`. On timeout: log and
|
||||
`/api/ps` every `UNLOAD_POLL_INTERVAL` (default 500 ms) until empty or
|
||||
`UNLOAD_TIMEOUT`. On timeout: log and
|
||||
continue (degrade, don't fail the user's request).
|
||||
3. Forward the original request body to ComfyUI `/prompt`, return status,
|
||||
headers and body to the caller unchanged, flush.
|
||||
4. If the response is 200 and contains `prompt_id`: in a goroutine, poll
|
||||
`GET /history/<prompt_id>` every 1 s until the entry has
|
||||
`GET /history/<prompt_id>` every `HISTORY_POLL_INTERVAL` (default 1 s)
|
||||
until the entry has
|
||||
`status.completed == true`, `status.status_str == "error"`, or
|
||||
`JOB_TIMEOUT`. Then `POST /free {"unload_models":true,"free_memory":true}`.
|
||||
Then release the image lock.
|
||||
@@ -110,7 +112,16 @@ 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
|
||||
start (`/api/version`, `/system_stats`); failure is logged, not fatal.
|
||||
@@ -125,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
|
||||
|
||||
@@ -137,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.
|
||||
@@ -185,8 +210,12 @@ are new.
|
||||
available in the runner image
|
||||
2. on a version tag only (`vX.Y.Z`, enforced): build the image with buildx
|
||||
and push it to the Gitea registry
|
||||
`git.rambossek.at/<owner>/gpu-turnstile` tagged `:<tag>` and `:latest`,
|
||||
using the workflow token (`${{ secrets.GITEA_TOKEN }}` / `gitea.actor`)
|
||||
`git.rambossek.at/<owner>/gpu-turnstile` tagged `:<tag>` and `:latest`
|
||||
(the repository path is lowercased in the workflow; Docker registry
|
||||
names must be lowercase).
|
||||
Login uses the repo secret `REGISTRY_TOKEN` (an access token with
|
||||
`write:package` scope) because the automatic `GITEA_TOKEN` cannot push
|
||||
packages; the username is just `gitea.actor`.
|
||||
- Release: a git tag `vX.Y.Z` produces the versioned image; the Open WebUI
|
||||
compose pins that tag. No images are built from branches.
|
||||
|
||||
@@ -194,7 +223,7 @@ are new.
|
||||
|
||||
```yaml
|
||||
gpu-turnstile:
|
||||
image: git.rambossek.at/<owner>/gpu-turnstile:v0.1.0
|
||||
image: git.rambossek.at/<owner>/gpu-turnstile:v0.1.0 # owner lowercased, e.g. "public"
|
||||
environment:
|
||||
OLLAMA_URL: http://<workstation-ip>:11435
|
||||
COMFY_URL: http://<workstation-ip>:8189
|
||||
|
||||
+89
-21
@@ -10,6 +10,7 @@ import (
|
||||
"net/http"
|
||||
"os"
|
||||
"os/signal"
|
||||
"strconv"
|
||||
"strings"
|
||||
"syscall"
|
||||
"time"
|
||||
@@ -32,9 +33,20 @@ type config struct {
|
||||
unloadTimeout time.Duration
|
||||
jobTimeout time.Duration
|
||||
llmWaitTimeout time.Duration
|
||||
warmModel string
|
||||
logLevel slog.Level
|
||||
logJSON bool
|
||||
|
||||
unloadPollInterval time.Duration
|
||||
historyPollInterval time.Duration
|
||||
probeTimeout time.Duration
|
||||
freeTimeout time.Duration
|
||||
warmTimeout time.Duration
|
||||
shutdownTimeout time.Duration
|
||||
backoffInitial time.Duration
|
||||
backoffMax time.Duration
|
||||
promptCaptureLimit int64
|
||||
|
||||
warmModel string
|
||||
logLevel slog.Level
|
||||
logJSON bool
|
||||
}
|
||||
|
||||
func envDuration(getenv func(string) string, name string, dst *time.Duration) error {
|
||||
@@ -59,7 +71,18 @@ func loadConfig(getenv func(string) string) (config, error) {
|
||||
unloadTimeout: time.Minute,
|
||||
jobTimeout: 15 * time.Minute,
|
||||
llmWaitTimeout: 10 * time.Minute,
|
||||
logLevel: slog.LevelInfo,
|
||||
|
||||
unloadPollInterval: 500 * time.Millisecond,
|
||||
historyPollInterval: time.Second,
|
||||
probeTimeout: 5 * time.Second,
|
||||
freeTimeout: 30 * time.Second,
|
||||
warmTimeout: 2 * time.Minute,
|
||||
shutdownTimeout: 10 * time.Second,
|
||||
backoffInitial: time.Second,
|
||||
backoffMax: time.Minute,
|
||||
promptCaptureLimit: 64 * 1024,
|
||||
|
||||
logLevel: slog.LevelWarn,
|
||||
}
|
||||
for _, e := range []struct {
|
||||
name string
|
||||
@@ -82,15 +105,35 @@ func loadConfig(getenv func(string) string) (config, error) {
|
||||
{"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},
|
||||
} {
|
||||
if err := envDuration(getenv, e.name, e.dst); err != nil {
|
||||
return cfg, err
|
||||
}
|
||||
}
|
||||
if v := getenv("LOG_LEVEL"); v != "" {
|
||||
if v := getenv("PROMPT_CAPTURE_LIMIT"); v != "" {
|
||||
n, err := strconv.ParseInt(v, 10, 64)
|
||||
if err != nil || n < 0 {
|
||||
return cfg, fmt.Errorf("PROMPT_CAPTURE_LIMIT: must be a non-negative integer (bytes)")
|
||||
}
|
||||
cfg.promptCaptureLimit = n
|
||||
}
|
||||
// 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
|
||||
}
|
||||
@@ -119,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)
|
||||
@@ -139,17 +199,25 @@ func main() {
|
||||
}
|
||||
|
||||
srv, err := proxy.New(proxy.Config{
|
||||
OllamaURL: cfg.ollamaURL,
|
||||
ComfyURL: cfg.comfyURL,
|
||||
Lock: lock.New(log),
|
||||
Ollama: ollamaClient,
|
||||
Comfy: comfyClient,
|
||||
Metrics: metrics.New(),
|
||||
Log: log,
|
||||
LLMWaitTimeout: cfg.llmWaitTimeout,
|
||||
UnloadTimeout: cfg.unloadTimeout,
|
||||
JobTimeout: cfg.jobTimeout,
|
||||
WarmModel: cfg.warmModel,
|
||||
OllamaURL: cfg.ollamaURL,
|
||||
ComfyURL: cfg.comfyURL,
|
||||
Lock: lock.New(log),
|
||||
Ollama: ollamaClient,
|
||||
Comfy: comfyClient,
|
||||
Metrics: metrics.New(),
|
||||
Log: log,
|
||||
LogColor: !cfg.logJSON && os.Getenv("NO_COLOR") == "",
|
||||
LLMWaitTimeout: cfg.llmWaitTimeout,
|
||||
UnloadTimeout: cfg.unloadTimeout,
|
||||
JobTimeout: cfg.jobTimeout,
|
||||
UnloadPollInterval: cfg.unloadPollInterval,
|
||||
HistoryPollInterval: cfg.historyPollInterval,
|
||||
FreeTimeout: cfg.freeTimeout,
|
||||
WarmTimeout: cfg.warmTimeout,
|
||||
BackoffInitial: cfg.backoffInitial,
|
||||
BackoffMax: cfg.backoffMax,
|
||||
PromptCaptureLimit: cfg.promptCaptureLimit,
|
||||
WarmModel: cfg.warmModel,
|
||||
})
|
||||
if err != nil {
|
||||
log.Error("invalid configuration", "err", err)
|
||||
@@ -157,7 +225,7 @@ func main() {
|
||||
}
|
||||
|
||||
// Probe both upstreams once; failure is logged, not fatal.
|
||||
probeCtx, probeCancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
probeCtx, probeCancel := context.WithTimeout(context.Background(), cfg.probeTimeout)
|
||||
if err := ollamaClient.Probe(probeCtx); err != nil {
|
||||
log.Warn("ollama probe failed", "url", cfg.ollamaURL, "err", err)
|
||||
}
|
||||
@@ -186,7 +254,7 @@ func main() {
|
||||
log.Info("shutting down")
|
||||
}
|
||||
|
||||
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), cfg.shutdownTimeout)
|
||||
defer shutdownCancel()
|
||||
ollamaSrv.Shutdown(shutdownCtx)
|
||||
comfySrv.Shutdown(shutdownCtx)
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
# 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 (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.2
|
||||
restart: unless-stopped
|
||||
environment:
|
||||
# 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
|
||||
# WARM_MODEL: qwen3:14b
|
||||
# LOG_LEVEL: info
|
||||
# LOG_FORMAT: json
|
||||
ports:
|
||||
- "11434:11434" # LiteLLM api_base -> http://gpu-turnstile:11434
|
||||
- "8188:8188" # Open WebUI COMFYUI_BASE_URL -> http://gpu-turnstile:8188
|
||||
networks: [internal]
|
||||
|
||||
networks:
|
||||
internal:
|
||||
external: true
|
||||
+285
-21
@@ -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"
|
||||
@@ -21,10 +28,10 @@ import (
|
||||
"gpu-turnstile/internal/ollama"
|
||||
)
|
||||
|
||||
// captureLimit bounds how much of a /prompt response body is buffered while
|
||||
// looking for prompt_id. The body still passes through to the client
|
||||
// unchanged regardless of size.
|
||||
const captureLimit = 64 * 1024
|
||||
// defaultCaptureLimit bounds how much of a /prompt response body is
|
||||
// buffered while looking for prompt_id. The body still passes through to
|
||||
// the client unchanged regardless of size.
|
||||
const defaultCaptureLimit = 64 * 1024
|
||||
|
||||
// Config wires a Server.
|
||||
type Config struct {
|
||||
@@ -37,16 +44,44 @@ 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
|
||||
WarmModel string
|
||||
|
||||
// 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
|
||||
HistoryPollInterval time.Duration
|
||||
// FreeTimeout and WarmTimeout bound the /free call and the warm-model
|
||||
// reload; zero selects the defaults.
|
||||
FreeTimeout time.Duration
|
||||
WarmTimeout time.Duration
|
||||
// PromptCaptureLimit overrides defaultCaptureLimit when > 0.
|
||||
PromptCaptureLimit int64
|
||||
|
||||
WarmModel string
|
||||
}
|
||||
|
||||
// Server serves both gpu-turnstile listeners.
|
||||
type Server struct {
|
||||
cfg Config
|
||||
log *slog.Logger
|
||||
cfg Config
|
||||
log *slog.Logger
|
||||
freeTimeout time.Duration
|
||||
warmTimeout time.Duration
|
||||
captureLimit int64
|
||||
backoffInitial time.Duration
|
||||
backoffMax time.Duration
|
||||
|
||||
ollamaProxy *httputil.ReverseProxy
|
||||
comfyProxy *httputil.ReverseProxy
|
||||
@@ -66,16 +101,160 @@ func New(cfg Config) (*Server, error) {
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
if cfg.UnloadPollInterval > 0 {
|
||||
cfg.Ollama.PollInterval = cfg.UnloadPollInterval
|
||||
}
|
||||
if cfg.HistoryPollInterval > 0 {
|
||||
cfg.Comfy.PollInterval = cfg.HistoryPollInterval
|
||||
}
|
||||
freeTimeout := cfg.FreeTimeout
|
||||
if freeTimeout <= 0 {
|
||||
freeTimeout = 30 * time.Second
|
||||
}
|
||||
warmTimeout := cfg.WarmTimeout
|
||||
if warmTimeout <= 0 {
|
||||
warmTimeout = 2 * time.Minute
|
||||
}
|
||||
captureLimit := cfg.PromptCaptureLimit
|
||||
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,
|
||||
ollamaProxy: newReverseProxy(ollamaURL, log.With("upstream", "ollama")),
|
||||
comfyProxy: newReverseProxy(comfyURL, log.With("upstream", "comfy")),
|
||||
cfg: cfg,
|
||||
log: log,
|
||||
freeTimeout: freeTimeout,
|
||||
warmTimeout: warmTimeout,
|
||||
captureLimit: captureLimit,
|
||||
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()
|
||||
@@ -106,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{
|
||||
@@ -124,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)
|
||||
@@ -151,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
|
||||
@@ -166,15 +429,16 @@ func (s *Server) ComfyHandler() http.Handler {
|
||||
return
|
||||
}
|
||||
s.comfyProxy.ServeHTTP(w, r)
|
||||
})
|
||||
}))
|
||||
}
|
||||
|
||||
// captureWriter passes the response through unchanged while recording the
|
||||
// status code and the first captureLimit bytes of the body.
|
||||
// status code and the first limit bytes of the body.
|
||||
type captureWriter struct {
|
||||
http.ResponseWriter
|
||||
status int
|
||||
buf bytes.Buffer
|
||||
limit int64
|
||||
}
|
||||
|
||||
func (w *captureWriter) WriteHeader(code int) {
|
||||
@@ -183,7 +447,7 @@ func (w *captureWriter) WriteHeader(code int) {
|
||||
}
|
||||
|
||||
func (w *captureWriter) Write(p []byte) (int, error) {
|
||||
if w.buf.Len() < captureLimit {
|
||||
if int64(w.buf.Len()) < w.limit {
|
||||
w.buf.Write(p)
|
||||
}
|
||||
return w.ResponseWriter.Write(p)
|
||||
@@ -228,7 +492,7 @@ func (s *Server) handlePrompt(w http.ResponseWriter, r *http.Request) {
|
||||
log.Info("ollama models unloaded", "seconds", elapsed.Seconds())
|
||||
}
|
||||
|
||||
cw := &captureWriter{ResponseWriter: w, status: http.StatusOK}
|
||||
cw := &captureWriter{ResponseWriter: w, status: http.StatusOK, limit: s.captureLimit}
|
||||
s.comfyProxy.ServeHTTP(cw, r)
|
||||
|
||||
var accepted struct {
|
||||
@@ -262,7 +526,7 @@ func (s *Server) finishImageJob(promptID string) {
|
||||
log.Info("image job completed")
|
||||
}
|
||||
|
||||
freeCtx, freeCancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
freeCtx, freeCancel := context.WithTimeout(context.Background(), s.freeTimeout)
|
||||
if err := s.cfg.Comfy.Free(freeCtx); err != nil {
|
||||
log.Warn("failed to free ComfyUI models", "err", err)
|
||||
}
|
||||
@@ -273,7 +537,7 @@ func (s *Server) finishImageJob(promptID string) {
|
||||
|
||||
if s.cfg.WarmModel != "" {
|
||||
if state, _, _ := s.cfg.Lock.Snapshot(); state == lock.StateIdle {
|
||||
wctx, wcancel := context.WithTimeout(context.Background(), 2*time.Minute)
|
||||
wctx, wcancel := context.WithTimeout(context.Background(), s.warmTimeout)
|
||||
if err := s.cfg.Ollama.Warm(wctx, s.cfg.WarmModel); err != nil {
|
||||
log.Warn("warm model reload failed", "model", s.cfg.WarmModel, "err", err)
|
||||
} else {
|
||||
|
||||
@@ -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 }
|
||||
Reference in New Issue
Block a user