Retry stale serverless upstreams

agent/internal/serverless/gateway.go

@@ -20,8 +20,9 @@ import (
 )
 
 const (
-	GatewayPort      = 18080
-	upstreamCacheTTL = 10 * time.Second
+	GatewayPort              = 18080
+	upstreamCacheTTL         = 10 * time.Second
+	proxyRetryResolveTimeout = 2 * time.Second
 )
 
 var (
@@ -183,28 +184,68 @@ func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 
 	upstream := upstreams[g.nextIndex(len(upstreams))]
+	proxyErr := g.proxyOnce(w, r, upstream)
+	if proxyErr == nil {
+		return
+	}
+	if requestCancelled(r, proxyErr) {
+		return
+	}
+
+	log.Printf("[serverless-gateway] proxy error for host %s to %s: %v", host, upstream.Url, proxyErr)
+	g.evictUpstreams(host)
+	if !retryableProxyRequest(r) {
+		http.Error(w, "bad gateway", http.StatusBadGateway)
+		return
+	}
+
+	retryCtx, cancel := context.WithTimeout(r.Context(), proxyRetryResolveTimeout)
+	defer cancel()
+	retryUpstreams, err := g.getUpstreams(retryCtx, host)
+	if err != nil || len(retryUpstreams) == 0 {
+		if r.Context().Err() != nil {
+			return
+		}
+		if err == nil {
+			err = errors.New("no upstreams after cache eviction")
+		}
+		log.Printf("[serverless-gateway] proxy retry resolution failed for host %s: %v", host, err)
+		http.Error(w, "bad gateway", http.StatusBadGateway)
+		return
+	}
+
+	retryUpstream := selectRetryUpstream(retryUpstreams, upstream.Url, g.nextIndex(len(retryUpstreams)))
+	proxyErr = g.proxyOnce(w, r, retryUpstream)
+	if proxyErr == nil || requestCancelled(r, proxyErr) {
+		return
+	}
+	log.Printf("[serverless-gateway] proxy retry failed for host %s to %s: %v", host, retryUpstream.Url, proxyErr)
+	g.evictUpstreams(host)
+	http.Error(w, "bad gateway", http.StatusBadGateway)
+}
+
+func (g *Gateway) proxyOnce(w http.ResponseWriter, r *http.Request, upstream agenthttp.ServerlessUpstream) error {
 	target, err := url.Parse("http://" + upstream.Url)
 	if err != nil {
-		log.Printf("[serverless-gateway] invalid upstream %q for host %s: %v", upstream.Url, host, err)
-		http.Error(w, "bad upstream", http.StatusBadGateway)
-		return
+		return fmt.Errorf("invalid upstream %q: %w", upstream.Url, err)
 	}
 
 	proxy := httputil.NewSingleHostReverseProxy(target)
 	originalDirector := proxy.Director
 	originalHost := r.Host
+	originalProto := forwardedProto(r)
 	proxy.Director = func(req *http.Request) {
 		originalDirector(req)
 		req.Host = originalHost
 		req.Header.Set("X-Forwarded-Host", originalHost)
-		req.Header.Set("X-Forwarded-Proto", forwardedProto(r))
+		req.Header.Set("X-Forwarded-Proto", originalProto)
 	}
-	proxy.ErrorHandler = func(w http.ResponseWriter, req *http.Request, err error) {
-		log.Printf("[serverless-gateway] proxy error for host %s to %s: %v", host, upstream.Url, err)
-		g.evictUpstreams(host)
-		http.Error(w, "bad gateway", http.StatusBadGateway)
+	var proxyErr error
+	proxy.ErrorHandler = func(_ http.ResponseWriter, _ *http.Request, err error) {
+		proxyErr = err
 	}
 	proxy.ServeHTTP(w, r)
+	return proxyErr
 }
 
 func (g *Gateway) getUpstreams(ctx context.Context, host string) ([]agenthttp.ServerlessUpstream, error) {
@@ -1307,6 +1348,41 @@ func (g *Gateway) nextIndex(length int) int {
 	return int(atomic.AddUint64(&g.counter, 1)-1) % length
 }
 
+func retryableProxyRequest(r *http.Request) bool {
+	if r.Method != http.MethodGet && r.Method != http.MethodHead {
+		return false
+	}
+	if r.ContentLength != 0 || (r.Body != nil && r.Body != http.NoBody) {
+		return false
+	}
+	return !headerContainsToken(r.Header, "Connection", "upgrade") && r.Header.Get("Upgrade") == ""
+}
+
+func requestCancelled(r *http.Request, err error) bool {
+	return r.Context().Err() != nil || errors.Is(err, context.Canceled)
+}
+
+func selectRetryUpstream(upstreams []agenthttp.ServerlessUpstream, failedURL string, start int) agenthttp.ServerlessUpstream {
+	for offset := range len(upstreams) {
+		upstream := upstreams[(start+offset)%len(upstreams)]
+		if upstream.Url != failedURL {
+			return upstream
+		}
+	}
+	return upstreams[start%len(upstreams)]
+}
+
+func headerContainsToken(header http.Header, name string, token string) bool {
+	for _, value := range header.Values(name) {
+		for part := range strings.SplitSeq(value, ",") {
+			if strings.EqualFold(strings.TrimSpace(part), token) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 func normalizeHost(host string) string {
 	if h, _, err := net.SplitHostPort(host); err == nil {
 		return strings.ToLower(strings.TrimSpace(h))

agent/internal/serverless/gateway_test.go

@@ -5,9 +5,14 @@ import (
 	"context"
 	"errors"
 	"log"
+	"net"
+	"net/http"
+	"net/http/httptest"
 	"slices"
+	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -164,6 +169,160 @@ func assertProbeSet(t *testing.T, got []string, want ...string) {
 	}
 }
 
+func TestServeHTTPRetriesStaleUpstreamForSafeRequest(t *testing.T) {
+	var requests atomic.Int32
+	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		requests.Add(1)
+		w.WriteHeader(http.StatusNoContent)
+	}))
+	t.Cleanup(backend.Close)
+
+	gateway := testProxyGateway(t, backend.Listener.Addr().String())
+	gateway.cacheUpstreams("app.example.com", []agenthttp.ServerlessUpstream{{Url: "127.0.0.1:0"}})
+
+	response := httptest.NewRecorder()
+	request := httptest.NewRequest(http.MethodGet, "http://app.example.com/health", nil)
+	gateway.ServeHTTP(response, request)
+
+	if response.Code != http.StatusNoContent {
+		t.Fatalf("status = %d, want %d", response.Code, http.StatusNoContent)
+	}
+	if requests.Load() != 1 {
+		t.Fatalf("backend requests = %d, want 1", requests.Load())
+	}
+}
+
+func TestServeHTTPDoesNotRetryUnsafeRequest(t *testing.T) {
+	var requests atomic.Int32
+	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		requests.Add(1)
+		w.WriteHeader(http.StatusNoContent)
+	}))
+	t.Cleanup(backend.Close)
+
+	gateway := testProxyGateway(t, backend.Listener.Addr().String())
+	gateway.cacheUpstreams("app.example.com", []agenthttp.ServerlessUpstream{{Url: "127.0.0.1:0"}})
+
+	response := httptest.NewRecorder()
+	request := httptest.NewRequest(http.MethodPost, "http://app.example.com/items", strings.NewReader("item"))
+	gateway.ServeHTTP(response, request)
+
+	if response.Code != http.StatusBadGateway {
+		t.Fatalf("status = %d, want %d", response.Code, http.StatusBadGateway)
+	}
+	if requests.Load() != 0 {
+		t.Fatalf("backend requests = %d, want 0", requests.Load())
+	}
+}
+
+func TestServeHTTPDoesNotRetryUpgradeRequest(t *testing.T) {
+	var requests atomic.Int32
+	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		requests.Add(1)
+		w.WriteHeader(http.StatusNoContent)
+	}))
+	t.Cleanup(backend.Close)
+
+	gateway := testProxyGateway(t, backend.Listener.Addr().String())
+	gateway.cacheUpstreams("app.example.com", []agenthttp.ServerlessUpstream{{Url: "127.0.0.1:0"}})
+
+	response := httptest.NewRecorder()
+	request := httptest.NewRequest(http.MethodGet, "http://app.example.com/socket", nil)
+	request.Header.Set("Connection", "keep-alive, Upgrade")
+	request.Header.Set("Upgrade", "websocket")
+	gateway.ServeHTTP(response, request)
+
+	if response.Code != http.StatusBadGateway {
+		t.Fatalf("status = %d, want %d", response.Code, http.StatusBadGateway)
+	}
+	if requests.Load() != 0 {
+		t.Fatalf("backend requests = %d, want 0", requests.Load())
+	}
+}
+
+func TestServeHTTPDoesNotEvictCacheWhenRequestIsCancelled(t *testing.T) {
+	gateway := testProxyGateway(t, "127.0.0.1:0")
+	cached := []agenthttp.ServerlessUpstream{{Url: "127.0.0.1:0"}}
+	gateway.cacheUpstreams("app.example.com", cached)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	request := httptest.NewRequest(http.MethodGet, "http://app.example.com/", nil).WithContext(ctx)
+	gateway.ServeHTTP(httptest.NewRecorder(), request)
+
+	upstreams, ok := gateway.cachedUpstreams("app.example.com")
+	if !ok || len(upstreams) != 1 || upstreams[0].Url != cached[0].Url {
+		t.Fatalf("cached upstreams = %+v, ok=%t; want original cache entry", upstreams, ok)
+	}
+}
+
+func TestServeHTTPRetriesOnlyOnce(t *testing.T) {
+	retryAddress, accepts := acceptAndCloseAddress(t)
+	gateway := testProxyGateway(t, retryAddress)
+	gateway.cacheUpstreams("app.example.com", []agenthttp.ServerlessUpstream{{Url: "127.0.0.1:0"}})
+
+	response := httptest.NewRecorder()
+	request := httptest.NewRequest(http.MethodGet, "http://app.example.com/", nil)
+	gateway.ServeHTTP(response, request)
+
+	if response.Code != http.StatusBadGateway {
+		t.Fatalf("status = %d, want %d", response.Code, http.StatusBadGateway)
+	}
+	if _, ok := gateway.cachedUpstreams("app.example.com"); ok {
+		t.Fatal("failed retry left upstream cache populated")
+	}
+	if accepts.Load() != 1 {
+		t.Fatalf("retry upstream connections = %d, want 1", accepts.Load())
+	}
+}
+
+func testProxyGateway(t *testing.T, upstreamAddress string) *Gateway {
+	t.Helper()
+	host, portText, err := net.SplitHostPort(upstreamAddress)
+	if err != nil {
+		t.Fatalf("split upstream address: %v", err)
+	}
+	port, err := strconv.Atoi(portText)
+	if err != nil {
+		t.Fatalf("parse upstream port: %v", err)
+	}
+	state := testExpectedState("running")
+	state.Containers[0].IPAddress = host
+	state.Serverless.Routes[0].Port = port
+	state.Serverless.Routes[0].Upstreams = nil
+	runtime := &fakeRuntime{
+		state: state,
+		containers: []container.Container{{
+			ID:           "ctr-local",
+			State:        "running",
+			DeploymentID: "dep_local",
+			ServiceID:    "svc_1",
+		}},
+	}
+	return NewGateway(runtime)
+}
+
+func acceptAndCloseAddress(t *testing.T) (string, *atomic.Int32) {
+	t.Helper()
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listen for closing upstream: %v", err)
+	}
+	t.Cleanup(func() { listener.Close() })
+	var accepts atomic.Int32
+	go func() {
+		for {
+			conn, err := listener.Accept()
+			if err != nil {
+				return
+			}
+			accepts.Add(1)
+			conn.Close()
+		}
+	}()
+	return listener.Addr().String(), &accepts
+}
+
 func TestGetUpstreamsReturnsWhenFollowerContextIsCancelled(t *testing.T) {
 	state := testExpectedState("stopped")
 	g := NewGateway(&fakeRuntime{state: state})