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.
This commit is contained in:
mram
+93
-13
@@ -13,6 +13,7 @@ import (
|
||||
"net/http"
|
||||
"net/http/httputil"
|
||||
"net/url"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"gpu-turnstile/internal/comfy"
|
||||
@@ -41,6 +42,13 @@ type Config struct {
|
||||
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
|
||||
@@ -57,11 +65,13 @@ type Config struct {
|
||||
|
||||
// Server serves both gpu-turnstile listeners.
|
||||
type Server struct {
|
||||
cfg Config
|
||||
log *slog.Logger
|
||||
freeTimeout time.Duration
|
||||
warmTimeout time.Duration
|
||||
captureLimit int64
|
||||
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
|
||||
@@ -99,19 +109,89 @@ 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")),
|
||||
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 that fail with "connection refused" —
|
||||
// typically the upstream service simply not being up yet — using an
|
||||
// exponential backoff: initial, doubling each attempt, capped at max. The
|
||||
// retry loop runs until the request succeeds, fails with a different
|
||||
// error, or the client's context is cancelled. Retrying a refused
|
||||
// connection is always safe: no request bytes were ever sent.
|
||||
type retryTransport struct {
|
||||
base http.RoundTripper
|
||||
initial time.Duration
|
||||
max time.Duration
|
||||
log *slog.Logger
|
||||
}
|
||||
|
||||
// wsaECONNREFUSED is the Windows error code for a refused connection.
|
||||
// syscall.ECONNREFUSED is the POSIX value and never matches on Windows,
|
||||
// where dial errors carry the WSA code.
|
||||
const wsaECONNREFUSED syscall.Errno = 10061
|
||||
|
||||
func isConnRefused(err error) bool {
|
||||
if errors.Is(err, syscall.ECONNREFUSED) {
|
||||
return true
|
||||
}
|
||||
var errno syscall.Errno
|
||||
return errors.As(err, &errno) && errno == wsaECONNREFUSED
|
||||
}
|
||||
|
||||
func (t *retryTransport) logWarn(msg string, args ...any) {
|
||||
if t.log != nil {
|
||||
t.log.Warn(msg, args...)
|
||||
}
|
||||
}
|
||||
|
||||
func (t *retryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
wait := t.initial
|
||||
for {
|
||||
resp, err := t.base.RoundTrip(req)
|
||||
if err == nil || !isConnRefused(err) {
|
||||
return resp, err
|
||||
}
|
||||
t.logWarn("upstream connection refused; backing off",
|
||||
"path", req.URL.Path, "retry_in", wait)
|
||||
select {
|
||||
case <-req.Context().Done():
|
||||
return nil, 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()
|
||||
|
||||
@@ -0,0 +1,122 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"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 || !isConnRefused(err) {
|
||||
t.Fatalf("err = %v, want connection refused", err)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user