From 3f8a34303ce1b9a7cb0879a7f88f031d50e431e9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jean-Philippe=20Qu=C3=A9m=C3=A9ner?=
 <JohnnyQQQQ@users.noreply.github.com>
Date: Wed, 19 Feb 2025 19:58:10 +0100
Subject: [PATCH] feat(unified-storage): enhance gRPC client with dskit
 (#100993)

---
 pkg/storage/unified/client.go | 52 +++++++++++++++++++++++++++++++----
 1 file changed, 47 insertions(+), 5 deletions(-)

diff --git a/pkg/storage/unified/client.go b/pkg/storage/unified/client.go
index 7b34b95ab92..75febc247e0 100644
--- a/pkg/storage/unified/client.go
+++ b/pkg/storage/unified/client.go
@@ -6,6 +6,7 @@ import (
 	"path/filepath"
 
 	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
 	"gocloud.dev/blob/fileblob"
 	"google.golang.org/grpc"
@@ -13,6 +14,8 @@ import (
 
 	authnlib "github.com/grafana/authlib/authn"
 	"github.com/grafana/authlib/types"
+	"github.com/grafana/dskit/grpcclient"
+	"github.com/grafana/dskit/middleware"
 
 	infraDB "github.com/grafana/grafana/pkg/infra/db"
 	"github.com/grafana/grafana/pkg/infra/tracing"
@@ -39,6 +42,10 @@ type Options struct {
 	Docs     resource.DocumentBuilderSupplier
 }
 
+type ClientMetrics struct {
+	requestDuration *prometheus.HistogramVec
+}
+
 // This adds a UnifiedStorage client into the wire dependency tree
 func ProvideUnifiedStorageClient(opts *Options) (resource.ResourceClient, error) {
 	// See: apiserver.ApplyGrafanaConfig(cfg, features, o)
@@ -104,11 +111,8 @@ func newClient(opts options.StorageOptions,
 			return nil, fmt.Errorf("expecting address for storage_type: %s", opts.StorageType)
 		}
 
-		// Create a connection to the gRPC server
-		conn, err := grpc.NewClient(opts.Address,
-			grpc.WithStatsHandler(otelgrpc.NewClientHandler()),
-			grpc.WithTransportCredentials(insecure.NewCredentials()),
-		)
+		// Create a connection to the gRPC server.
+		conn, err := grpcConn(opts.Address, reg)
 		if err != nil {
 			return nil, err
 		}
@@ -153,3 +157,41 @@ func newResourceClient(conn *grpc.ClientConn, cfg *setting.Cfg, features feature
 	}
 	return resource.NewRemoteResourceClient(tracer, conn, clientCfgMapping(grpcutils.ReadGrpcClientConfig(cfg)), cfg.Env == setting.Dev)
 }
+
+// grpcConn creates a new gRPC connection to the provided address.
+func grpcConn(address string, reg prometheus.Registerer) (*grpc.ClientConn, error) {
+	// This works for now as the Provide function is only called once during startup.
+	// We might eventually want to tie this factory to a struct for more runtime control.
+	metrics := ClientMetrics{
+		requestDuration: promauto.With(reg).NewHistogramVec(prometheus.HistogramOpts{
+			Name:    "resource_server_client_request_duration_seconds",
+			Help:    "Time spent executing requests to the resource server.",
+			Buckets: prometheus.ExponentialBuckets(0.008, 4, 7),
+		}, []string{"operation", "status_code"}),
+	}
+
+	// Report gRPC status code errors as labels.
+	var instrumentationOptions []middleware.InstrumentationOption
+	instrumentationOptions = append(instrumentationOptions, middleware.ReportGRPCStatusOption)
+	unary, stream := grpcclient.Instrument(metrics.requestDuration, instrumentationOptions...)
+
+	// We can later pass in the gRPC config here, i.e. to set MaxRecvMsgSize etc.
+	cfg := grpcclient.Config{}
+	opts, err := cfg.DialOption(unary, stream)
+	if err != nil {
+		return nil, fmt.Errorf("could not instrument grpc client: %w", err)
+	}
+
+	opts = append(opts, grpc.WithStatsHandler(otelgrpc.NewClientHandler()))
+	opts = append(opts, grpc.WithTransportCredentials(insecure.NewCredentials()))
+
+	// Use round_robin to balance requests more evenly over the available Storage server.
+	opts = append(opts, grpc.WithDefaultServiceConfig(`{"loadBalancingPolicy":"round_robin"}`))
+
+	// Disable looking up service config from TXT DNS records.
+	// This reduces the number of requests made to the DNS servers.
+	opts = append(opts, grpc.WithDisableServiceConfig())
+
+	// Create a connection to the gRPC server
+	return grpc.NewClient(address, opts...)
+}