stats/opentelemetry: separate out interceptors for tracing and metrics

stats/opentelemetry/client_metrics.go

@@ -67,18 +67,25 @@
 	rm.registerMetrics(metrics, meter)
 }
 
-func (h *clientMetricsHandler) unaryInterceptor(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+// getOrCreateCallInfo returns the existing callInfo from context if present,
+// or creates and attaches a new one.
+func getOrCreateCallInfo(ctx context.Context, cc *grpc.ClientConn, method string, opts ...grpc.CallOption) (context.Context, *callInfo) {
 	ci := getCallInfo(ctx)
 	if ci == nil {
 		if logger.V(2) {
-			logger.Info("Creating new CallInfo since its not present in context in clientStatsHandler unaryInterceptor")
+			logger.Info("Creating new CallInfo since its not present in context")
 		}
 		ci = &callInfo{
 			target: cc.CanonicalTarget(),
 			method: determineMethod(method, opts...),
 		}
 		ctx = setCallInfo(ctx, ci)
 	}
+	return ctx, ci
+}
+
+func (h *clientMetricsHandler) unaryInterceptor(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+	ctx, ci := getOrCreateCallInfo(ctx, cc, method, opts...)
 
 	if h.options.MetricsOptions.pluginOption != nil {
 		md := h.options.MetricsOptions.pluginOption.GetMetadata()
@@ -108,17 +115,7 @@
 }
 
 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...),
-		}
-		ctx = setCallInfo(ctx, ci)
-	}
+	ctx, ci := getOrCreateCallInfo(ctx, cc, method, opts...)
 
 	if h.options.MetricsOptions.pluginOption != nil {
 		md := h.options.MetricsOptions.pluginOption.GetMetadata()
@@ -156,6 +153,20 @@
 // HandleConn exists to satisfy stats.Handler.
 func (h *clientMetricsHandler) HandleConn(context.Context, stats.ConnStats) {}
 
+// getOrCreateRPCAttemptInfo retrieves or creates an rpc attemptInfo object
+// and ensures it is set in the context along with the rpcInfo.
+func getOrCreateRPCAttemptInfo(ctx context.Context) (context.Context, *attemptInfo) {
+	ri := getRPCInfo(ctx)
+	if ri == nil {
+		ri = &rpcInfo{}
+	}
+	ai := ri.ai
+	if ai == nil {
+		ai = &attemptInfo{}
+	}
+	return setRPCInfo(ctx, &rpcInfo{ai: ai}), ai
+}
+
 // TagRPC implements per RPC attempt context management for metrics.
 func (h *clientMetricsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
 	// Numerous stats handlers can be used for the same channel. The cluster
@@ -174,16 +185,7 @@
 		}
 		ctx = istats.SetLabels(ctx, labels)
 	}
-	ri := getRPCInfo(ctx)
-	if ri == nil {
-		ri = &rpcInfo{}
-	}
-	var ai *attemptInfo
-	if ri.ai == nil {
-		ai = &attemptInfo{}
-	} else {
-		ai = ri.ai
-	}
+	ctx, ai := getOrCreateRPCAttemptInfo(ctx)
 	ai.startTime = time.Now()
 	ai.xdsLabels = labels.TelemetryLabels
 	ai.method = removeLeadingSlash(info.FullMethodName)

stats/opentelemetry/client_tracing.go

@@ -41,23 +41,13 @@
 
 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)
-	}
+	ctx, _ = getOrCreateCallInfo(ctx, cc, method, opts...)
 
 	var span trace.Span
 	ctx, span = h.createCallTraceSpan(ctx, method)
@@ -67,17 +57,7 @@
 }
 
 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)
-	}
+	ctx, _ = getOrCreateCallInfo(ctx, cc, method, opts...)
 
 	var span trace.Span
 	ctx, span = h.createCallTraceSpan(ctx, method)
@@ -124,7 +104,7 @@
 // createCallTraceSpan creates a call span to put in the provided context using
 // provided TraceProvider. If TraceProvider is nil, it returns context as is.
 func (h *clientTracingHandler) createCallTraceSpan(ctx context.Context, method string) (context.Context, trace.Span) {
-	mn := strings.Replace(removeLeadingSlash(method), "/", ".", -1)
+	mn := "Sent." + strings.Replace(removeLeadingSlash(method), "/", ".", -1)
 	tracer := h.options.TraceOptions.TracerProvider.Tracer(tracerName, trace.WithInstrumentationVersion(grpc.Version))
 	ctx, span := tracer.Start(ctx, mn, trace.WithSpanKind(trace.SpanKindClient))
 	return ctx, span
@@ -140,13 +120,7 @@
 
 // TagRPC implements per RPC attempt context management for traces.
 func (h *clientTracingHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
-	ri := getRPCInfo(ctx)
-	var ai *attemptInfo
-	if ri == nil {
-		ai = &attemptInfo{}
-	} else {
-		ai = ri.ai
-	}
+	ctx, ai := getOrCreateRPCAttemptInfo(ctx)
 	ctx, ai = h.traceTagRPC(ctx, ai, info.NameResolutionDelay)
 	return setRPCInfo(ctx, &rpcInfo{ai: ai})
 }
@@ -155,8 +129,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/server_metrics.go

@@ -196,16 +196,7 @@ func (h *serverMetricsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInf
 			method = "other"
 		}
 	}
-	ri := getRPCInfo(ctx)
-	if ri == nil {
-		ri = &rpcInfo{}
-	}
-	var ai *attemptInfo
-	if ri.ai == nil {
-		ai = &attemptInfo{}
-	} else {
-		ai = ri.ai
-	}
+	ctx, ai := getOrCreateRPCAttemptInfo(ctx)
 	ai.startTime = time.Now()
 	ai.method = removeLeadingSlash(method)
 

stats/opentelemetry/server_tracing.go

@@ -33,20 +33,14 @@
 
 func (h *serverTracingHandler) initializeTraces() {
 	if h.options.TraceOptions.TracerProvider == nil {
 		log.Printf("TracerProvider is not provided in server TraceOptions")
 		return
 	}
 }
 
 // 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 {
-		ai = &attemptInfo{}
-	} else {
-		ai = ri.ai
-	}
+	ctx, ai := getOrCreateRPCAttemptInfo(ctx)
 	ctx, ai = h.traceTagRPC(ctx, ai)
 	return setRPCInfo(ctx, &rpcInfo{ai: ai})
 }
@@ -59,7 +53,7 @@
 // If TextMapPropagator is not provided in the trace options, it returns context
 // as is.
 func (h *serverTracingHandler) traceTagRPC(ctx context.Context, ai *attemptInfo) (context.Context, *attemptInfo) {
-	mn := strings.Replace(ai.method, "/", ".", -1)
+	mn := "Recv." + strings.Replace(ai.method, "/", ".", -1)
 	var span trace.Span
 	tracer := h.options.TraceOptions.TracerProvider.Tracer(tracerName, trace.WithInstrumentationVersion(grpc.Version))
 	ctx = h.options.TraceOptions.TextMapPropagator.Extract(ctx, otelinternaltracing.NewIncomingCarrier(ctx))
@@ -75,9 +69,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)
 }