otelgrpc: add custom attributes to the stats handler

Fixes #3894

Author: inigohu

instrumentation/google.golang.org/grpc/otelgrpc/stats_handler.go

@@ -30,6 +30,16 @@ type gRPCContext struct {
 	record           bool
 }
 
+// AddAttrs adds attributes to the given context.
+func AddAttrs(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,17 +68,21 @@ 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)
 	}
@@ -98,17 +112,21 @@ 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)
 	}