rpc,server,tracing: refactor logic of tracing interceptors

These changes are part of DRPC migration.
As part of a plan to implement DRPC compliant interceptors,
the common logic of tracing interceptors are moved to a common
interceptorutil module. This enables us to reuse tracing logic
in drpc interceptors.

Epic: none
Part of: CRDB-49356

Release note: None

Author: shaikzakiriitm

pkg/BUILD.bazel

@@ -2754,6 +2754,7 @@ GO_TARGETS = [
     "//pkg/util/tracing/tracingpb:tracingpb",
     "//pkg/util/tracing/tracingservicepb:tracingservicepb",
     "//pkg/util/tracing/tracingui:tracingui",
+    "//pkg/util/tracing/tracingutil:tracingutil",
     "//pkg/util/tracing/zipper:zipper",
     "//pkg/util/tracing:tracing",
     "//pkg/util/tracing:tracing_test",

pkg/rpc/BUILD.bazel

@@ -65,6 +65,7 @@ go_library(
         "//pkg/util/timeutil",
         "//pkg/util/tracing",
         "//pkg/util/tracing/grpcinterceptor",
+        "//pkg/util/tracing/tracingutil",
         "//pkg/util/uuid",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_cockroachdb_logtags//:logtags",

pkg/rpc/context.go

@@ -39,6 +39,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing/grpcinterceptor"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing/tracingutil"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/logtags"
@@ -728,7 +729,7 @@ func makeInternalClientAdapter(
 	batchServerHandler := chainUnaryServerInterceptors(
 		&grpc.UnaryServerInfo{
 			Server:     server,
-			FullMethod: grpcinterceptor.BatchMethodName,
+			FullMethod: tracingutil.BatchMethodName,
 		},
 		serverUnaryInterceptors,
 		func(ctx context.Context, req interface{}) (interface{}, error) {
@@ -812,7 +813,7 @@ func makeInternalClientAdapter(
 				ctx = tracing.ContextWithSpan(ctx, nil)
 			}
 
-			err := batchClientHandler(ctx, grpcinterceptor.BatchMethodName, ba, reply, nil /* ClientConn */)
+			err := batchClientHandler(ctx, tracingutil.BatchMethodName, ba, reply, nil /* ClientConn */)
 			return reply, err
 		},
 	}

pkg/server/BUILD.bazel

@@ -342,11 +342,11 @@ go_library(
         "//pkg/util/timeutil/ptp",
         "//pkg/util/tracing",
         "//pkg/util/tracing/collector",
-        "//pkg/util/tracing/grpcinterceptor",
         "//pkg/util/tracing/service",
         "//pkg/util/tracing/tracingpb",
         "//pkg/util/tracing/tracingservicepb",
         "//pkg/util/tracing/tracingui",
+        "//pkg/util/tracing/tracingutil",
         "//pkg/util/tracing/zipper",
         "//pkg/util/uint128",
         "//pkg/util/unique",

pkg/server/node.go

@@ -74,8 +74,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
-	"github.com/cockroachdb/cockroach/pkg/util/tracing/grpcinterceptor"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing/tracingpb"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing/tracingutil"
 	"github.com/cockroachdb/cockroach/pkg/util/unique"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
 	"github.com/cockroachdb/crlib/crtime"
@@ -2004,7 +2004,7 @@ func setupSpanForIncomingRPC(
 		// request that didn't specify tracing information. We make a child span
 		// if the incoming request would like to be traced.
 		ctx, newSpan = tracing.ChildSpan(ctx,
-			grpcinterceptor.BatchMethodName, tracing.WithServerSpanKind)
+			tracingutil.BatchMethodName, tracing.WithServerSpanKind)
 	} else {
 		// Non-local call. Tracing information comes from the request proto.
 
@@ -2016,7 +2016,7 @@ func setupSpanForIncomingRPC(
 		}
 
 		ctx, newSpan = tr.StartSpanCtx(
-			ctx, grpcinterceptor.BatchMethodName,
+			ctx, tracingutil.BatchMethodName,
 			tracing.WithRemoteParentFromTraceInfo(ba.TraceInfo),
 			tracing.WithServerSpanKind)
 	}

pkg/sql/distsql/BUILD.bazel

@@ -33,6 +33,7 @@ go_library(
         "//pkg/util/tochar",
         "//pkg/util/tracing",
         "//pkg/util/tracing/grpcinterceptor",
+        "//pkg/util/tracing/tracingutil",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_cockroachdb_logtags//:logtags",
         "@com_github_cockroachdb_redact//:redact",

pkg/sql/distsql/server.go

@@ -37,6 +37,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/tochar"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing/grpcinterceptor"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing/tracingutil"
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/logtags"
 	"github.com/cockroachdb/redact"
@@ -599,13 +600,13 @@ func (ds *ServerImpl) setupSpanForIncomingRPC(
 		// It's not expected to have a span in the context since the gRPC server
 		// interceptor that generally opens spans exempts this particular RPC. Note
 		// that this method is not called for flows local to the gateway.
-		return tr.StartSpanCtx(ctx, grpcinterceptor.SetupFlowMethodName,
+		return tr.StartSpanCtx(ctx, tracingutil.SetupFlowMethodName,
 			tracing.WithParent(parentSpan),
 			tracing.WithServerSpanKind)
 	}
 
 	if !req.TraceInfo.Empty() {
-		return tr.StartSpanCtx(ctx, grpcinterceptor.SetupFlowMethodName,
+		return tr.StartSpanCtx(ctx, tracingutil.SetupFlowMethodName,
 			tracing.WithRemoteParentFromTraceInfo(req.TraceInfo),
 			tracing.WithServerSpanKind)
 	}
@@ -615,7 +616,7 @@ func (ds *ServerImpl) setupSpanForIncomingRPC(
 	if err != nil {
 		log.Warningf(ctx, "error extracting tracing info from gRPC: %s", err)
 	}
-	return tr.StartSpanCtx(ctx, grpcinterceptor.SetupFlowMethodName,
+	return tr.StartSpanCtx(ctx, tracingutil.SetupFlowMethodName,
 		tracing.WithRemoteParentFromSpanMeta(remoteParent),
 		tracing.WithServerSpanKind)
 }

pkg/util/tracing/BUILD.bazel

@@ -87,8 +87,8 @@ go_test(
         "//pkg/util/log",
         "//pkg/util/randutil",
         "//pkg/util/timeutil",
-        "//pkg/util/tracing/grpcinterceptor",
         "//pkg/util/tracing/tracingpb",
+        "//pkg/util/tracing/tracingutil",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_cockroachdb_logtags//:logtags",
         "@com_github_gogo_protobuf//proto",

pkg/util/tracing/grpcinterceptor/BUILD.bazel

@@ -9,6 +9,7 @@ go_library(
         "//pkg/util/ctxutil",
         "//pkg/util/grpcutil",
         "//pkg/util/tracing",
+        "//pkg/util/tracing/tracingutil",
         "@com_github_cockroachdb_errors//:errors",
         "@io_opentelemetry_go_otel//attribute",
         "@io_opentelemetry_go_otel//codes",
@@ -30,6 +31,7 @@ go_test(
         "//pkg/util/stop",
         "//pkg/util/tracing",
         "//pkg/util/tracing/tracingpb",
+        "//pkg/util/tracing/tracingutil",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_gogo_protobuf//types",
         "@com_github_stretchr_testify//require",

pkg/util/tracing/grpcinterceptor/grpc_interceptor.go

@@ -13,6 +13,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/ctxutil"
 	"github.com/cockroachdb/cockroach/pkg/util/grpcutil"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing/tracingutil"
 	"github.com/cockroachdb/errors"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/codes"
@@ -33,43 +34,6 @@ func ExtractSpanMetaFromGRPCCtx(
 	return tracer.ExtractMetaFrom(tracing.MetadataCarrier{MD: md})
 }
 
-// setGRPCErrorTag sets an error tag on the span.
-func setGRPCErrorTag(sp *tracing.Span, err error) {
-	if err == nil {
-		return
-	}
-	s, _ := status.FromError(err)
-	sp.SetTag("response_code", attribute.IntValue(int(codes.Error)))
-	sp.SetOtelStatus(codes.Error, s.Message())
-}
-
-// BatchMethodName is the method name of Internal.Batch RPC.
-const BatchMethodName = "/cockroach.roachpb.Internal/Batch"
-
-// BatchStreamMethodName is the method name of the Internal.BatchStream RPC.
-const BatchStreamMethodName = "/cockroach.roachpb.Internal/BatchStream"
-
-// sendKVBatchMethodName is the method name for adminServer.SendKVBatch.
-const sendKVBatchMethodName = "/cockroach.server.serverpb.Admin/SendKVBatch"
-
-// SetupFlowMethodName is the method name of DistSQL.SetupFlow RPC.
-const SetupFlowMethodName = "/cockroach.sql.distsqlrun.DistSQL/SetupFlow"
-const flowStreamMethodName = "/cockroach.sql.distsqlrun.DistSQL/FlowStream"
-
-// methodExcludedFromTracing returns true if a call to the given RPC method does
-// not need to propagate tracing info. Some RPCs (Internal.Batch,
-// DistSQL.SetupFlow) have dedicated fields for passing along the tracing
-// context in the request, which is more efficient than letting the RPC
-// interceptors deal with it. Others (DistSQL.FlowStream) are simply exempt from
-// tracing because it's not worth it.
-func methodExcludedFromTracing(method string) bool {
-	return method == BatchMethodName ||
-		method == BatchStreamMethodName ||
-		method == sendKVBatchMethodName ||
-		method == SetupFlowMethodName ||
-		method == flowStreamMethodName
-}
-
 // ServerInterceptor returns a grpc.UnaryServerInterceptor suitable
 // for use in a grpc.NewServer call.
 //
@@ -92,7 +56,7 @@ func ServerInterceptor(tracer *tracing.Tracer) grpc.UnaryServerInterceptor {
 		info *grpc.UnaryServerInfo,
 		handler grpc.UnaryHandler,
 	) (interface{}, error) {
-		if methodExcludedFromTracing(info.FullMethod) {
+		if tracingutil.MethodExcludedFromTracing(info.FullMethod) {
 			return handler(ctx, req)
 		}
 
@@ -139,7 +103,7 @@ func ServerInterceptor(tracer *tracing.Tracer) grpc.UnaryServerInterceptor {
 // application-specific gRPC handler(s) to access.
 func StreamServerInterceptor(tracer *tracing.Tracer) grpc.StreamServerInterceptor {
 	return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
-		if methodExcludedFromTracing(info.FullMethod) {
+		if tracingutil.MethodExcludedFromTracing(info.FullMethod) {
 			return handler(srv, ss)
 		}
 		spanMeta, err := ExtractSpanMetaFromGRPCCtx(ss.Context(), tracer)
@@ -188,19 +152,6 @@ func (ss *tracingServerStream) Context() context.Context {
 	return ss.ctx
 }
 
-func injectSpanMeta(
-	ctx context.Context, tracer *tracing.Tracer, clientSpan *tracing.Span,
-) context.Context {
-	md, ok := metadata.FromOutgoingContext(ctx)
-	if !ok {
-		md = metadata.New(nil)
-	} else {
-		md = md.Copy()
-	}
-	tracer.InjectMetaInto(clientSpan.Meta(), tracing.MetadataCarrier{MD: md})
-	return metadata.NewOutgoingContext(ctx, md)
-}
-
 // ClientInterceptor returns a grpc.UnaryClientInterceptor suitable
 // for use in a grpc.Dial call.
 //
@@ -229,17 +180,13 @@ func ClientInterceptor(
 		invoker grpc.UnaryInvoker,
 		opts ...grpc.CallOption,
 	) error {
-		// Local RPCs don't need any special tracing, since the caller's context
-		// will be used on the "server".
-		_, localRequest := grpcutil.IsLocalRequestContext(ctx)
-		if localRequest {
-			return invoker(ctx, method, req, resp, cc, opts...)
-		}
-		parent := tracing.SpanFromContext(ctx)
-		if !tracing.SpanInclusionFuncForClient(parent) {
+		skipTracing := tracingutil.ShouldSkipClientTracing(ctx)
+		if skipTracing {
 			return invoker(ctx, method, req, resp, cc, opts...)
 		}
 
+		// Create clientSpan here after determining that we shouldn't skip tracing
+		parent := tracing.SpanFromContext(ctx)
 		clientSpan := tracer.StartSpan(
 			method,
 			tracing.WithParent(parent),
@@ -248,11 +195,11 @@ func ClientInterceptor(
 		init(clientSpan)
 		defer clientSpan.Finish()
 
-		// For most RPCs we pass along tracing info as gRPC metadata. Some select
-		// RPCs carry the tracing in the request protos, which is more efficient.
-		if !methodExcludedFromTracing(method) {
-			ctx = injectSpanMeta(ctx, tracer, clientSpan)
+		// For most RPCs we pass along tracing info as metadata
+		if !tracingutil.MethodExcludedFromTracing(method) {
+			ctx = tracingutil.InjectSpanMeta(ctx, tracer, clientSpan)
 		}
+
 		if invoker != nil {
 			err := invoker(ctx, method, req, resp, cc, opts...)
 			if err != nil {
@@ -294,27 +241,23 @@ func StreamClientInterceptor(
 		streamer grpc.Streamer,
 		opts ...grpc.CallOption,
 	) (grpc.ClientStream, error) {
-		// Local RPCs don't need any special tracing, since the caller's context
-		// will be used on the "server".
-		_, localRequest := grpcutil.IsLocalRequestContext(ctx)
-		if localRequest {
-			return streamer(ctx, desc, cc, method, opts...)
-		}
-		parent := tracing.SpanFromContext(ctx)
-		if !tracing.SpanInclusionFuncForClient(parent) {
+		skipTracing := tracingutil.ShouldSkipClientTracing(ctx)
+		if skipTracing {
 			return streamer(ctx, desc, cc, method, opts...)
 		}
 
-		// Create a span that will live for the life of the stream.
+		// Create clientSpan here after determining that we shouldn't skip tracing
+		parent := tracing.SpanFromContext(ctx)
 		clientSpan := tracer.StartSpan(
 			method,
 			tracing.WithParent(parent),
 			tracing.WithClientSpanKind,
 		)
 		init(clientSpan)
 
-		if !methodExcludedFromTracing(method) {
-			ctx = injectSpanMeta(ctx, tracer, clientSpan)
+		// For most RPCs we pass along tracing info as metadata
+		if !tracingutil.MethodExcludedFromTracing(method) {
+			ctx = tracingutil.InjectSpanMeta(ctx, tracer, clientSpan)
 		}
 
 		cs, err := streamer(ctx, desc, cc, method, opts...)
@@ -421,3 +364,13 @@ func (cs *tracingClientStream) CloseSend() error {
 	}
 	return errors.Wrap(err, "close send error")
 }
+
+// setGRPCErrorTag sets an error tag on the span.
+func setGRPCErrorTag(sp *tracing.Span, err error) {
+	if err == nil {
+		return
+	}
+	s, _ := status.FromError(err)
+	sp.SetTag("response_code", attribute.IntValue(int(codes.Error)))
+	sp.SetOtelStatus(codes.Error, s.Message())
+}