stats/opentelemetry: separate out interceptors for tracing and metrics

stats/opentelemetry/client_metrics.go

@@ -21,27 +21,23 @@
 	"sync/atomic"
 	"time"
 
-	otelcodes "go.opentelemetry.io/otel/codes"
-	"go.opentelemetry.io/otel/trace"
+	otelattribute "go.opentelemetry.io/otel/attribute"
+	otelmetric "go.opentelemetry.io/otel/metric"
 	"google.golang.org/grpc"
-	grpccodes "google.golang.org/grpc/codes"
 	estats "google.golang.org/grpc/experimental/stats"
 	istats "google.golang.org/grpc/internal/stats"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/stats"
 	"google.golang.org/grpc/status"
-
-	otelattribute "go.opentelemetry.io/otel/attribute"
-	otelmetric "go.opentelemetry.io/otel/metric"
 )
 
-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 {
@@ -71,12 +67,18 @@
 	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 {
-	ci := &callInfo{
-		target: cc.CanonicalTarget(),
-		method: h.determineMethod(method, opts...),
+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...),
+		}
+		ctx = setCallInfo(ctx, ci)
 	}
-	ctx = setCallInfo(ctx, ci)
 
 	if h.options.MetricsOptions.pluginOption != nil {
 		md := h.options.MetricsOptions.pluginOption.GetMetadata()
@@ -88,19 +90,15 @@
 	}
 
 	startTime := time.Now()
-	var span trace.Span
-	if h.options.isTracingEnabled() {
-		ctx, span = h.createCallTraceSpan(ctx, method)
-	}
 	err := invoker(ctx, method, req, reply, cc, opts...)
-	h.perCallTracesAndMetrics(ctx, err, startTime, ci, span)
+	h.perCallMetrics(ctx, err, startTime, ci)
 	return err
 }
 
 // determineMethod determines the method to record attributes with. This will be
 // "other" if StaticMethod isn't specified or if method filter is set and
 // specifies, the method name as is otherwise.
-func (h *clientStatsHandler) determineMethod(method string, opts ...grpc.CallOption) string {
+func determineMethod(method string, opts ...grpc.CallOption) string {
 	for _, opt := range opts {
 		if _, ok := opt.(grpc.StaticMethodCallOption); ok {
 			return removeLeadingSlash(method)
@@ -109,12 +107,18 @@
 	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) {
-	ci := &callInfo{
-		target: cc.CanonicalTarget(),
-		method: h.determineMethod(method, opts...),
+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 = setCallInfo(ctx, ci)
 
 	if h.options.MetricsOptions.pluginOption != nil {
 		md := h.options.MetricsOptions.pluginOption.GetMetadata()
@@ -126,49 +130,34 @@
 	}
 
 	startTime := time.Now()
-	var span trace.Span
-	if h.options.isTracingEnabled() {
-		ctx, span = h.createCallTraceSpan(ctx, method)
-	}
 	callback := func(err error) {
-		h.perCallTracesAndMetrics(ctx, err, startTime, ci, span)
+		h.perCallMetrics(ctx, err, startTime, ci)
 	}
 	opts = append([]grpc.CallOption{grpc.OnFinish(callback)}, opts...)
 	return streamer(ctx, desc, cc, method, opts...)
 }
 
-// perCallTracesAndMetrics records per call trace spans and metrics.
-func (h *clientStatsHandler) perCallTracesAndMetrics(ctx context.Context, err error, startTime time.Time, ci *callInfo, ts trace.Span) {
-	if h.options.isTracingEnabled() {
-		s := status.Convert(err)
-		if s.Code() == grpccodes.OK {
-			ts.SetStatus(otelcodes.Ok, s.Message())
-		} else {
-			ts.SetStatus(otelcodes.Error, s.Message())
-		}
-		ts.End()
-	}
-	if h.options.isMetricsEnabled() {
-		callLatency := float64(time.Since(startTime)) / float64(time.Second)
-		attrs := otelmetric.WithAttributeSet(otelattribute.NewSet(
-			otelattribute.String("grpc.method", ci.method),
-			otelattribute.String("grpc.target", ci.target),
-			otelattribute.String("grpc.status", canonicalString(status.Code(err))),
-		))
-		h.clientMetrics.callDuration.Record(ctx, callLatency, attrs)
-	}
+// perCallMetrics records per call metrics for both unary and stream calls.
+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),
+		otelattribute.String("grpc.target", ci.target),
+		otelattribute.String("grpc.status", canonicalString(status.Code(err))),
+	))
+	h.clientMetrics.callDuration.Record(ctx, callLatency, attrs)
 }
 
 // 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.
-func (h *clientStatsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
+// 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
 	// 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
@@ -185,34 +174,34 @@
 		}
 		ctx = istats.SetLabels(ctx, labels)
 	}
-	ai := &attemptInfo{
-		startTime: time.Now(),
-		xdsLabels: labels.TelemetryLabels,
-		method:    removeLeadingSlash(info.FullMethodName),
+	ri := getRPCInfo(ctx)
+	if ri == nil {
+		ri = &rpcInfo{}
 	}
-	if h.options.isTracingEnabled() {
-		ctx, ai = h.traceTagRPC(ctx, ai, info.NameResolutionDelay)
+	var ai *attemptInfo
+	if ri.ai == nil {
+		ai = &attemptInfo{}
+	} else {
+		ai = ri.ai
 	}
-	return setRPCInfo(ctx, &rpcInfo{
-		ai: ai,
-	})
+	ai.startTime = time.Now()
+	ai.xdsLabels = labels.TelemetryLabels
+	ai.method = removeLeadingSlash(info.FullMethodName)
+
+	return setRPCInfo(ctx, &rpcInfo{ai: ai})
 }
 
-func (h *clientStatsHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
+// HandleRPC handles per RPC stats implementation.
+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")
 		return
 	}
-	if h.options.isMetricsEnabled() {
-		h.processRPCEvent(ctx, rs, ri.ai)
-	}
-	if h.options.isTracingEnabled() {
-		populateSpan(rs, ri.ai)
-	}
+	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)
@@ -240,7 +229,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 {
@@ -250,7 +239,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

@@ -18,24 +18,92 @@
 
 import (
 	"context"
+	"log"
 	"strings"
 
+	otelcodes "go.opentelemetry.io/otel/codes"
 	"go.opentelemetry.io/otel/trace"
 	"google.golang.org/grpc"
+	grpccodes "google.golang.org/grpc/codes"
+	"google.golang.org/grpc/stats"
 	otelinternaltracing "google.golang.org/grpc/stats/opentelemetry/internal/tracing"
+	"google.golang.org/grpc/status"
 )
 
 const (
 	delayedResolutionEventName = "Delayed name resolution complete"
 	tracerName                 = "grpc-go"
 )
 
+type clientTracingHandler struct {
+	options Options
+}
+
+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.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.finishTrace(err, span) }
+	opts = append([]grpc.CallOption{grpc.OnFinish(callback)}, opts...)
+	return streamer(ctx, desc, cc, method, opts...)
+}
+
+// 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) finishTrace(err error, ts trace.Span) {
+	s := status.Convert(err)
+	if s.Code() == grpccodes.OK {
+		ts.SetStatus(otelcodes.Ok, s.Message())
+	} else {
+		ts.SetStatus(otelcodes.Error, s.Message())
+	}
+	ts.End()
+}
+
 // traceTagRPC populates provided context with a new span using the
 // TextMapPropagator supplied in trace options and internal itracing.carrier.
 // It creates a new outgoing carrier which serializes information about this
 // span into gRPC Metadata, if TextMapPropagator is provided in the trace
 // options. if TextMapPropagator is not provided, it returns the context as is.
-func (h *clientStatsHandler) traceTagRPC(ctx context.Context, ai *attemptInfo, nameResolutionDelayed bool) (context.Context, *attemptInfo) {
+func (h *clientTracingHandler) traceTagRPC(ctx context.Context, ai *attemptInfo, nameResolutionDelayed bool) (context.Context, *attemptInfo) {
 	// Add a "Delayed name resolution complete" event to the call span
 	// if there was name resolution delay. In case of multiple retry attempts,
 	// ensure that event is added only once.
@@ -55,13 +123,40 @@
 
 // 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 *clientStatsHandler) createCallTraceSpan(ctx context.Context, method string) (context.Context, trace.Span) {
-	if h.options.TraceOptions.TracerProvider == nil {
-		logger.Error("TraceProvider is not provided in trace options")
-		return ctx, nil
-	}
+func (h *clientTracingHandler) createCallTraceSpan(ctx context.Context, method string) (context.Context, trace.Span) {
 	mn := 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
 }
+
+// TagConn exists to satisfy stats.Handler for tracing.
+func (h *clientTracingHandler) TagConn(ctx context.Context, _ *stats.ConnTagInfo) context.Context {
+	return ctx
+}
+
+// HandleConn exists to satisfy stats.Handler for tracing.
+func (h *clientTracingHandler) HandleConn(context.Context, stats.ConnStats) {}
+
+// 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 = h.traceTagRPC(ctx, ai, info.NameResolutionDelay)
+	return setRPCInfo(ctx, &rpcInfo{ai: ai})
+}
+
+// HandleRPC handles per RPC tracing implementation.
+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