From 21ac48b2a0848739502cc81632a448a793038852 Mon Sep 17 00:00:00 2001
From: Kennedy Oliveira <kennedyroliveira@gmail.com>
Date: Fri, 8 Nov 2024 00:21:39 -0300
Subject: [PATCH] Allow custom Registry on HttpMetricsLayerBuilder

---
 src/lib.rs | 31 ++++++++++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/src/lib.rs b/src/lib.rs
index d579e02..e9499b0 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -244,6 +244,7 @@ pub struct HttpMetricsLayerBuilder {
     skipper: PathSkipper,
     is_tls: bool,
     exporter: Option<String>,
+    registry: Option<Registry>,
 }
 
 impl Default for HttpMetricsLayerBuilder {
@@ -257,6 +258,7 @@ impl Default for HttpMetricsLayerBuilder {
             skipper: PathSkipper::default(),
             is_tls: false,
             exporter: Some("prometheus".to_string()),
+            registry: None,
         }
     }
 }
@@ -301,6 +303,11 @@ impl HttpMetricsLayerBuilder {
         self
     }
 
+    pub fn with_registry(mut self, registry: Registry) -> Self {
+        self.registry = Some(registry);
+        self
+    }
+
     pub fn build(self) -> HttpMetricsLayer {
         let mut resource = vec![];
 
@@ -426,7 +433,9 @@ impl HttpMetricsLayerBuilder {
     }
 
     fn build_prometheus(&self) -> (Registry, impl opentelemetry_sdk::metrics::reader::MetricReader) {
-        let registry = if let Some(prefix) = self.prefix.clone() {
+        let registry = if let Some(registry) = self.registry.clone() {
+            registry
+        } else if let Some(prefix) = self.prefix.clone() {
             Registry::new_custom(Some(prefix), self.labels.clone()).expect("create prometheus registry")
         } else {
             Registry::new()
@@ -757,4 +766,24 @@ mod tests {
             "<h1>Hello, World!</h1>"
         }
     }
+
+    #[test]
+    fn test_builder_with_custom_registry() {
+        let metrics = HttpMetricsLayerBuilder::new()
+            .with_registry(prometheus::default_registry().to_owned())
+            .build();
+
+        let _app = Router::new()
+            // export metrics at `/metrics` endpoint
+            .merge(metrics.routes::<()>())
+            .route("/", get(handler))
+            .route("/hello", get(handler))
+            .route("/world", get(handler))
+            // add the metrics middleware
+            .layer(metrics);
+
+        async fn handler() -> &'static str {
+            "<h1>Hello, World!</h1>"
+        }
+    }
 }