mram
20d5439b5f
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.
123 lines
3.3 KiB
Go
123 lines
3.3 KiB
Go
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)
|
|
}
|
|
}
|