From 1aaf523ed2625b9d703ae1851f644a545d5d930b Mon Sep 17 00:00:00 2001
From: Jorge Alberto Cricelli <jorge.cricelli@vemobility.net>
Date: Tue, 10 Sep 2024 10:30:29 -0300
Subject: [PATCH] feat(metrics): custom histogram buckets configuration (#489)

---
 components/metrics/builder.go |  7 +++++++
 components/metrics/handler.go | 10 +++++++---
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/components/metrics/builder.go b/components/metrics/builder.go
index 201b26195..1ed1b8bb9 100644
--- a/components/metrics/builder.go
+++ b/components/metrics/builder.go
@@ -24,10 +24,16 @@ type PrometheusMetricsBuilder struct {
 
 	Namespace string
 	Subsystem string
+	// PublishBuckets defines the histogram buckets for publish time histogram, defaulted if nil.
+	PublishBuckets []float64
+	// HandlerBuckets defines the histogram buckets for handle execution time histogram, defaulted to watermill's default.
+	HandlerBuckets []float64
 }
 
 // AddPrometheusRouterMetrics is a convenience function that acts on the message router to add the metrics middleware
 // to all its handlers. The handlers' publishers and subscribers are also decorated.
+// The default buckets are used for the handler execution time histogram (use your own provisioning
+// with NewRouterMiddlewareWithConfig if needed).
 func (b PrometheusMetricsBuilder) AddPrometheusRouterMetrics(r *message.Router) {
 	r.AddPublisherDecorators(b.DecoratePublisher)
 	r.AddSubscriberDecorators(b.DecorateSubscriber)
@@ -48,6 +54,7 @@ func (b PrometheusMetricsBuilder) DecoratePublisher(pub message.Publisher) (mess
 			Subsystem: b.Subsystem,
 			Name:      "publish_time_seconds",
 			Help:      "The time that a publishing attempt (success or not) took in seconds",
+			Buckets:   b.PublishBuckets,
 		},
 		publisherLabelKeys,
 	))
diff --git a/components/metrics/handler.go b/components/metrics/handler.go
index 92d11ba19..73eb47afa 100644
--- a/components/metrics/handler.go
+++ b/components/metrics/handler.go
@@ -15,9 +15,9 @@ var (
 		labelSuccess,
 	}
 
-	// handlerExecutionTimeBuckets are one order of magnitude smaller than default buckets (5ms~10s),
+	// defaultHandlerExecutionTimeBuckets are one order of magnitude smaller than default buckets (5ms~10s),
 	// because the handler execution times are typically shorter (µs~ms range).
-	handlerExecutionTimeBuckets = []float64{
+	defaultHandlerExecutionTimeBuckets = []float64{
 		0.0005,
 		0.001,
 		0.0025,
@@ -64,13 +64,17 @@ func (b PrometheusMetricsBuilder) NewRouterMiddleware() HandlerPrometheusMetrics
 	var err error
 	m := HandlerPrometheusMetricsMiddleware{}
 
+	if b.HandlerBuckets == nil {
+		b.HandlerBuckets = defaultHandlerExecutionTimeBuckets
+	}
+
 	m.handlerExecutionTimeSeconds, err = b.registerHistogramVec(prometheus.NewHistogramVec(
 		prometheus.HistogramOpts{
 			Namespace: b.Namespace,
 			Subsystem: b.Subsystem,
 			Name:      "handler_execution_time_seconds",
 			Help:      "The total time elapsed while executing the handler function in seconds",
-			Buckets:   handlerExecutionTimeBuckets,
+			Buckets:   b.HandlerBuckets,
 		},
 		handlerLabelKeys,
 	))