feat(unified-storage): add grpc connection pooling (#102575)

go.mod

@@ -559,6 +559,7 @@ require (
 )
 
 require (
+	github.com/1NCE-GmbH/grpc-go-pool v0.0.0-20231117122434-2a5bb974daa2 // indirect
 	github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.49.0 // indirect
 	github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.49.0 // indirect
 	github.com/RoaringBitmap/roaring/v2 v2.4.5 // indirect

go.sum

@@ -643,6 +643,8 @@ filippo.io/age v1.2.1 h1:X0TZjehAZylOIj4DubWYU1vWQxv9bJpo+Uu2/LGhi1o=
 filippo.io/age v1.2.1/go.mod h1:JL9ew2lTN+Pyft4RiNGguFfOpewKwSHm5ayKD/A4004=
 gioui.org v0.0.0-20210308172011-57750fc8a0a6/go.mod h1:RSH6KIUZ0p2xy5zHDxgAM4zumjgTw83q2ge/PI+yyw8=
 git.sr.ht/~sbinet/gg v0.3.1/go.mod h1:KGYtlADtqsqANL9ueOFkWymvzUvLMQllU5Ixo+8v3pc=
+github.com/1NCE-GmbH/grpc-go-pool v0.0.0-20231117122434-2a5bb974daa2 h1:qFYgLH2zZe3WHpQgUrzeazC+ebDebwAQqS9yE1cP5Bs=
+github.com/1NCE-GmbH/grpc-go-pool v0.0.0-20231117122434-2a5bb974daa2/go.mod h1:09/ALd1AXCTCOfcJYD8+jIYKmFmi6PVCkTsipC18F7E=
 github.com/Azure/azure-pipeline-go v0.2.3 h1:7U9HBg1JFK3jHl5qmo4CTZKFTVgMwdFHMVtCdfBE21U=
 github.com/Azure/azure-pipeline-go v0.2.3/go.mod h1:x841ezTBIMG6O3lAcl8ATHnsOPVl2bqk7S3ta6S6u4k=
 github.com/Azure/azure-sdk-for-go v23.2.0+incompatible/go.mod h1:9XXNKU+eRnpl9moKnB4QOLf1HestfXbmab5FXxiDBjc=

packages/grafana-data/src/types/featureToggles.gen.ts

@@ -1060,4 +1060,8 @@ export interface FeatureToggles {
   * Enables the unified storage history pruner
   */
   unifiedStorageHistoryPruner?: boolean;
+  /**
+  * Enables the unified storage grpc connection pool
+  */
+  unifiedStorageGrpcConnectionPool?: boolean;
 }

pkg/services/featuremgmt/registry.go

@@ -1830,6 +1830,14 @@ var (
 			HideFromAdminPage: true,
 			HideFromDocs:      true,
 		},
+		{
+			Name:              "unifiedStorageGrpcConnectionPool",
+			Description:       "Enables the unified storage grpc connection pool",
+			Stage:             FeatureStageExperimental,
+			Owner:             grafanaSearchAndStorageSquad,
+			HideFromAdminPage: true,
+			HideFromDocs:      true,
+		},
 	}
 )
 

pkg/services/featuremgmt/toggles_gen.csv

@@ -241,3 +241,4 @@ extraLanguages,experimental,@grafana/grafana-frontend-platform,false,false,true
 noBackdropBlur,experimental,@grafana/grafana-frontend-platform,false,false,true
 alertingMigrationUI,experimental,@grafana/alerting-squad,false,false,true
 unifiedStorageHistoryPruner,experimental,@grafana/search-and-storage,false,false,false
+unifiedStorageGrpcConnectionPool,experimental,@grafana/search-and-storage,false,false,false

pkg/services/featuremgmt/toggles_gen.go

@@ -974,4 +974,8 @@ const (
 	// FlagUnifiedStorageHistoryPruner
 	// Enables the unified storage history pruner
 	FlagUnifiedStorageHistoryPruner = "unifiedStorageHistoryPruner"
+
+	// FlagUnifiedStorageGrpcConnectionPool
+	// Enables the unified storage grpc connection pool
+	FlagUnifiedStorageGrpcConnectionPool = "unifiedStorageGrpcConnectionPool"
 )

pkg/services/featuremgmt/toggles_gen.json

@@ -4203,6 +4203,20 @@
         "codeowner": "@grafana/search-and-storage"
       }
     },
+    {
+      "metadata": {
+        "name": "unifiedStorageGrpcConnectionPool",
+        "resourceVersion": "1742549790491",
+        "creationTimestamp": "2025-03-21T09:36:30Z"
+      },
+      "spec": {
+        "description": "Enables the unified storage grpc connection pool",
+        "stage": "experimental",
+        "codeowner": "@grafana/search-and-storage",
+        "hideFromAdminPage": true,
+        "hideFromDocs": true
+      }
+    },
     {
       "metadata": {
         "name": "unifiedStorageHistoryPruner",

pkg/storage/unified/client.go

@@ -117,10 +117,31 @@ 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 := GrpcConn(opts.Address, reg)
-		if err != nil {
-			return nil, err
+		var (
+			conn    grpc.ClientConnInterface
+			err     error
+			metrics = newClientMetrics(reg)
+		)
+		// Create either a connection pool or a single connection.
+		// The connection pool __can__ be useful when connection to
+		// server side load balancers like kube-proxy.
+		if features.IsEnabledGlobally(featuremgmt.FlagUnifiedStorageGrpcConnectionPool) {
+			conn, err = newPooledConn(&poolOpts{
+				initialCapacity: 3,
+				maxCapacity:     6,
+				idleTimeout:     time.Minute,
+				factory: func() (*grpc.ClientConn, error) {
+					return grpcConn(opts.Address, metrics)
+				},
+			})
+			if err != nil {
+				return nil, err
+			}
+		} else {
+			conn, err = grpcConn(opts.Address, metrics)
+			if err != nil {
+				return nil, err
+			}
 		}
 
 		// Create a client instance
@@ -144,7 +165,7 @@ func newClient(opts options.StorageOptions,
 	}
 }
 
-func newResourceClient(conn *grpc.ClientConn, cfg *setting.Cfg, features featuremgmt.FeatureToggles, tracer tracing.Tracer) (resource.ResourceClient, error) {
+func newResourceClient(conn grpc.ClientConnInterface, cfg *setting.Cfg, features featuremgmt.FeatureToggles, tracer tracing.Tracer) (resource.ResourceClient, error) {
 	if !features.IsEnabledGlobally(featuremgmt.FlagAppPlatformGrpcClientAuth) {
 		return resource.NewLegacyResourceClient(conn), nil
 	}
@@ -160,22 +181,8 @@ func newResourceClient(conn *grpc.ClientConn, cfg *setting.Cfg, features feature
 	})
 }
 
-// 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 tight 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"}),
-		requestRetries: promauto.With(reg).NewCounterVec(prometheus.CounterOpts{
-			Name: "resource_server_client_request_retries_total",
-			Help: "Total number of retries for requests to the resource server.",
-		}, []string{"operation"}),
-	}
-
+// grpcConn creates a new gRPC connection to the provided address.
+func grpcConn(address string, metrics *clientMetrics) (*grpc.ClientConn, error) {
 	// Report gRPC status code errors as labels.
 	unary, stream := instrument(metrics.requestDuration, middleware.ReportGRPCStatusOption)
 
@@ -212,6 +219,12 @@ func GrpcConn(address string, reg prometheus.Registerer) (*grpc.ClientConn, erro
 	return grpc.NewClient(address, opts...)
 }
 
+// GrpcConn is the public constructor that can be used for testing.
+func GrpcConn(address string, reg prometheus.Registerer) (*grpc.ClientConn, error) {
+	metrics := newClientMetrics(reg)
+	return grpcConn(address, metrics)
+}
+
 // instrument is the same as grpcclient.Instrument but without the middleware.ClientUserHeaderInterceptor
 // and middleware.StreamClientUserHeaderInterceptor as we don't need them.
 func instrument(requestDuration *prometheus.HistogramVec, instrumentationLabelOptions ...middleware.InstrumentationOption) ([]grpc.UnaryClientInterceptor, []grpc.StreamClientInterceptor) {
@@ -223,3 +236,19 @@ func instrument(requestDuration *prometheus.HistogramVec, instrumentationLabelOp
 			middleware.StreamClientInstrumentInterceptor(requestDuration, instrumentationLabelOptions...),
 		}
 }
+
+func newClientMetrics(reg prometheus.Registerer) *clientMetrics {
+	// This works for now as the Provide function is only called once during startup.
+	// We might eventually want to tight this factory to a struct for more runtime control.
+	return &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"}),
+		requestRetries: promauto.With(reg).NewCounterVec(prometheus.CounterOpts{
+			Name: "resource_server_client_request_retries_total",
+			Help: "Total number of retries for requests to the resource server.",
+		}, []string{"operation"}),
+	}
+}

pkg/storage/unified/grpc_pool.go

@@ -0,0 +1,82 @@
+package unified
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	grpcpool "github.com/1NCE-GmbH/grpc-go-pool"
+	"google.golang.org/grpc"
+)
+
+// pooledClientConn implements grpc.ClientConnInterface using a connection from a pool.
+type pooledClientConn struct {
+	pool *grpcpool.Pool
+	// For streaming we want to keep a single connection, as otherwise we saturate the pool.
+	// Streaming should only be used for watching.
+	streamConn grpc.ClientConnInterface
+}
+
+// Invoke implements the grpc.ClientConnInterface.Invoke method.
+func (pc *pooledClientConn) Invoke(ctx context.Context, method string, args interface{}, reply interface{}, opts ...grpc.CallOption) error {
+	conn, err := pc.pool.Get(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to create grpc conn in pooled client: %w", err)
+	}
+	// Return connection to pool when done.
+	defer func() {
+		_ = conn.Close()
+	}()
+	return conn.ClientConn.Invoke(ctx, method, args, reply, opts...)
+}
+
+// NewStream implements the grpc.ClientConnInterface.NewStream method.
+func (pc *pooledClientConn) NewStream(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
+	stream, err := pc.streamConn.NewStream(ctx, desc, method, opts...)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create grpc stream in pooled client: %w", err)
+	}
+	return stream, nil
+}
+
+type poolOpts struct {
+	initialCapacity int
+	maxCapacity     int
+	idleTimeout     time.Duration
+	factory         func() (*grpc.ClientConn, error)
+}
+
+func (opts *poolOpts) validate() error {
+	if opts.initialCapacity <= 0 {
+		return errors.New("initial capacity is required")
+	}
+	if opts.maxCapacity < opts.initialCapacity {
+		return errors.New("max capacity is less than initial capacity")
+	}
+	if opts.idleTimeout <= 0 {
+		return errors.New("idle timeout is required")
+	}
+	if opts.factory == nil {
+		return errors.New("factory is required")
+	}
+	return nil
+}
+
+func newPooledConn(opts *poolOpts) (grpc.ClientConnInterface, error) {
+	if err := opts.validate(); err != nil {
+		return nil, fmt.Errorf("failed to validate grpc connection pool options: %w", err)
+	}
+	pool, err := grpcpool.New(opts.factory, opts.initialCapacity, opts.maxCapacity, opts.idleTimeout)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create grpc connection pool: %w", err)
+	}
+	streamConn, err := opts.factory()
+	if err != nil {
+		return nil, fmt.Errorf("failed to create groc streaming connection: %w", err)
+	}
+	return &pooledClientConn{
+		pool:       pool,
+		streamConn: streamConn,
+	}, nil
+}

pkg/storage/unified/resource/client.go

@@ -14,6 +14,7 @@ import (
 
 	authnlib "github.com/grafana/authlib/authn"
 	"github.com/grafana/authlib/types"
+
 	"github.com/grafana/grafana/pkg/apimachinery/identity"
 
 	"github.com/grafana/grafana/pkg/infra/tracing"
@@ -40,7 +41,7 @@ type resourceClient struct {
 	DiagnosticsClient
 }
 
-func NewLegacyResourceClient(channel *grpc.ClientConn) ResourceClient {
+func NewLegacyResourceClient(channel grpc.ClientConnInterface) ResourceClient {
 	cc := grpchan.InterceptClientConn(channel, grpcUtils.UnaryClientInterceptor, grpcUtils.StreamClientInterceptor)
 	return &resourceClient{
 		ResourceStoreClient:      NewResourceStoreClient(cc),
@@ -99,7 +100,7 @@ type RemoteResourceClientConfig struct {
 	AllowInsecure    bool
 }
 
-func NewRemoteResourceClient(tracer tracing.Tracer, conn *grpc.ClientConn, cfg RemoteResourceClientConfig) (ResourceClient, error) {
+func NewRemoteResourceClient(tracer tracing.Tracer, conn grpc.ClientConnInterface, cfg RemoteResourceClientConfig) (ResourceClient, error) {
 	exchangeOpts := []authnlib.ExchangeClientOpts{}
 
 	if cfg.AllowInsecure {