From 23877987191a8458ef674792649cdb872d161e76 Mon Sep 17 00:00:00 2001
From: owensmallwood <owen.smallwood@grafana.com>
Date: Tue, 6 Aug 2024 09:52:11 -0600
Subject: [PATCH] Unified Storage: Propagate traces from the apiserver grpc
 client (#91226)

---
 pkg/services/apiserver/service.go             |  8 ++++++-
 .../grpcserver/interceptors/tracing.go        | 21 +------------------
 pkg/services/grpcserver/service.go            |  3 ++-
 3 files changed, 10 insertions(+), 22 deletions(-)

diff --git a/pkg/services/apiserver/service.go b/pkg/services/apiserver/service.go
index 0afa1621fc9..92b1d659ebf 100644
--- a/pkg/services/apiserver/service.go
+++ b/pkg/services/apiserver/service.go
@@ -8,6 +8,7 @@ import (
 
 	"github.com/grafana/dskit/services"
 	"github.com/prometheus/client_golang/prometheus"
+	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -281,8 +282,13 @@ func (s *service) start(ctx context.Context) error {
 		if !s.features.IsEnabledGlobally(featuremgmt.FlagUnifiedStorage) {
 			return fmt.Errorf("unified storage requires the unifiedStorage feature flag")
 		}
+
+		opts := []grpc.DialOption{
+			grpc.WithStatsHandler(otelgrpc.NewClientHandler()),
+			grpc.WithTransportCredentials(insecure.NewCredentials()),
+		}
 		// Create a connection to the gRPC server
-		conn, err := grpc.NewClient(o.StorageOptions.Address, grpc.WithTransportCredentials(insecure.NewCredentials()))
+		conn, err := grpc.NewClient(o.StorageOptions.Address, opts...)
 		if err != nil {
 			return err
 		}
diff --git a/pkg/services/grpcserver/interceptors/tracing.go b/pkg/services/grpcserver/interceptors/tracing.go
index e7bd3b6a925..f885d01407a 100644
--- a/pkg/services/grpcserver/interceptors/tracing.go
+++ b/pkg/services/grpcserver/interceptors/tracing.go
@@ -3,34 +3,15 @@ package interceptors
 import (
 	"context"
 
+	"github.com/grafana/grafana/pkg/infra/tracing"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/propagation"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/metadata"
-
-	"github.com/grafana/grafana/pkg/infra/tracing"
 )
 
 const tracingPrefix = "gRPC Server "
 
-func TracingUnaryInterceptor(tracer tracing.Tracer) grpc.UnaryServerInterceptor {
-	return func(
-		ctx context.Context,
-		req any,
-		info *grpc.UnaryServerInfo,
-		handler grpc.UnaryHandler,
-	) (resp any, err error) {
-		if md, ok := metadata.FromIncomingContext(ctx); ok {
-			ctx = otel.GetTextMapPropagator().Extract(ctx, propagation.HeaderCarrier(md))
-		}
-
-		ctx, span := tracer.Start(ctx, tracingPrefix+info.FullMethod)
-		defer span.End()
-		resp, err = handler(ctx, req)
-		return resp, err
-	}
-}
-
 func TracingStreamInterceptor(tracer tracing.Tracer) grpc.StreamServerInterceptor {
 	return func(srv any, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
 		ctx := stream.Context()
diff --git a/pkg/services/grpcserver/service.go b/pkg/services/grpcserver/service.go
index f4777b44356..611746ee88b 100644
--- a/pkg/services/grpcserver/service.go
+++ b/pkg/services/grpcserver/service.go
@@ -11,6 +11,7 @@ import (
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	grpcAuth "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/auth"
 	"github.com/prometheus/client_golang/prometheus"
+	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
 
@@ -74,9 +75,9 @@ func ProvideService(cfg *setting.Cfg, features featuremgmt.FeatureToggles, authe
 	// services which implement ServiceAuthFuncOverride interface.
 	// See https://github.com/grpc-ecosystem/go-grpc-middleware/blob/main/interceptors/auth/auth.go#L30.
 	opts = append(opts, []grpc.ServerOption{
+		grpc.StatsHandler(otelgrpc.NewServerHandler()),
 		grpc.ChainUnaryInterceptor(
 			grpcAuth.UnaryServerInterceptor(authenticator.Authenticate),
-			interceptors.TracingUnaryInterceptor(tracer),
 			interceptors.LoggingUnaryInterceptor(s.cfg, s.logger), // needs to be registered after tracing interceptor to get trace id
 			middleware.UnaryServerInstrumentInterceptor(grpcRequestDuration),
 		),