From 9d7101649259bb605d08a432d2cbbb7232e3af39 Mon Sep 17 00:00:00 2001
From: Matthieu Bessat <mbess@mbess.net>
Date: Thu, 29 Aug 2024 11:46:43 +0200
Subject: [PATCH] feat: add helper middleware to easily define metrics config
 on a route

Adding MetricsConfiguratorMiddleware to allow upstream code to
provide a quick way to configure cardinality_keep_params with wrap
even on route where wrap_fn is not as easy.
example: wrap in route macro attribute.
---
 README.md  | 17 +++++++++++++++++
 src/lib.rs | 56 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+)

diff --git a/README.md b/README.md
index 2f509ca..64923e1 100644
--- a/README.md
+++ b/README.md
@@ -260,6 +260,23 @@ web::resource("/posts/{language}/{slug}")
 
 See the full example `with_cardinality_on_params.rs`.
 
+Actix-web-prom also provide a helper middleware which is a little bit more compact :
+
+```rust
+use actix_web::{route, HttpRequest, HttpResponse, Result as ActixResult};
+use actix_web_prom::{MetricsConfiguratorMiddleware, MetricsConfig};
+
+#[route(
+    "/posts/{language}/{slug}",
+    method = "GET",
+    method = "HEAD",
+    wrap = "(MetricsConfiguratorMiddleware(MetricsConfig { cardinality_keep_params: vec![\"language\".into()] })).into_middleware()"
+)]
+async fn get_source_info(req: HttpRequest) -> ActixResult<HttpResponse> {
+    Ok(HttpResponse::Ok().into())
+}
+```
+
 ### Configurable metric names
 
 If you want to rename the default metrics, you can use `ActixMetricsConfiguration` to do so.
diff --git a/src/lib.rs b/src/lib.rs
index 68e2fcd..8f52ae2 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -316,6 +316,7 @@ use std::collections::{HashMap, HashSet};
 use std::future::{ready, Future, Ready};
 use std::marker::PhantomData;
 use std::pin::Pin;
+use std::rc::Rc;
 use std::sync::Arc;
 use std::task::{Context, Poll};
 use std::time::Instant;
@@ -327,6 +328,7 @@ use actix_web::{
         header::{HeaderValue, CONTENT_TYPE},
         Method, StatusCode, Version,
     },
+    middleware::{from_fn, Next},
     web::Bytes,
     Error, HttpMessage,
 };
@@ -924,6 +926,60 @@ impl<B: MessageBody> MessageBody for StreamLog<B> {
     }
 }
 
+/// Helper middleware to allow user that cannot use wrap_fn directly (eg. downstream code is defining a route with attribute macro)
+///
+/// Example use:
+/// ```rust
+/// use actix_web::{route, HttpRequest, HttpResponse, Result as ActixResult};
+/// use actix_web_prom::{MetricsConfiguratorMiddleware, MetricsConfig};
+///
+/// #[route(
+///     "/posts/{language}/{slug}",
+///     method = "GET",
+///     method = "HEAD",
+///     wrap = "(MetricsConfiguratorMiddleware(MetricsConfig { cardinality_keep_params: vec![\"language\".into()] })).into_middleware()"
+/// )]
+/// async fn get_source_info(req: HttpRequest) -> ActixResult<HttpResponse> {
+///     Ok(HttpResponse::Ok().into())
+/// }
+/// ```
+pub struct MetricsConfiguratorMiddleware(pub MetricsConfig);
+
+impl MetricsConfiguratorMiddleware {
+    async fn append_metrics_config(
+        &self,
+        req: ServiceRequest,
+        next: Next<impl MessageBody + 'static>,
+    ) -> Result<ServiceResponse<impl MessageBody>, Error> {
+        let req = req;
+        req.extensions_mut().insert::<MetricsConfig>(MetricsConfig {
+            cardinality_keep_params: self.0.cardinality_keep_params.clone(),
+        });
+        next.call(req).await
+    }
+
+    /// Build into helper middleware to configure metrics behaviour for a specific path or service
+    pub fn into_middleware<S, B>(
+        self,
+    ) -> impl Transform<
+        S,
+        ServiceRequest,
+        Response = ServiceResponse<impl MessageBody>,
+        Error = Error,
+        InitError = (),
+    >
+    where
+        S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
+        B: MessageBody + 'static,
+    {
+        let this = Rc::new(self);
+        from_fn(move |req, next| {
+            let this = Rc::clone(&this);
+            async move { Self::append_metrics_config(&this, req, next).await }
+        })
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;