feat(storage): WithMeterProvider allows custom meter provider configuration (#12668)

cojenco · BrennaEpp · web-flow · commit 7f574b01e0b4 · 2025-08-26T11:56:01.000-07:00
* feat: WithMeterProvider allows custom meter provider configurated

* update meterProvider to metricsContext

* Update storage/grpc_metrics.go

Co-authored-by: Brenna N Epp &lt;brennanatashaepp@gmail.com&gt;

* update docstrings and tests

---------

Co-authored-by: Brenna N Epp &lt;brennanatashaepp@gmail.com&gt;
diff --git a/storage/experimental/experimental.go b/storage/experimental/experimental.go
@@ -33,17 +33,25 @@ import (
 // It sets how often to emit metrics [metric.WithInterval] when using
 // [metric.NewPeriodicReader]
 // When using Cloud Monitoring interval must be at minimum 1 [time.Minute].
+// This option is ignored if WithMeterProvider is also set.
 func WithMetricInterval(metricInterval time.Duration) option.ClientOption {
 	return internal.WithMetricInterval.(func(time.Duration) option.ClientOption)(metricInterval)
 }
 
 // WithMetricExporter provides a [option.ClientOption] that may be passed to [storage.NewGRPCClient].
 // Set an alternate client-side metric Exporter to emit metrics through.
-// Must implement [metric.Exporter]
+// Must implement [metric.Exporter]. This option is ignored if WithMeterProvider is also set.
 func WithMetricExporter(ex *metric.Exporter) option.ClientOption {
 	return internal.WithMetricExporter.(func(*metric.Exporter) option.ClientOption)(ex)
 }
 
+// WithMeterProvider provides a [option.ClientOption] that may be passed to [storage.NewGRPCClient].
+// Set an alternate client-side meter provider to emit metrics through.
+// This option takes precedence over WithMetricExporter and WithMetricInterval.
+func WithMeterProvider(mp *metric.MeterProvider) option.ClientOption {
+	return internal.WithMeterProvider.(func(*metric.MeterProvider) option.ClientOption)(mp)
+}
+
 // WithReadStallTimeout provides a [option.ClientOption] that may be passed to [storage.NewClient].
 // It enables the client to retry stalled requests when starting a download from
 // Cloud Storage. If the timeout elapses with no response from the server, the request
diff --git a/storage/grpc_client.go b/storage/grpc_client.go
@@ -125,9 +125,11 @@ func enableClientMetrics(ctx context.Context, s *settings, config storageConfig)
 		project = c.ProjectID
 	}
 	metricsContext, err := newGRPCMetricContext(ctx, metricsConfig{
-		project:      project,
-		interval:     config.metricInterval,
-		manualReader: config.manualReader},
+		project:       project,
+		interval:      config.metricInterval,
+		manualReader:  config.manualReader,
+		meterProvider: config.meterProvider,
+	},
 	)
 	if err != nil {
 		return nil, fmt.Errorf("gRPC Metrics: %w", err)
diff --git a/storage/grpc_metrics.go b/storage/grpc_metrics.go
@@ -126,6 +126,7 @@ type metricsConfig struct {
 	project         string
 	interval        time.Duration
 	customExporter  *metric.Exporter
+	meterProvider   *metric.MeterProvider
 	manualReader    *metric.ManualReader // used by tests
 	disableExporter bool                 // used by tests disables exports
 	resourceOpts    []resource.Option    // used by tests
@@ -172,7 +173,10 @@ func newGRPCMetricContext(ctx context.Context, cfg metricsConfig) (*metricsConte
 		meterOpts = append(meterOpts, metric.WithReader(
 			metric.NewPeriodicReader(&exporterLogSuppressor{Exporter: exporter}, metric.WithInterval(interval))))
 	}
-	provider := metric.NewMeterProvider(meterOpts...)
+	provider := cfg.meterProvider
+	if provider == nil {
+		provider = metric.NewMeterProvider(meterOpts...)
+	}
 	mo := opentelemetry.MetricsOptions{
 		MeterProvider: provider,
 		Metrics: stats.NewMetrics(
diff --git a/storage/grpc_metrics_test.go b/storage/grpc_metrics_test.go
@@ -201,3 +201,66 @@ type failingExporter struct {
 func (f *failingExporter) Export(ctx context.Context, rm *metricdata.ResourceMetrics) error {
 	return fmt.Errorf("PermissionDenied")
 }
+
+func TestNewGRPCMetricContextWithCustomProvider(t *testing.T) {
+	// Create a manual reader for collecting metrics.
+	mr := metric.NewManualReader()
+
+	// Create a custom meter provider with the manual reader.
+	customProvider := metric.NewMeterProvider(
+		metric.WithReader(mr),
+	)
+	ctx := context.Background()
+
+	cfg := metricsConfig{
+		project:         "project-id",
+		meterProvider:   customProvider, // use custom provider
+		disableExporter: true,           // disable since this is a unit test
+	}
+
+	mc, err := newGRPCMetricContext(ctx, cfg)
+	if err != nil {
+		t.Errorf("newGRPCMetricContext: %v", err)
+	}
+	defer mc.close()
+
+	// Verify the provider in metricsContext is our custom provider.
+	if mc.provider != customProvider {
+		t.Errorf("metricsContext.provider = %v, want %v", mc.provider, customProvider)
+	}
+
+	// Now we can collect metrics from the manual reader that's registered with our custom provider.
+	rm := metricdata.ResourceMetrics{}
+	if err := mr.Collect(ctx, &rm); err != nil {
+		t.Errorf("ManualReader.Collect: %v", err)
+	}
+}
+
+func TestNewGRPCMetricContextWithCustomProviderExporter(t *testing.T) {
+	// 1. Create a custom meter provider.
+	mr := metric.NewManualReader()
+	customProvider := metric.NewMeterProvider(
+		metric.WithReader(mr),
+	)
+	// 2. Create a custom `failingExporter`.
+	failingExp := &failingExporter{}
+	customExporter := metric.Exporter(failingExp)
+
+	// 3. Create the metrics context with both the custom provider and the custom exporter.
+	// The custom provider should take precedence and prevent the failing exporter from being used.
+	ctx := context.Background()
+	cfg := metricsConfig{
+		project:        "project-id",
+		meterProvider:  customProvider,
+		customExporter: &customExporter,
+	}
+
+	mc, err := newGRPCMetricContext(ctx, cfg)
+	if err != nil {
+		t.Errorf("newGRPCMetricContext: %v", err)
+	}
+	// 4. Verification: The `failingExporter` should not have been used.
+	// We expect the close to succeed, which means the failing exporter was not registered.
+	// If the failing exporter had been used, mc.close() would likely panic or return an error.
+	mc.close()
+}
diff --git a/storage/internal/experimental.go b/storage/internal/experimental.go
@@ -26,6 +26,10 @@ var (
 	// Set an alternate client-side metric Exporter to emit metrics through.
 	WithMetricExporter any // func (*metric.Exporter) option.ClientOption
 
+	// WithMeterProvider is a function which is implemented by storage package.
+	// Set an alternate client-side meter provider to emit metrics through.
+	WithMeterProvider any // func (*metric.MeterProvider) option.ClientOption
+
 	// WithReadStallTimeout is a function which is implemented by storage package.
 	// It takes ReadStallTimeoutConfig as inputs and returns a option.ClientOption.
 	WithReadStallTimeout any // func (*ReadStallTimeoutConfig) option.ClientOption
diff --git a/storage/option.go b/storage/option.go
@@ -39,6 +39,7 @@ func init() {
 	// initialize experimental options
 	storageinternal.WithMetricExporter = withMetricExporter
 	storageinternal.WithMetricInterval = withMetricInterval
+	storageinternal.WithMeterProvider = withMeterProvider
 	storageinternal.WithReadStallTimeout = withReadStallTimeout
 	storageinternal.WithGRPCBidiReads = withGRPCBidiReads
 	storageinternal.WithZonalBucketAPIs = withZonalBucketAPIs
@@ -81,6 +82,7 @@ type storageConfig struct {
 	disableClientMetrics   bool
 	metricExporter         *metric.Exporter
 	metricInterval         time.Duration
+	meterProvider          *metric.MeterProvider
 	manualReader           *metric.ManualReader
 	readStallTimeoutConfig *experimental.ReadStallTimeoutConfig
 	grpcBidiReads          bool
@@ -203,6 +205,20 @@ type withTestMetricReaderConfig struct {
 	metricReader *metric.ManualReader
 }
 
+type withMeterProviderConfig struct {
+	internaloption.EmbeddableAdapter
+	// meter provider override
+	meterProvider *metric.MeterProvider
+}
+
+func withMeterProvider(provider *metric.MeterProvider) option.ClientOption {
+	return &withMeterProviderConfig{meterProvider: provider}
+}
+
+func (w *withMeterProviderConfig) ApplyStorageOpt(c *storageConfig) {
+	c.meterProvider = w.meterProvider
+}
+
 func withTestMetricReader(ex *metric.ManualReader) option.ClientOption {
 	return &withTestMetricReaderConfig{metricReader: ex}
 }
