stats/opentelemetry: separate out interceptors for tracing and metrics

stats/opentelemetry/client_metrics.go

@@ -31,13 +31,13 @@
 	"google.golang.org/grpc/status"
 )
 
-type clientStatsHandler struct {
+type clientMetricsHandler struct {
 	estats.MetricsRecorder
 	options       Options
 	clientMetrics clientMetrics
 }
 
-func (h *clientStatsHandler) initializeMetrics() {
+func (h *clientMetricsHandler) initializeMetrics() {
 	// Will set no metrics to record, logically making this stats handler a
 	// no-op.
 	if h.options.MetricsOptions.MeterProvider == nil {
@@ -67,12 +67,12 @@
 	rm.registerMetrics(metrics, meter)
 }
 
-func (h *clientStatsHandler) unaryInterceptor(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+func (h *clientMetricsHandler) unaryInterceptor(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
 	ci := getCallInfo(ctx)
 	if ci == nil {
 		if logger.V(2) {
 			logger.Info("Creating new CallInfo since its not present in context in clientStatsHandler unaryInterceptor")
 		}
 		ci = &callInfo{
 			target: cc.CanonicalTarget(),
 			method: determineMethod(method, opts...),
@@ -107,12 +107,12 @@
 	return "other"
 }
 
-func (h *clientStatsHandler) streamInterceptor(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
+func (h *clientMetricsHandler) streamInterceptor(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
 	ci := getCallInfo(ctx)
 	if ci == nil {
 		if logger.V(2) {
 			logger.Info("Creating new CallInfo since its not present in context in clientStatsHandler streamInterceptor")
 		}
 		ci = &callInfo{
 			target: cc.CanonicalTarget(),
 			method: determineMethod(method, opts...),
@@ -138,7 +138,7 @@
 }
 
 // perCallMetrics records per call metrics for both unary and stream calls.
-func (h *clientStatsHandler) perCallMetrics(ctx context.Context, err error, startTime time.Time, ci *callInfo) {
+func (h *clientMetricsHandler) perCallMetrics(ctx context.Context, err error, startTime time.Time, ci *callInfo) {
 	callLatency := float64(time.Since(startTime)) / float64(time.Second)
 	attrs := otelmetric.WithAttributeSet(otelattribute.NewSet(
 		otelattribute.String("grpc.method", ci.method),
@@ -149,15 +149,15 @@
 }
 
 // TagConn exists to satisfy stats.Handler.
-func (h *clientStatsHandler) TagConn(ctx context.Context, _ *stats.ConnTagInfo) context.Context {
+func (h *clientMetricsHandler) TagConn(ctx context.Context, _ *stats.ConnTagInfo) context.Context {
 	return ctx
 }
 
 // HandleConn exists to satisfy stats.Handler.
-func (h *clientStatsHandler) HandleConn(context.Context, stats.ConnStats) {}
+func (h *clientMetricsHandler) HandleConn(context.Context, stats.ConnStats) {}
 
 // TagRPC implements per RPC attempt context management for metrics.
-func (h *clientStatsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
+func (h *clientMetricsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
 	// Numerous stats handlers can be used for the same channel. The cluster
 	// impl balancer which writes to this will only write once, thus have this
 	// stats handler's per attempt scoped context point to the same optional
@@ -175,21 +175,19 @@
 		ctx = istats.SetLabels(ctx, labels)
 	}
 	ri := getRPCInfo(ctx)
-	var ai *attemptInfo
 	if ri == nil {
-		ai = &attemptInfo{}
-	} else {
-		ai = ri.ai
+		ri = &rpcInfo{ai: &attemptInfo{}}
 	}
+	ai := ri.ai
 	ai.startTime = time.Now()
 	ai.xdsLabels = labels.TelemetryLabels
 	ai.method = removeLeadingSlash(info.FullMethodName)
 
-	return setRPCInfo(ctx, &rpcInfo{ai: ai})
+	return setRPCInfo(ctx, ri)
 }
 
 // HandleRPC handles per RPC stats implementation.
-func (h *clientStatsHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
+func (h *clientMetricsHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
 	ri := getRPCInfo(ctx)
 	if ri == nil {
 		logger.Error("ctx passed into client side stats handler metrics event handling has no client attempt data present")
@@ -198,7 +196,7 @@
 	h.processRPCEvent(ctx, rs, ri.ai)
 }
 
-func (h *clientStatsHandler) processRPCEvent(ctx context.Context, s stats.RPCStats, ai *attemptInfo) {
+func (h *clientMetricsHandler) processRPCEvent(ctx context.Context, s stats.RPCStats, ai *attemptInfo) {
 	switch st := s.(type) {
 	case *stats.Begin:
 		ci := getCallInfo(ctx)
@@ -226,7 +224,7 @@
 	}
 }
 
-func (h *clientStatsHandler) setLabelsFromPluginOption(ai *attemptInfo, incomingMetadata metadata.MD) {
+func (h *clientMetricsHandler) setLabelsFromPluginOption(ai *attemptInfo, incomingMetadata metadata.MD) {
 	if ai.pluginOptionLabels == nil && h.options.MetricsOptions.pluginOption != nil {
 		labels := h.options.MetricsOptions.pluginOption.GetLabels(incomingMetadata)
 		if labels == nil {
@@ -236,7 +234,7 @@
 	}
 }
 
-func (h *clientStatsHandler) processRPCEnd(ctx context.Context, ai *attemptInfo, e *stats.End) {
+func (h *clientMetricsHandler) processRPCEnd(ctx context.Context, ai *attemptInfo, e *stats.End) {
 	ci := getCallInfo(ctx)
 	if ci == nil {
 		logger.Error("ctx passed into client side stats handler metrics event handling has no metrics data present")

stats/opentelemetry/client_tracing.go

@@ -38,54 +38,54 @@
 
 func (h *clientTracingHandler) initializeTraces() {
 	if h.options.TraceOptions.TracerProvider == nil {
 		log.Printf("TracerProvider is not provided in client TraceOptions")
 		return
 	}
 }
 
 func (h *clientTracingHandler) unaryInterceptor(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
 	ci := getCallInfo(ctx)
 	if ci == nil {
 		if logger.V(2) {
 			logger.Info("Creating new CallInfo since its not present in context in clientTracingHandler unaryInterceptor")
 		}
 		ci = &callInfo{
 			target: cc.CanonicalTarget(),
 			method: determineMethod(method, opts...),
 		}
 		ctx = setCallInfo(ctx, ci)
 	}
 
 	var span trace.Span
 	ctx, span = h.createCallTraceSpan(ctx, method)
 	err := invoker(ctx, method, req, reply, cc, opts...)
-	h.perCallTraces(err, span)
+	h.finishTrace(err, span)
 	return err
 }
 
 func (h *clientTracingHandler) streamInterceptor(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
 	ci := getCallInfo(ctx)
 	if ci == nil {
 		if logger.V(2) {
 			logger.Info("Creating new CallInfo since its not present in context in clientTracingHandler streamInterceptor")
 		}
 		ci = &callInfo{
 			target: cc.CanonicalTarget(),
 			method: determineMethod(method, opts...),
 		}
 		ctx = setCallInfo(ctx, ci)
 	}
 
 	var span trace.Span
 	ctx, span = h.createCallTraceSpan(ctx, method)
-	callback := func(err error) { h.perCallTraces(err, span) }
+	callback := func(err error) { h.finishTrace(err, span) }
 	opts = append([]grpc.CallOption{grpc.OnFinish(callback)}, opts...)
 	return streamer(ctx, desc, cc, method, opts...)
 }
 
-// perCallTraces sets the span status based on the RPC result and ends the span.
+// finishTrace sets the span status based on the RPC result and ends the span.
 // It is used to finalize tracing for both unary and streaming calls.
-func (h *clientTracingHandler) perCallTraces(err error, ts trace.Span) {
+func (h *clientTracingHandler) finishTrace(err error, ts trace.Span) {
 	s := status.Convert(err)
 	if s.Code() == grpccodes.OK {
 		ts.SetStatus(otelcodes.Ok, s.Message())
@@ -130,8 +130,8 @@
 func (h *clientTracingHandler) TagRPC(ctx context.Context, _ *stats.RPCTagInfo) context.Context {
 	ri := getRPCInfo(ctx)
 	var ai *attemptInfo
-	if ri == nil {
+	if ri.ai == nil {
 		ai = &attemptInfo{}
 	} else {
 		ai = ri.ai
 	}
@@ -143,8 +143,8 @@
 func (h *clientTracingHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
 	ri := getRPCInfo(ctx)
 	if ri == nil {
 		logger.Error("ctx passed into client side tracing handler trace event handling has no client attempt data present")
 		return
 	}
 	populateSpan(rs, ri.ai)
 }

stats/opentelemetry/metricsregistry_test.go

@@ -47,11 +47,11 @@ type metricsRecorderForTest interface {
 }
 
 func newClientStatsHandler(options MetricsOptions) metricsRecorderForTest {
-	return &clientStatsHandler{options: Options{MetricsOptions: options}}
+	return &clientMetricsHandler{options: Options{MetricsOptions: options}}
 }
 
 func newServerStatsHandler(options MetricsOptions) metricsRecorderForTest {
-	return &serverStatsHandler{options: Options{MetricsOptions: options}}
+	return &serverMetricsHandler{options: Options{MetricsOptions: options}}
 }
 
 // TestMetricsRegistryMetrics tests the OpenTelemetry behavior with respect to

stats/opentelemetry/opentelemetry.go

@@ -65,6 +65,10 @@ type Options struct {
 	TraceOptions experimental.TraceOptions
 }
 
+func (o *Options) isMetricsEnabled() bool {
+	return o.MetricsOptions.MeterProvider != nil
+}
+
 func (o *Options) isTracingEnabled() bool {
 	return o.TraceOptions.TracerProvider != nil
 }
@@ -117,15 +121,31 @@ type MetricsOptions struct {
 // For the traces supported by this instrumentation code, provide an
 // implementation of a TextMapPropagator and OpenTelemetry TracerProvider.
 func DialOption(o Options) grpc.DialOption {
-	csh := &clientStatsHandler{options: o}
-	csh.initializeMetrics()
-	do := joinDialOptions(grpc.WithChainUnaryInterceptor(csh.unaryInterceptor), grpc.WithChainStreamInterceptor(csh.streamInterceptor), grpc.WithStatsHandler(csh))
-	if !o.isTracingEnabled() {
-		return do
+	var unaryInterceptors []grpc.UnaryClientInterceptor
+	var streamInterceptors []grpc.StreamClientInterceptor
+	var do []grpc.DialOption
+
+	if o.isMetricsEnabled() {
+		metricsHandler := &clientMetricsHandler{options: o}
+		metricsHandler.initializeMetrics()
+		unaryInterceptors = append(unaryInterceptors, metricsHandler.unaryInterceptor)
+		streamInterceptors = append(streamInterceptors, metricsHandler.streamInterceptor)
+		do = append(do, grpc.WithStatsHandler(metricsHandler))
+	}
+	if o.isTracingEnabled() {
+		tracingHandler := &clientTracingHandler{options: o}
+		tracingHandler.initializeTraces()
+		unaryInterceptors = append(unaryInterceptors, tracingHandler.unaryInterceptor)
+		streamInterceptors = append(streamInterceptors, tracingHandler.streamInterceptor)
+		do = append(do, grpc.WithStatsHandler(tracingHandler))
+	}
+	if len(unaryInterceptors) > 0 {
+		do = append(do, grpc.WithChainUnaryInterceptor(unaryInterceptors...))
 	}
-	tracingHandler := &clientTracingHandler{options: o}
-	tracingHandler.initializeTraces()
-	return joinDialOptions(do, grpc.WithChainUnaryInterceptor(tracingHandler.unaryInterceptor), grpc.WithChainStreamInterceptor(tracingHandler.streamInterceptor), grpc.WithStatsHandler(tracingHandler))
+	if len(streamInterceptors) > 0 {
+		do = append(do, grpc.WithChainStreamInterceptor(streamInterceptors...))
+	}
+	return joinDialOptions(do...)
 }
 
 var joinServerOptions = internal.JoinServerOptions.(func(...grpc.ServerOption) grpc.ServerOption)
@@ -146,15 +166,19 @@ var joinServerOptions = internal.JoinServerOptions.(func(...grpc.ServerOption) g
 // For the traces supported by this instrumentation code, provide an
 // implementation of a TextMapPropagator and OpenTelemetry TracerProvider.
 func ServerOption(o Options) grpc.ServerOption {
-	ssh := &serverStatsHandler{options: o}
-	ssh.initializeMetrics()
-	so := joinServerOptions(grpc.ChainUnaryInterceptor(ssh.unaryInterceptor), grpc.ChainStreamInterceptor(ssh.streamInterceptor), grpc.StatsHandler(ssh))
-	if !o.isTracingEnabled() {
-		return so
+	var so []grpc.ServerOption
+
+	if o.isMetricsEnabled() {
+		metricsHandler := &serverMetricsHandler{options: o}
+		metricsHandler.initializeMetrics()
+		so = append(so, grpc.ChainUnaryInterceptor(metricsHandler.unaryInterceptor), grpc.ChainStreamInterceptor(metricsHandler.streamInterceptor), grpc.StatsHandler(metricsHandler))
+	}
+	if o.isTracingEnabled() {
+		tracingHandler := &serverTracingHandler{options: o}
+		tracingHandler.initializeTraces()
+		so = append(so, grpc.StatsHandler(tracingHandler))
 	}
-	tracingHandler := &serverTracingHandler{options: o}
-	tracingHandler.initializeTraces()
-	return joinServerOptions(so, grpc.ChainUnaryInterceptor(tracingHandler.unaryInterceptor), grpc.ChainStreamInterceptor(tracingHandler.streamInterceptor), grpc.StatsHandler(tracingHandler))
+	return joinServerOptions(so...)
 }
 
 // callInfo is information pertaining to the lifespan of the RPC client side.

stats/opentelemetry/server_metrics.go

@@ -32,13 +32,13 @@ import (
 	"google.golang.org/grpc/status"
 )
 
-type serverStatsHandler struct {
+type serverMetricsHandler struct {
 	estats.MetricsRecorder
 	options       Options
 	serverMetrics serverMetrics
 }
 
-func (h *serverStatsHandler) initializeMetrics() {
+func (h *serverMetricsHandler) initializeMetrics() {
 	// Will set no metrics to record, logically making this stats handler a
 	// no-op.
 	if h.options.MetricsOptions.MeterProvider == nil {
@@ -90,7 +90,7 @@ func (s *attachLabelsTransportStream) SendHeader(md metadata.MD) error {
 	return s.ServerTransportStream.SendHeader(md)
 }
 
-func (h *serverStatsHandler) unaryInterceptor(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
+func (h *serverMetricsHandler) unaryInterceptor(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
 	var metadataExchangeLabels metadata.MD
 	if h.options.MetricsOptions.pluginOption != nil {
 		metadataExchangeLabels = h.options.MetricsOptions.pluginOption.GetMetadata()
@@ -151,7 +151,7 @@ func (s *attachLabelsStream) SendMsg(m any) error {
 	return s.ServerStream.SendMsg(m)
 }
 
-func (h *serverStatsHandler) streamInterceptor(srv any, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+func (h *serverMetricsHandler) streamInterceptor(srv any, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
 	var metadataExchangeLabels metadata.MD
 	if h.options.MetricsOptions.pluginOption != nil {
 		metadataExchangeLabels = h.options.MetricsOptions.pluginOption.GetMetadata()
@@ -171,15 +171,15 @@ func (h *serverStatsHandler) streamInterceptor(srv any, ss grpc.ServerStream, _
 }
 
 // TagConn exists to satisfy stats.Handler.
-func (h *serverStatsHandler) TagConn(ctx context.Context, _ *stats.ConnTagInfo) context.Context {
+func (h *serverMetricsHandler) TagConn(ctx context.Context, _ *stats.ConnTagInfo) context.Context {
 	return ctx
 }
 
 // HandleConn exists to satisfy stats.Handler.
-func (h *serverStatsHandler) HandleConn(context.Context, stats.ConnStats) {}
+func (h *serverMetricsHandler) HandleConn(context.Context, stats.ConnStats) {}
 
 // TagRPC implements per RPC context management for metrics.
-func (h *serverStatsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
+func (h *serverMetricsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
 	method := info.FullMethodName
 	if h.options.MetricsOptions.MethodAttributeFilter != nil {
 		if !h.options.MetricsOptions.MethodAttributeFilter(method) {
@@ -197,20 +197,18 @@ func (h *serverStatsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo)
 		}
 	}
 	ri := getRPCInfo(ctx)
-	var ai *attemptInfo
 	if ri == nil {
-		ai = &attemptInfo{}
-	} else {
-		ai = ri.ai
+		ri = &rpcInfo{ai: &attemptInfo{}}
 	}
+	ai := ri.ai
 	ai.startTime = time.Now()
 	ai.method = removeLeadingSlash(method)
 
 	return setRPCInfo(ctx, &rpcInfo{ai: ai})
 }
 
 // HandleRPC handles per RPC stats implementation.
-func (h *serverStatsHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
+func (h *serverMetricsHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
 	ri := getRPCInfo(ctx)
 	if ri == nil {
 		logger.Error("ctx passed into server side stats handler metrics event handling has no server call data present")
@@ -219,7 +217,7 @@ func (h *serverStatsHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
 	h.processRPCData(ctx, rs, ri.ai)
 }
 
-func (h *serverStatsHandler) processRPCData(ctx context.Context, s stats.RPCStats, ai *attemptInfo) {
+func (h *serverMetricsHandler) processRPCData(ctx context.Context, s stats.RPCStats, ai *attemptInfo) {
 	switch st := s.(type) {
 	case *stats.InHeader:
 		if ai.pluginOptionLabels == nil && h.options.MetricsOptions.pluginOption != nil {
@@ -243,7 +241,7 @@ func (h *serverStatsHandler) processRPCData(ctx context.Context, s stats.RPCStat
 	}
 }
 
-func (h *serverStatsHandler) processRPCEnd(ctx context.Context, ai *attemptInfo, e *stats.End) {
+func (h *serverMetricsHandler) processRPCEnd(ctx context.Context, ai *attemptInfo, e *stats.End) {
 	latency := float64(time.Since(ai.startTime)) / float64(time.Second)
 	st := "OK"
 	if e.Error != nil {

stats/opentelemetry/server_tracing.go

@@ -33,25 +33,17 @@
 
 func (h *serverTracingHandler) initializeTraces() {
 	if h.options.TraceOptions.TracerProvider == nil {
 		log.Printf("TracerProvider is not provided in server TraceOptions")
 		return
 	}
 }
 
-func (h *serverTracingHandler) unaryInterceptor(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
-	return handler(ctx, req)
-}
-
-func (h *serverTracingHandler) streamInterceptor(srv any, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
-	return handler(srv, ss)
-}
-
 // TagRPC implements per RPC attempt context management for traces.
 func (h *serverTracingHandler) TagRPC(ctx context.Context, _ *stats.RPCTagInfo) context.Context {
 	ri := getRPCInfo(ctx)
 	var ai *attemptInfo
-	if ri == nil {
+	if ri.ai == nil {
 		ai = &attemptInfo{}
 	} else {
 		ai = ri.ai
 	}
@@ -83,9 +75,9 @@
 func (h *serverTracingHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
 	ri := getRPCInfo(ctx)
 	if ri == nil {
 		logger.Error("ctx passed into server side tracing handler trace event handling has no server call data present")
 		return
 	}
 	populateSpan(rs, ri.ai)
 }