otelgrpc: add custom attributes to the stats handler

Fixes #3894
open-telemetry · Jun 6, 2024 · a6b8d67 · a6b8d67
1 parent 0430e80
commit a6b8d67
Show file tree

Hide file tree

Showing 4 changed files with 140 additions and 52 deletions.
diff --git a/instrumentation/google.golang.org/grpc/otelgrpc/config.go b/instrumentation/google.golang.org/grpc/otelgrpc/config.go
@@ -42,6 +42,7 @@ type config struct {
  TracerProvider trace.TracerProvider
  MeterProvider metric.MeterProvider
  SpanStartOptions []trace.SpanStartOption
+ Attributes []attribute.KeyValue
 
  ReceivedEvent bool
  SentEvent bool
@@ -257,3 +258,16 @@ func (o spanStartOption) apply(c *config) {
 func WithSpanOptions(opts ...trace.SpanStartOption) Option {
  return spanStartOption{opts}
 }
+
+type attributesOption struct{ a []attribute.KeyValue }
+
+func (o attributesOption) apply(c *config) {
+ if o.a != nil {
+ c.Attributes = o.a
+ }
+}
+
+// WithAttributes returns an Option to add custom attributes to the spans and metrics.
+func WithAttributes(a ...attribute.KeyValue) Option {
+ return attributesOption{a: a}
+}
diff --git a/instrumentation/google.golang.org/grpc/otelgrpc/stats_handler.go b/instrumentation/google.golang.org/grpc/otelgrpc/stats_handler.go
@@ -58,6 +58,9 @@ func (h *serverHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) cont
 
  name, attrs := internal.ParseFullMethod(info.FullMethodName)
  attrs = append(attrs, RPCSystemGRPC)
+
+ attrs = append(attrs, h.config.Attributes...)
+
  ctx, _ = h.tracer.Start(
  trace.ContextWithRemoteSpanContext(ctx, trace.SpanContextFromContext(ctx)),
  name,
@@ -98,6 +101,9 @@ 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)
+
+ attrs = append(attrs, h.config.Attributes...)
+
  ctx, _ = h.tracer.Start(
  ctx,
  name,