refactor: extract endpoint resolution into dedicated middleware (#1558)

Author: linkvt

interceptor/forward_wait_func.go

@@ -4,13 +4,15 @@ import (
 	"errors"
 	"net/http"
 	"net/http/httputil"
+	"time"
 
-	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
 	"go.opentelemetry.io/otel/attribute"
-	"go.opentelemetry.io/otel/propagation"
 	"go.opentelemetry.io/otel/trace"
 
 	"github.com/kedacore/http-add-on/interceptor/config"
+	kedahttp "github.com/kedacore/http-add-on/pkg/http"
+	"github.com/kedacore/http-add-on/pkg/k8s"
 	"github.com/kedacore/http-add-on/pkg/util"
 )
 
@@ -21,16 +23,16 @@ var (
 )
 
 type Upstream struct {
-	roundTripper   http.RoundTripper
-	tracingCfg     config.Tracing
-	shouldFallback bool
+	transportPool     *kedahttp.TransportPool
+	tracingCfg        config.Tracing
+	respHeaderTimeout time.Duration
 }
 
-func NewUpstream(roundTripper http.RoundTripper, tracingCfg config.Tracing, shouldFallback bool) *Upstream {
+func NewUpstream(baseTransport *http.Transport, tracingCfg config.Tracing, respHeaderTimeout time.Duration) *Upstream {
 	return &Upstream{
-		roundTripper:   roundTripper,
-		tracingCfg:     tracingCfg,
-		shouldFallback: shouldFallback,
+		transportPool:     kedahttp.NewTransportPool(baseTransport),
+		tracingCfg:        tracingCfg,
+		respHeaderTimeout: respHeaderTimeout,
 	}
 }
 
@@ -41,24 +43,14 @@ func (uh *Upstream) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 
 	if uh.tracingCfg.Enabled {
-		p := otel.GetTextMapPropagator()
-		ctx = p.Extract(ctx, propagation.HeaderCarrier(r.Header))
-
-		p.Inject(ctx, propagation.HeaderCarrier(w.Header()))
-
 		span := trace.SpanFromContext(ctx)
-		defer span.End()
-
-		serviceValAttr := attribute.String("service", "keda-http-interceptor-proxy-upstream")
-		coldStartValAttr := attribute.String("cold-start", w.Header().Get("X-KEDA-HTTP-Cold-Start"))
-
-		span.SetAttributes(serviceValAttr, coldStartValAttr)
+		span.SetAttributes(
+			attribute.String("service", "keda-http-interceptor-proxy-upstream"),
+			attribute.String("cold-start", w.Header().Get(kedahttp.HeaderColdStart)),
+		)
 	}
 
 	url := util.UpstreamURLFromContext(ctx)
-	if uh.shouldFallback {
-		url = util.FallbackURLFromContext(ctx)
-	}
 
 	if url == nil {
 		sh := NewStatic(http.StatusInternalServerError, errNilUpstreamURL)
@@ -67,6 +59,27 @@ func (uh *Upstream) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	respHeaderTimeout := uh.respHeaderTimeout
+	// TODO(v1): remove timeout compatibility fallback for HTTPSO before v1 release
+	if ir := util.InterceptorRouteFromContext(ctx); ir != nil {
+		if v, ok := ir.Annotations[k8s.AnnotationResponseHeaderTimeout]; ok {
+			if d, err := time.ParseDuration(v); err == nil && d > 0 {
+				respHeaderTimeout = d
+			}
+		}
+	}
+	transport := uh.transportPool.Get(respHeaderTimeout)
+
+	var rt http.RoundTripper = transport
+	if uh.tracingCfg.Enabled {
+		rt = otelhttp.NewTransport(transport)
+	}
+
+	rc := http.NewResponseController(w)
+	if err := rc.EnableFullDuplex(); err != nil {
+		util.LoggerFromContext(ctx).Error(err, "could not enable full duplex on response writer, continuing")
+	}
+
 	proxy := &httputil.ReverseProxy{
 		Rewrite: func(pr *httputil.ProxyRequest) {
 			pr.SetURL(url)
@@ -84,7 +97,7 @@ func (uh *Upstream) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			}
 		},
 		BufferPool: bufferPool,
-		Transport:  uh.roundTripper,
+		Transport:  rt,
 		ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
 			sh := NewStatic(http.StatusBadGateway, err)
 			sh.ServeHTTP(w, r)

interceptor/handler/upstream.go

@@ -55,7 +55,7 @@ func TestW3CPropagation(t *testing.T) {
 	receivedRequest := microservice.IncomingRequests()[0]
 	receivedHeaders := receivedRequest.Header
 
-	r.Equal(receivedHeaders.Get("Traceparent"), traceParent)
+	r.Contains(receivedHeaders.Get("Traceparent"), fullW3CLengthTraceID)
 
 	r.NotContains(receivedHeaders, "B3")
 	r.NotContains(receivedHeaders, "b3")
@@ -67,9 +67,7 @@ func TestW3CPropagation(t *testing.T) {
 	_ = tracerProvider.ForceFlush(request.Context())
 
 	exportedSpans := exporter.GetSpans()
-	if len(exportedSpans) != 1 {
-		t.Fatalf("Expected 1 Span, got %d", len(exportedSpans))
-	}
+	r.GreaterOrEqual(len(exportedSpans), 1, "expected at least 1 span")
 	sc := exportedSpans[0].SpanContext
 	r.Equal(fullW3CLengthTraceID, sc.TraceID().String())
 	r.True(sc.IsSampled())
@@ -99,7 +97,7 @@ func TestPropagationWhenNoHeaders(t *testing.T) {
 	receivedRequest := microservice.IncomingRequests()[0]
 	receivedHeaders := receivedRequest.Header
 
-	r.NotContains(receivedHeaders, "Traceparent")
+	r.Contains(receivedHeaders, "Traceparent")
 	r.NotContains(receivedHeaders, "B3")
 	r.NotContains(receivedHeaders, "b3")
 	r.NotContains(receivedHeaders, "X-B3-Parentspanid")
@@ -110,22 +108,21 @@ func TestPropagationWhenNoHeaders(t *testing.T) {
 	_ = tracerProvider.ForceFlush(request.Context())
 
 	exportedSpans := exporter.GetSpans()
-	if len(exportedSpans) != 1 {
-		t.Fatalf("Expected 1 Span, got %d", len(exportedSpans))
-	}
+	r.GreaterOrEqual(len(exportedSpans), 1, "expected at least 1 span")
 	sc := exportedSpans[0].SpanContext
 	r.NotEmpty(sc.SpanID())
 	r.NotEmpty(sc.TraceID())
 
 	hasServiceAttribute := false
 	hasColdStartAttribute := false
-	for _, attribute := range exportedSpans[0].Attributes {
-		if attribute.Key == "service" && attribute.Value.AsString() == "keda-http-interceptor-proxy-upstream" {
-			hasServiceAttribute = true
-		}
-
-		if attribute.Key == "cold-start" {
-			hasColdStartAttribute = true
+	for _, s := range exportedSpans {
+		for _, attribute := range s.Attributes {
+			if attribute.Key == "service" && attribute.Value.AsString() == "keda-http-interceptor-proxy-upstream" {
+				hasServiceAttribute = true
+			}
+			if attribute.Key == "cold-start" {
+				hasColdStartAttribute = true
+			}
 		}
 	}
 	r.True(hasServiceAttribute)
@@ -158,8 +155,7 @@ func TestForwarderSuccess(t *testing.T) {
 	req = util.RequestWithUpstreamURL(req, forwardURL)
 	timeouts := defaultTimeouts()
 	dialCtxFunc := retryDialContextFunc(timeouts)
-	rt := newRoundTripper(dialCtxFunc, timeouts.ResponseHeader)
-	uh := NewUpstream(rt, config.Tracing{}, false)
+	uh := NewUpstream(newTestTransport(dialCtxFunc), config.Tracing{}, timeouts.ResponseHeader)
 	uh.ServeHTTP(res, req)
 
 	r.True(
@@ -202,8 +198,7 @@ func TestForwarderHeaderTimeout(t *testing.T) {
 	res, req, err := reqAndRes("/testfwd")
 	r.NoError(err)
 	req = util.RequestWithUpstreamURL(req, originURL)
-	rt := newRoundTripper(dialCtxFunc, timeouts.ResponseHeader)
-	uh := NewUpstream(rt, config.Tracing{}, false)
+	uh := NewUpstream(newTestTransport(dialCtxFunc), config.Tracing{}, timeouts.ResponseHeader)
 	uh.ServeHTTP(res, req)
 
 	forwardedRequests := hdl.IncomingRequests()
@@ -252,8 +247,7 @@ func TestForwarderWaitsForSlowOrigin(t *testing.T) {
 	res, req, err := reqAndRes(path)
 	r.NoError(err)
 	req = util.RequestWithUpstreamURL(req, originURL)
-	rt := newRoundTripper(dialCtxFunc, timeouts.ResponseHeader)
-	uh := NewUpstream(rt, config.Tracing{}, false)
+	uh := NewUpstream(newTestTransport(dialCtxFunc), config.Tracing{}, timeouts.ResponseHeader)
 	uh.ServeHTTP(res, req)
 	// wait for the goroutine above to finish, with a little cusion
 	ensureSignalBeforeTimeout(originWaitCh, originDelay*2)
@@ -272,8 +266,7 @@ func TestForwarderConnectionRetryAndTimeout(t *testing.T) {
 	res, req, err := reqAndRes("/test")
 	r.NoError(err)
 	req = util.RequestWithUpstreamURL(req, noSuchURL)
-	rt := newRoundTripper(dialCtxFunc, timeouts.ResponseHeader)
-	uh := NewUpstream(rt, config.Tracing{}, false)
+	uh := NewUpstream(newTestTransport(dialCtxFunc), config.Tracing{}, timeouts.ResponseHeader)
 
 	start := time.Now()
 	uh.ServeHTTP(res, req)
@@ -321,8 +314,7 @@ func TestForwardRequestRedirectAndHeaders(t *testing.T) {
 	res, req, err := reqAndRes("/testfwd")
 	r.NoError(err)
 	req = util.RequestWithUpstreamURL(req, srvURL)
-	rt := newRoundTripper(dialCtxFunc, timeouts.ResponseHeader)
-	uh := NewUpstream(rt, config.Tracing{}, false)
+	uh := NewUpstream(newTestTransport(dialCtxFunc), config.Tracing{}, timeouts.ResponseHeader)
 	uh.ServeHTTP(res, req)
 	r.Equal(301, res.Code)
 	r.Equal("abc123.com", res.Header().Get("Location"))
@@ -379,7 +371,7 @@ func TestUpstreamPreservesXForwardedHeaders(t *testing.T) {
 			}
 
 			// Configure the Upstream and send a dummy request
-			upstream := NewUpstream(http.DefaultTransport, config.Tracing{}, false)
+			upstream := NewUpstream(http.DefaultTransport.(*http.Transport), config.Tracing{}, 500*time.Millisecond)
 
 			req := httptest.NewRequest("GET", "/test", nil)
 			if tt.forwardedFor != "" {
@@ -435,13 +427,9 @@ func TestUpstreamPreservesXForwardedHeaders(t *testing.T) {
 	}
 }
 
-func newRoundTripper(
-	dialCtxFunc kedanet.DialContextFunc,
-	httpRespHeaderTimeout time.Duration,
-) http.RoundTripper {
+func newTestTransport(dialCtxFunc kedanet.DialContextFunc) *http.Transport {
 	return &http.Transport{
-		DialContext:           dialCtxFunc,
-		ResponseHeaderTimeout: httpRespHeaderTimeout,
+		DialContext: dialCtxFunc,
 	}
 }
 
@@ -485,8 +473,8 @@ func ensureSignalBeforeTimeout(signalCh <-chan struct{}, timeout time.Duration)
 func serveHTTP(w http.ResponseWriter, r *http.Request) {
 	timeouts := defaultTimeouts()
 	dialCtxFunc := retryDialContextFunc(timeouts)
-	rt := newRoundTripper(dialCtxFunc, timeouts.ResponseHeader)
-	upstream := NewUpstream(rt, config.Tracing{Enabled: true}, false)
+	transport := newTestTransport(dialCtxFunc)
+	upstream := NewUpstream(transport, config.Tracing{Enabled: true}, timeouts.ResponseHeader)
 
 	upstream.ServeHTTP(w, r)
 }

interceptor/handler/upstream_test.go

@@ -101,8 +101,6 @@ func main() {
 		runtime.Goexit()
 	}
 
-	waitFunc := newWorkloadReplicasForwardWaitFunc(readyCache)
-
 	queues := queue.NewMemory()
 	routingTable := routing.NewTable(ctrlCache, queues)
 
@@ -218,7 +216,7 @@ func main() {
 
 			setupLog.Info("starting the proxy server with TLS enabled", "port", proxyTLSPort)
 
-			if err := runProxyServer(ctx, ctrl.Log, queues, waitFunc, routingTable, ctrlCache, timeoutCfg, servingCfg, proxyTLSPort, tlsCfg, tracingCfg); !util.IsIgnoredErr(err) {
+			if err := runProxyServer(ctx, ctrl.Log, queues, readyCache, routingTable, ctrlCache, timeoutCfg, servingCfg, proxyTLSPort, tlsCfg, tracingCfg); !util.IsIgnoredErr(err) {
 				setupLog.Error(err, "tls proxy server failed")
 				return err
 			}
@@ -230,7 +228,7 @@ func main() {
 	eg.Go(func() error {
 		setupLog.Info("starting the proxy server with TLS disabled", "port", servingCfg.ProxyPort)
 
-		if err := runProxyServer(ctx, ctrl.Log, queues, waitFunc, routingTable, ctrlCache, timeoutCfg, servingCfg, servingCfg.ProxyPort, nil, tracingCfg); !util.IsIgnoredErr(err) {
+		if err := runProxyServer(ctx, ctrl.Log, queues, readyCache, routingTable, ctrlCache, timeoutCfg, servingCfg, servingCfg.ProxyPort, nil, tracingCfg); !util.IsIgnoredErr(err) {
 			setupLog.Error(err, "proxy server failed")
 			return err
 		}
@@ -292,7 +290,7 @@ func runProxyServer(
 	ctx context.Context,
 	logger logr.Logger,
 	q queue.Counter,
-	waitFunc forwardWaitFunc,
+	readyCache *k8s.ReadyEndpointsCache,
 	routingTable routing.Table,
 	reader client.Reader,
 	timeouts config.Timeouts,
@@ -305,7 +303,7 @@ func runProxyServer(
 	rootHandler := BuildProxyHandler(&ProxyHandlerConfig{
 		Logger:       logger,
 		Queue:        q,
-		WaitFunc:     waitFunc,
+		ReadyCache:   readyCache,
 		RoutingTable: routingTable,
 		Reader:       reader,
 		Timeouts:     timeouts,

interceptor/main.go

@@ -10,14 +10,14 @@ import (
 )
 
 type Counting struct {
-	queueCounter    queue.Counter
-	upstreamHandler http.Handler
+	next         http.Handler
+	queueCounter queue.Counter
 }
 
-func NewCountingMiddleware(queueCounter queue.Counter, upstreamHandler http.Handler) *Counting {
+func NewCountingMiddleware(next http.Handler, queueCounter queue.Counter) *Counting {
 	return &Counting{
-		queueCounter:    queueCounter,
-		upstreamHandler: upstreamHandler,
+		next:         next,
+		queueCounter: queueCounter,
 	}
 }
 
@@ -32,7 +32,7 @@ func (cm *Counting) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	if err := cm.queueCounter.Increase(key, 1); err != nil {
 		util.LoggerFromContext(ctx).Error(err, "error incrementing queue counter", "key", key)
-		cm.upstreamHandler.ServeHTTP(w, r)
+		cm.next.ServeHTTP(w, r)
 		return
 	}
 	metrics.RecordPendingRequestCount(key, 1)
@@ -44,5 +44,5 @@ func (cm *Counting) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		metrics.RecordPendingRequestCount(key, -1)
 	}()
 
-	cm.upstreamHandler.ServeHTTP(w, r)
+	cm.next.ServeHTTP(w, r)
 }

interceptor/middleware/counting.go

@@ -39,7 +39,6 @@ func TestCountMiddleware(t *testing.T) {
 
 	var concurrencyDuringRequest int
 	middleware := NewCountingMiddleware(
-		queueCounter,
 		http.HandlerFunc(func(wr http.ResponseWriter, _ *http.Request) {
 			counts, err := queueCounter.Current()
 			if err == nil {
@@ -48,6 +47,7 @@ func TestCountMiddleware(t *testing.T) {
 			wr.WriteHeader(200)
 			_, _ = wr.Write([]byte("OK"))
 		}),
+		queueCounter,
 	)
 
 	req, err := http.NewRequest("GET", "/something", nil)