Merge pull request #5687 from Turbo87/update-metrics

middleware/update_metrics: Migrate to `axum`

Author: Turbo87

Cargo.toml

@@ -24,7 +24,7 @@ path = "src/tests/all.rs"
 [dependencies]
 anyhow = "=1.0.67"
 aws-sigv4 = "=0.52.0"
-axum = { version = "=0.6.1", features = ["headers", "macros"] }
+axum = { version = "=0.6.1", features = ["headers", "macros", "matched-path"] }
 axum-extra = { version = "=0.4.2", features = ["cookie-signed"] }
 base64 = "=0.13.1"
 cargo-registry-index = { path = "cargo-registry-index" }

conduit-axum/Cargo.toml

@@ -11,6 +11,7 @@ rust-version = "1.56.0"
 [dependencies]
 axum = "=0.6.1"
 conduit = "=0.10.0"
+conduit-router = "=0.10.0"
 hyper = { version = "=0.14.23", features = ["server", "stream"] }
 http = "=0.2.8"
 percent-encoding = "=2.2.0"

conduit-axum/src/fallback.rs

@@ -11,7 +11,8 @@ use axum::body::{Body, HttpBody};
 use axum::extract::{ConnectInfo, Extension};
 use axum::handler::Handler as AxumHandler;
 use axum::response::IntoResponse;
-use conduit::{Handler, StartInstant};
+use conduit::{Handler, RequestExt, StartInstant};
+use conduit_router::RoutePattern;
 use http::header::CONTENT_LENGTH;
 use http::StatusCode;
 use hyper::{Request, Response};
@@ -58,7 +59,7 @@ async fn fallback_to_conduit(
             let mut request = ConduitRequest::new(request, remote_addr, now);
             handler
                 .call(&mut request)
-                .map(conduit_into_axum)
+                .map(|response| conduit_into_axum(response, request))
                 .unwrap_or_else(|e| server_error_response(&*e))
         })
     })
@@ -67,9 +68,13 @@ async fn fallback_to_conduit(
 }
 
 /// Turns a `ConduitResponse` into a `AxumResponse`
-pub fn conduit_into_axum(response: ConduitResponse) -> AxumResponse {
+fn conduit_into_axum(mut response: ConduitResponse, mut request: ConduitRequest) -> AxumResponse {
     use conduit::Body::*;
 
+    if let Some(pattern) = request.mut_extensions().remove::<RoutePattern>() {
+        response.extensions_mut().insert(pattern);
+    }
+
     let (parts, body) = response.into_parts();
     match body {
         Static(slice) => Response::from_parts(parts, axum::body::Body::from(slice)).into_response(),

src/middleware.rs

@@ -10,7 +10,6 @@ use self::ember_html::EmberHtml;
 use self::head::Head;
 use self::known_error_to_json::KnownErrorToJson;
 use self::static_or_continue::StaticOrContinue;
-use self::update_metrics::UpdateMetrics;
 
 pub mod app;
 mod balance_capacity;
@@ -49,6 +48,10 @@ pub fn apply_axum_middleware(state: AppState, router: Router) -> Router {
         .layer(sentry_tower::NewSentryLayer::<Request>::new_from_top())
         .layer(sentry_tower::SentryHttpLayer::with_transaction())
         .layer(from_fn(log_request::log_requests))
+        .layer(from_fn_with_state(
+            state.clone(),
+            update_metrics::update_metrics,
+        ))
         // The following layer is unfortunately necessary for `option_layer()` to work
         .layer(HandleErrorLayer::new(dummy_error_handler))
         // Optionally print debug information for each request
@@ -78,9 +81,6 @@ pub fn build_middleware(app: Arc<App>, endpoints: RouteBuilder) -> MiddlewareBui
     m.add(AppMiddleware::new(app));
     m.add(KnownErrorToJson);
 
-    // This is added *after* AppMiddleware to make sure the app is available.
-    m.add(UpdateMetrics);
-
     // Note: The following `m.around()` middleware is run from bottom to top
 
     // This is currently the final middleware to run. If a middleware layer requires a database

src/middleware/update_metrics.rs

@@ -1,44 +1,45 @@
-use super::app::RequestApp;
-use super::prelude::*;
+use crate::app::AppState;
+use axum::extract::{MatchedPath, State};
+use axum::middleware::Next;
+use axum::response::Response;
 use conduit_router::RoutePattern;
+use http::Request;
+use std::time::Instant;
 
-#[derive(Debug, Default)]
-pub(super) struct UpdateMetrics;
+pub async fn update_metrics<B>(
+    State(state): State<AppState>,
+    matched_path: Option<MatchedPath>,
+    req: Request<B>,
+    next: Next<B>,
+) -> Response {
+    let start_instant = Instant::now();
 
-impl Middleware for UpdateMetrics {
-    fn before(&self, req: &mut dyn RequestExt) -> BeforeResult {
-        let metrics = &req.app().instance_metrics;
+    let metrics = &state.instance_metrics;
+    metrics.requests_in_flight.inc();
 
-        metrics.requests_in_flight.inc();
+    let response = next.run(req).await;
 
-        Ok(())
-    }
+    metrics.requests_in_flight.dec();
+    metrics.requests_total.inc();
 
-    fn after(&self, req: &mut dyn RequestExt, res: AfterResult) -> AfterResult {
-        let metrics = &req.app().instance_metrics;
-
-        metrics.requests_in_flight.dec();
-        metrics.requests_total.inc();
-
-        let endpoint = req
+    let endpoint = match matched_path {
+        Some(ref matched_path) => matched_path.as_str(),
+        None => response
             .extensions()
             .get::<RoutePattern>()
-            .map(|p| p.pattern())
-            .unwrap_or("<unknown>");
-        metrics
-            .response_times
-            .with_label_values(&[endpoint])
-            .observe(req.elapsed().as_millis() as f64 / 1000.0);
-
-        let status = match &res {
-            Ok(res) => res.status().as_u16(),
-            Err(_) => 500,
-        };
-        metrics
-            .responses_by_status_code_total
-            .with_label_values(&[&status.to_string()])
-            .inc();
-
-        res
-    }
+            .map(|route_pattern| route_pattern.pattern())
+            .unwrap_or("<unknown>"),
+    };
+    metrics
+        .response_times
+        .with_label_values(&[endpoint])
+        .observe(start_instant.elapsed().as_millis() as f64 / 1000.0);
+
+    let status = response.status().as_u16();
+    metrics
+        .responses_by_status_code_total
+        .with_label_values(&[&status.to_string()])
+        .inc();
+
+    response
 }