Add Span Name Handler

Signed-off-by: Mahad Zaryab <mahadzaryab1@gmail.com>

cmd/query/app/apiv3/http_gateway.go

@@ -68,6 +68,7 @@ func (*HTTPGateway) addRoute(
 ) *mux.Route {
 	var handler http.Handler = http.HandlerFunc(f)
 	handler = otelhttp.WithRouteTag(route, handler)
+	handler = spanNameHandler(route, handler)
 	return router.HandleFunc(route, handler.ServeHTTP)
 }
 
@@ -239,3 +240,11 @@ func (h *HTTPGateway) getOperations(w http.ResponseWriter, r *http.Request) {
 	}
 	h.marshalResponse(&api_v3.GetOperationsResponse{Operations: apiOperations}, w)
 }
+
+func spanNameHandler(spanName string, handler http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		span := trace.SpanFromContext(r.Context())
+		span.SetName(spanName)
+		handler.ServeHTTP(w, r)
+	})
+}

cmd/query/app/http_handler.go

@@ -134,6 +134,7 @@ func (aH *APIHandler) handleFunc(
 	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)
 }
 
@@ -530,3 +531,11 @@ func traceResponseHandler(handler http.Handler) http.Handler {
 		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())
+		span.SetName(spanName)
+		handler.ServeHTTP(w, r)
+	})
+}

cmd/query/app/server_test.go

@@ -4,6 +4,7 @@
 package app
 
 import (
+	"bytes"
 	"context"
 	"crypto/tls"
 	"fmt"
@@ -12,6 +13,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/gogo/protobuf/jsonpb"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -20,6 +22,8 @@ import (
 	"go.opentelemetry.io/collector/config/confighttp"
 	"go.opentelemetry.io/collector/config/confignet"
 	"go.opentelemetry.io/collector/config/configtls"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/sdk/trace/tracetest"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zaptest"
 	"go.uber.org/zap/zaptest/observer"
@@ -855,3 +859,62 @@ func TestServerHTTPTenancy(t *testing.T) {
 		})
 	}
 }
+
+func TestServerHTTP_TracesRequest(t *testing.T) {
+	serverOptions := &QueryOptions{
+		HTTP: confighttp.ServerConfig{
+			Endpoint: ":8080",
+		},
+		GRPC: configgrpc.ServerConfig{
+			NetAddr: confignet.AddrConfig{
+				Endpoint:  ":8081",
+				Transport: confignet.TransportTypeTCP,
+			},
+		},
+	}
+
+	makeMockTrace := func(t *testing.T) *model.Trace {
+		out := new(bytes.Buffer)
+		err := new(jsonpb.Marshaler).Marshal(out, mockTrace)
+		require.NoError(t, err)
+		var trace model.Trace
+		require.NoError(t, jsonpb.Unmarshal(out, &trace))
+		trace.Spans[1].References = []model.SpanRef{
+			{TraceID: model.NewTraceID(0, 0)},
+		}
+		return &trace
+	}
+
+	exporter := tracetest.NewInMemoryExporter()
+	tracerProvider := sdktrace.NewTracerProvider(
+		sdktrace.WithSyncer(exporter),
+		sdktrace.WithSampler(sdktrace.AlwaysSample()),
+	)
+	tracer := jtracer.JTracer{OTEL: tracerProvider}
+
+	tenancyMgr := tenancy.NewManager(&serverOptions.Tenancy)
+	querySvc := makeQuerySvc()
+	querySvc.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), model.NewTraceID(0, 0x123456abc)).
+		Return(makeMockTrace(t), nil).Once()
+	telset := initTelSet(zaptest.NewLogger(t), &tracer, healthcheck.New())
+
+	server, err := NewServer(context.Background(), querySvc.qs,
+		nil, serverOptions, tenancyMgr, telset)
+	require.NoError(t, err)
+	require.NoError(t, server.Start(context.Background()))
+	t.Cleanup(func() {
+		require.NoError(t, server.Close())
+	})
+
+	queryString := "/api/traces/123456aBC"
+	req, err := http.NewRequest(http.MethodGet, "http://localhost:8080"+queryString, nil)
+	require.NoError(t, err)
+
+	client := &http.Client{}
+	resp, err := client.Do(req)
+	require.NoError(t, err)
+
+	assert.Len(t, exporter.GetSpans(), 1, "HTTP request was traced and span reported")
+	assert.Equal(t, "/api/traces/{traceID}", exporter.GetSpans()[0].Name)
+	require.NoError(t, resp.Body.Close())
+}