[query] Move trace handler to server level

cmd/query/app/http_handler.go

@@ -18,7 +18,6 @@ import (
 	"github.com/gogo/protobuf/proto"
 	"github.com/gorilla/mux"
 	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
-	"go.opentelemetry.io/otel/propagation"
 	"go.opentelemetry.io/otel/trace"
 	"go.uber.org/zap"
 
@@ -132,7 +131,6 @@ func (aH *APIHandler) handleFunc(
 ) *mux.Route {
 	route := aH.formatRoute(routeFmt, args...)
 	var handler http.Handler = http.HandlerFunc(f)
-	handler = traceResponseHandler(handler)
 	handler = otelhttp.WithRouteTag(route, handler)
 	handler = spanNameHandler(route, handler)
 	return router.HandleFunc(route, handler.ServeHTTP)
@@ -517,21 +515,6 @@ func (aH *APIHandler) writeJSON(w http.ResponseWriter, r *http.Request, response
 	}
 }
 
-// Returns a handler that generates a traceresponse header.
-// https://github.com/w3c/trace-context/blob/main/spec/21-http_response_header_format.md
-func traceResponseHandler(handler http.Handler) http.Handler {
-	// We use the standard TraceContext propagator, since the formats are identical.
-	// But the propagator uses "traceparent" header name, so we inject it into a map
-	// `carrier` and then use the result to set the "tracereponse" header.
-	var prop propagation.TraceContext
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		carrier := make(map[string]string)
-		prop.Inject(r.Context(), propagation.MapCarrier(carrier))
-		w.Header().Add("traceresponse", carrier["traceparent"])
-		handler.ServeHTTP(w, r)
-	})
-}
-
 func spanNameHandler(spanName string, handler http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		span := trace.SpanFromContext(r.Context())

cmd/query/app/server.go

@@ -201,6 +201,7 @@ func initRouter(
 	if tenancyMgr.Enabled {
 		handler = tenancy.ExtractTenantHTTPHandler(tenancyMgr, handler)
 	}
+	handler = traceResponseHandler(handler)
 	return handler, staticHandlerCloser
 }
 

cmd/query/app/trace_response_handler.go

@@ -0,0 +1,25 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package app
+
+import (
+	"net/http"
+
+	"go.opentelemetry.io/otel/propagation"
+)
+
+// Returns a handler that generates a traceresponse header.
+// https://github.com/w3c/trace-context/blob/main/spec/21-http_response_header_format.md
+func traceResponseHandler(handler http.Handler) http.Handler {
+	// We use the standard TraceContext propagator, since the formats are identical.
+	// But the propagator uses "traceparent" header name, so we inject it into a map
+	// `carrier` and then use the result to set the "tracereponse" header.
+	var prop propagation.TraceContext
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		carrier := make(map[string]string)
+		prop.Inject(r.Context(), propagation.MapCarrier(carrier))
+		w.Header().Add("traceresponse", carrier["traceparent"])
+		handler.ServeHTTP(w, r)
+	})
+}

cmd/query/app/trace_response_handler_test.go

@@ -0,0 +1,49 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package app
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/sdk/trace/tracetest"
+
+	"github.com/jaegertracing/jaeger/pkg/jtracer"
+)
+
+func TestTraceResponseHandler(t *testing.T) {
+	exporter := tracetest.NewInMemoryExporter()
+	tracerProvider := sdktrace.NewTracerProvider(
+		sdktrace.WithSyncer(exporter),
+		sdktrace.WithSampler(sdktrace.AlwaysSample()),
+	)
+	jTracer := jtracer.JTracer{OTEL: tracerProvider}
+	tracer := jTracer.OTEL.Tracer("instrumentation/trace_response_handler_test")
+
+	emptyHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
+	handlerWithTracing := traceResponseHandler(emptyHandler)
+	hanlderWithInstrumentation := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_, span := tracer.Start(r.Context(), "operation")
+		defer span.End()
+
+		handlerWithTracing.ServeHTTP(w, r)
+	})
+
+	server := httptest.NewServer(hanlderWithInstrumentation)
+	defer server.Close()
+
+	req, err := http.NewRequest(http.MethodGet, server.URL, nil)
+	require.NoError(t, err)
+
+	resp, err := server.Client().Do(req)
+	require.NoError(t, err)
+
+	traceResponse := resp.Header.Get("traceresponse")
+
+	// TODO: why isn't this populated?
+	require.Equal(t, "", traceResponse)
+}