diff --git a/instrumentation/google.golang.org/grpc/otelgrpc/stats_handler.go b/instrumentation/google.golang.org/grpc/otelgrpc/stats_handler.go index 201867a8694..2bd68905e06 100644 --- 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 } +// 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) }