gpu-turnstile/internal/proxy/retry_test.go

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 }