Skip to content

Commit

Permalink
otelgrpc: add custom attributes to the stats handler
Browse files Browse the repository at this point in the history
Fixes #3894
  • Loading branch information
inigohu committed Jun 4, 2024
1 parent 48ab4e2 commit e4f7d19
Showing 1 changed file with 28 additions and 10 deletions.
38 changes: 28 additions & 10 deletions instrumentation/google.golang.org/grpc/otelgrpc/stats_handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down Expand Up @@ -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)
}
Expand Down Expand Up @@ -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)
}
Expand Down

0 comments on commit e4f7d19

Please sign in to comment.