otelgrpc: add custom attributes to the stats handler

Fixes open-telemetry#3894
inigohu · Jun 5, 2024 · 26b7c36 · 26b7c36
1 parent a594aef
commit 26b7c36
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 13 deletions.
diff --git a/instrumentation/google.golang.org/grpc/otelgrpc/stats_handler.go b/instrumentation/google.golang.org/grpc/otelgrpc/stats_handler.go
@@ -30,6 +30,16 @@ type gRPCContext struct {
 	record           bool
 }
 
+// WithAttributes adds attributes to the given context.
+func WithAttributes(ctx context.Context, attrs ...attribute.KeyValue) context.Context {
+	gctx, ok := ctx.Value(gRPCContextKey{}).(*gRPCContext)
+	if !ok {
+		gctx = &gRPCContext{}
+	}
+	gctx.metricAttrs = append(gctx.metricAttrs, attrs...)
+	return context.WithValue(ctx, gRPCContextKey{}, gctx)
+}
+
 type serverHandler struct {
 	*config
 }
@@ -58,21 +68,25 @@ func (h *serverHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) cont
 
 	name, attrs := internal.ParseFullMethod(info.FullMethodName)
 	attrs = append(attrs, RPCSystemGRPC)
+
+	gctx, ok := ctx.Value(gRPCContextKey{}).(*gRPCContext)
+	if !ok {
+		gctx = &gRPCContext{}
+	}
+	gctx.metricAttrs = append(gctx.metricAttrs, attrs...)
+
 	ctx, _ = h.tracer.Start(
 		trace.ContextWithRemoteSpanContext(ctx, trace.SpanContextFromContext(ctx)),
 		name,
 		trace.WithSpanKind(trace.SpanKindServer),
-		trace.WithAttributes(attrs...),
+		trace.WithAttributes(gctx.metricAttrs...),
 	)
 
-	gctx := gRPCContext{
-		metricAttrs: attrs,
-		record:      true,
-	}
+	gctx.record = true
 	if h.config.Filter != nil {
 		gctx.record = h.config.Filter(info)
 	}
-	return context.WithValue(ctx, gRPCContextKey{}, &gctx)
+	return context.WithValue(ctx, gRPCContextKey{}, gctx)
 }
 
 // HandleRPC processes the RPC stats.
@@ -98,22 +112,26 @@ 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)
+
+	gctx, ok := ctx.Value(gRPCContextKey{}).(*gRPCContext)
+	if !ok {
+		gctx = &gRPCContext{}
+	}
+	gctx.metricAttrs = append(gctx.metricAttrs, attrs...)
+
 	ctx, _ = h.tracer.Start(
 		ctx,
 		name,
 		trace.WithSpanKind(trace.SpanKindClient),
-		trace.WithAttributes(attrs...),
+		trace.WithAttributes(gctx.metricAttrs...),
 	)
 
-	gctx := gRPCContext{
-		metricAttrs: attrs,
-		record:      true,
-	}
+	gctx.record = true
 	if h.config.Filter != nil {
 		gctx.record = h.config.Filter(info)
 	}
 
-	return inject(context.WithValue(ctx, gRPCContextKey{}, &gctx), h.config.Propagators)
+	return inject(context.WithValue(ctx, gRPCContextKey{}, gctx), h.config.Propagators)
 }
 
 // HandleRPC processes the RPC stats.

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
@@ -38,11 +38,12 @@ func TestStatsHandlerHandleRPCServerErrors(t *testing.T) {
 				otelgrpc.WithMeterProvider(mp),
 			)
 
+			ctx := otelgrpc.WithAttributes(context.Background(), attribute.String("custom", "OK"))
 			serviceName := "TestGrpcService"
 			methodName := serviceName + "/" + name
 			fullMethodName := "/" + methodName
 			// call the server handler
-			ctx := serverHandler.TagRPC(context.Background(), &stats.RPCTagInfo{
+			ctx = serverHandler.TagRPC(ctx, &stats.RPCTagInfo{
 				FullMethodName: fullMethodName,
 			})
 
@@ -79,6 +80,7 @@ func assertStatsHandlerServerMetrics(t *testing.T, reader metric.Reader, service
 								semconv.RPCService(serviceName),
 								otelgrpc.RPCSystemGRPC,
 								otelgrpc.GRPCStatusCodeKey.Int64(int64(code)),
+								attribute.String("custom", "OK"),
 							),
 						},
 					},
@@ -97,6 +99,7 @@ func assertStatsHandlerServerMetrics(t *testing.T, reader metric.Reader, service
 								semconv.RPCService(serviceName),
 								otelgrpc.RPCSystemGRPC,
 								otelgrpc.GRPCStatusCodeKey.Int64(int64(code)),
+								attribute.String("custom", "OK"),
 							),
 						},
 					},
@@ -115,6 +118,7 @@ func assertStatsHandlerServerMetrics(t *testing.T, reader metric.Reader, service
 								semconv.RPCService(serviceName),
 								otelgrpc.RPCSystemGRPC,
 								otelgrpc.GRPCStatusCodeKey.Int64(int64(code)),
+								attribute.String("custom", "OK"),
 							),
 						},
 					},