From 5b7888efbc26341c7541b62f198bcbd6ad77a13b Mon Sep 17 00:00:00 2001
From: jeremy <jbreiding@users.noreply.github.com>
Date: Thu, 9 Jun 2022 00:56:38 +0000
Subject: [PATCH] Adding serveroption for metrics.MetricHandler

---
 temporal/fx.go             | 20 +++++++++++++-------
 temporal/server_option.go  |  9 +++++++++
 temporal/server_options.go |  2 ++
 3 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/temporal/fx.go b/temporal/fx.go
index c012db60e55..57301718882 100644
--- a/temporal/fx.go
+++ b/temporal/fx.go
@@ -122,6 +122,7 @@ type (
 		TLSConfigProvider       encryption.TLSConfigProvider
 		EsConfig                *esclient.Config
 		EsClient                esclient.Client
+		MetricHandler           metrics.MetricHandler
 	}
 )
 
@@ -180,16 +181,20 @@ func ServerOptionsProvider(opts []ServerOption) (serverOptionsProvider, error) {
 		clientFactoryProvider = client.NewFactoryProvider()
 	}
 
+	// MetricHandler
+	handler := so.metricHandler
+	if handler == nil {
+		if so.config.Global.Metrics != nil {
+			handler = metrics.MetricHandlerFromConfig(logger, so.config.Global.Metrics)
+		} else {
+			handler = metrics.NoopMetricHandler
+		}
+	}
+
 	// ServerReporter
 	serverReporter := so.metricsReporter
 	if serverReporter == nil {
-		if so.config.Global.Metrics == nil {
-			logger.Warn("no metrics config provided, using noop reporter")
-			serverReporter = metrics.NoopReporter
-		} else {
-			handler := metrics.MetricHandlerFromConfig(logger, so.config.Global.Metrics)
-			serverReporter = metrics.NewEventsReporter(handler)
-		}
+		serverReporter = metrics.NewEventsReporter(handler)
 	}
 
 	// MetricsClient
@@ -275,6 +280,7 @@ func ServerOptionsProvider(opts []ServerOption) (serverOptionsProvider, error) {
 		TLSConfigProvider:       tlsConfigProvider,
 		EsConfig:                esConfig,
 		EsClient:                esClient,
+		MetricHandler:           handler,
 	}, nil
 }
 
diff --git a/temporal/server_option.go b/temporal/server_option.go
index da842692aef..d0de25bbdd5 100644
--- a/temporal/server_option.go
+++ b/temporal/server_option.go
@@ -129,6 +129,7 @@ func WithAudienceGetter(audienceGetter func(cfg *config.Config) authorization.JW
 // handler, err := NewCustomHandlerImplementation(logger)
 // reporter, err2 := metrics.NewEventsReporter(handler)
 // server := temporal.NewServer(temporal.WithCustomMetricsReporter(repoter))
+// Deprecated use WithCustomMetricHandler
 func WithCustomMetricsReporter(reporter metrics.Reporter) ServerOption {
 	return newApplyFuncContainer(func(s *serverOptions) {
 		s.metricsReporter = reporter
@@ -190,3 +191,11 @@ func WithChainedFrontendGrpcInterceptors(
 		s.customInterceptors = interceptors
 	})
 }
+
+// WithCustomerMetricsHandler sets a custom implementation of the metrics.MetricHandler interface
+// metrics.MetricHandler is the base interface for publishing metric events
+func WithCustomMetricHandler(handler metrics.MetricHandler) ServerOption {
+	return newApplyFuncContainer(func(s *serverOptions) {
+		s.metricHandler = handler
+	})
+}
diff --git a/temporal/server_options.go b/temporal/server_options.go
index 5a373407b0f..ad89eb20bcf 100644
--- a/temporal/server_options.go
+++ b/temporal/server_options.go
@@ -68,6 +68,7 @@ type (
 		clientFactoryProvider      client.FactoryProvider
 		searchAttributesMapper     searchattribute.Mapper
 		customInterceptors         []grpc.UnaryServerInterceptor
+		metricHandler              metrics.MetricHandler
 	}
 )
 
@@ -127,6 +128,7 @@ func (so *serverOptions) validateConfig() error {
 	}
 	return nil
 }
+
 func isValidService(service string) bool {
 	for _, s := range Services {
 		if s == service {