otelgrpc: Add SpanStartOptions behavior to stats_handlers

open-telemetry · Apr 26, 2024 · 0f1a838 · 0f1a838
1 parent fed6e67
commit 0f1a838
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 5 deletions.
diff --git a/instrumentation/google.golang.org/grpc/otelgrpc/stats_handler.go b/instrumentation/google.golang.org/grpc/otelgrpc/stats_handler.go
@@ -57,11 +57,18 @@ func (h *serverHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) cont
 
 	name, attrs := internal.ParseFullMethod(info.FullMethodName)
 	attrs = append(attrs, RPCSystemGRPC)
+
+	startOpts := append([]trace.SpanStartOption{
+		trace.WithSpanKind(trace.SpanKindServer),
+		trace.WithAttributes(attrs...),
+	},
+		h.config.SpanStartOptions...,
+	)
+
 	ctx, _ = h.tracer.Start(
 		trace.ContextWithRemoteSpanContext(ctx, trace.SpanContextFromContext(ctx)),
 		name,
-		trace.WithSpanKind(trace.SpanKindServer),
-		trace.WithAttributes(attrs...),
+		startOpts...,
 	)
 
 	gctx := gRPCContext{
@@ -93,11 +100,18 @@ func NewClientHandler(opts ...Option) stats.Handler {
 func (h *clientHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
 	name, attrs := internal.ParseFullMethod(info.FullMethodName)
 	attrs = append(attrs, RPCSystemGRPC)
+
+	startOpts := append([]trace.SpanStartOption{
+		trace.WithSpanKind(trace.SpanKindClient),
+		trace.WithAttributes(attrs...),
+	},
+		h.config.SpanStartOptions...,
+	)
+
 	ctx, _ = h.tracer.Start(
 		ctx,
 		name,
-		trace.WithSpanKind(trace.SpanKindClient),
-		trace.WithAttributes(attrs...),
+		startOpts...,
 	)
 
 	gctx := gRPCContext{

diff --git a/instrumentation/google.golang.org/grpc/otelgrpc/test/stats_handler_test.go b/instrumentation/google.golang.org/grpc/otelgrpc/test/stats_handler_test.go
@@ -21,6 +21,11 @@ import (
 	"go.opentelemetry.io/otel/sdk/trace"
 	"go.opentelemetry.io/otel/sdk/trace/tracetest"
 	semconv "go.opentelemetry.io/otel/semconv/v1.17.0"
+	oteltrace "go.opentelemetry.io/otel/trace"
+)
+
+const (
+	serviceName = "TestGrpcService"
 )
 
 func TestStatsHandlerHandleRPCServerErrors(t *testing.T) {
@@ -38,7 +43,6 @@ func TestStatsHandlerHandleRPCServerErrors(t *testing.T) {
 				otelgrpc.WithMeterProvider(mp),
 			)
 
-			serviceName := "TestGrpcService"
 			methodName := serviceName + "/" + name
 			fullMethodName := "/" + methodName
 			// call the server handler
@@ -62,6 +66,64 @@ func TestStatsHandlerHandleRPCServerErrors(t *testing.T) {
 	}
 }
 
+func TestStatsHandlerServerWithSpanOptions(t *testing.T) {
+	sr := tracetest.NewSpanRecorder()
+	tp := trace.NewTracerProvider(trace.WithSpanProcessor(sr))
+	serverHandler := otelgrpc.NewServerHandler(
+		otelgrpc.WithTracerProvider(tp),
+		otelgrpc.WithSpanOptions(oteltrace.WithAttributes(attribute.Bool("custom", true))),
+	)
+
+	methodName := serviceName + "/" + "test"
+	fullMethodName := "/" + methodName
+	// call the server handler
+	ctx := serverHandler.TagRPC(context.Background(), &stats.RPCTagInfo{
+		FullMethodName: fullMethodName,
+	})
+
+	serverHandler.HandleRPC(ctx, &stats.End{})
+
+	expected := []attribute.KeyValue{
+		semconv.RPCSystemGRPC,
+		semconv.RPCService("TestGrpcService"),
+		semconv.RPCMethod("test"),
+		otelgrpc.GRPCStatusCodeKey.Int64(0),
+		attribute.Bool("custom", true),
+	}
+	span, ok := getSpanFromRecorder(sr, methodName)
+	require.True(t, ok, "missing span %q", methodName)
+	assert.ElementsMatch(t, expected, span.Attributes())
+}
+
+func TestStatsHandlerClientWithSpanOptions(t *testing.T) {
+	sr := tracetest.NewSpanRecorder()
+	tp := trace.NewTracerProvider(trace.WithSpanProcessor(sr))
+	clientHandler := otelgrpc.NewClientHandler(
+		otelgrpc.WithTracerProvider(tp),
+		otelgrpc.WithSpanOptions(oteltrace.WithAttributes(attribute.Bool("custom", true))),
+	)
+
+	methodName := serviceName + "/" + "test"
+	fullMethodName := "/" + methodName
+	// call the client handler
+	ctx := clientHandler.TagRPC(context.Background(), &stats.RPCTagInfo{
+		FullMethodName: fullMethodName,
+	})
+
+	clientHandler.HandleRPC(ctx, &stats.End{})
+
+	expected := []attribute.KeyValue{
+		semconv.RPCSystemGRPC,
+		semconv.RPCService("TestGrpcService"),
+		semconv.RPCMethod("test"),
+		otelgrpc.GRPCStatusCodeKey.Int64(0),
+		attribute.Bool("custom", true),
+	}
+	span, ok := getSpanFromRecorder(sr, methodName)
+	require.True(t, ok, "missing span %q", methodName)
+	assert.ElementsMatch(t, expected, span.Attributes())
+}
+
 func assertStatsHandlerServerMetrics(t *testing.T, reader metric.Reader, serviceName, name string, code grpc_codes.Code) {
 	want := metricdata.ScopeMetrics{
 		Scope: wantInstrumentationScope,