From 8df333bb08fc677fe12f216ebd2546ab9ac5e5aa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Josef=20Kol=C3=A1=C5=99?= <josef.kolar@firma.seznam.cz>
Date: Fri, 6 Sep 2024 12:25:57 +0200
Subject: [PATCH] feat(metrics): custom histogram buckets configuration (#489)

---
 components/metrics/builder.go |  7 +++++++
 components/metrics/handler.go | 10 +++++++---
 docs/content/docs/metrics.md  |  4 ++++
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/components/metrics/builder.go b/components/metrics/builder.go
index c70e4755a..fb2cb0838 100644
--- a/components/metrics/builder.go
+++ b/components/metrics/builder.go
@@ -22,10 +22,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)
@@ -46,6 +52,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 b9240426f..4e3e5ea1a 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,
 	))
diff --git a/docs/content/docs/metrics.md b/docs/content/docs/metrics.md
index 2d59e6bac..8491d27fe 100644
--- a/docs/content/docs/metrics.md
+++ b/docs/content/docs/metrics.md
@@ -35,6 +35,10 @@ Example use of `AddPrometheusRouterMetrics`:
 
 In the snippet above, we have left the `namespace` and `subsystem` arguments empty. The Prometheus client library [uses these](https://godoc.org/github.com/prometheus/client_golang/prometheus#BuildFQName) to prefix the metric names. You may want to use namespace or subsystem, but be aware that this will impact the metric names and you will have to adjust the Grafana dashboard accordingly.
 
+The `PrometheusMetricsBuilder` allows for custom configuration of histogram buckets by setting the `PublishBuckets` or `HandlerBuckets` field.
+If `HandlerBuckets` is not provided, default watermill's values will be used, which are one order of magnitude smaller than default buckets (5ms~10s), because the handler execution times are typically shorter (µs~ms range).
+For `PublishBuckets`, the default values are the same as the default Prometheus buckets (5ms~10s).
+
 Standalone publishers and subscribers may also be decorated through the use of dedicated methods of `PrometheusMetricBuilder`:
 
 {{% render-md %}}