From cecc35ae85490e024b2865e10739e64d31de0e2d Mon Sep 17 00:00:00 2001
From: asonix <asonix@asonix.dog>
Date: Sat, 19 Nov 2022 20:35:45 -0600
Subject: [PATCH] Add timings metrics middleware

---
 src/main.rs               |  3 +-
 src/middleware.rs         |  2 ++
 src/middleware/timings.rs | 72 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 76 insertions(+), 1 deletion(-)
 create mode 100644 src/middleware/timings.rs

diff --git a/src/main.rs b/src/main.rs
index adeff75..0b309e2 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -33,7 +33,7 @@ use self::{
     data::{ActorCache, MediaCache, State},
     db::Db,
     jobs::create_workers,
-    middleware::{DebugPayload, RelayResolver},
+    middleware::{DebugPayload, RelayResolver, Timings},
     routes::{actor, inbox, index, nodeinfo, nodeinfo_meta, statics},
 };
 
@@ -182,6 +182,7 @@ async fn main() -> Result<(), anyhow::Error> {
         };
 
         app.wrap(TracingLogger::default())
+            .wrap(Timings)
             .service(web::resource("/").route(web::get().to(index)))
             .service(web::resource("/media/{path}").route(web::get().to(routes::media)))
             .service(
diff --git a/src/middleware.rs b/src/middleware.rs
index e11344d..93d6a68 100644
--- a/src/middleware.rs
+++ b/src/middleware.rs
@@ -1,7 +1,9 @@
 mod payload;
+mod timings;
 mod verifier;
 mod webfinger;
 
 pub(crate) use payload::DebugPayload;
+pub(crate) use timings::Timings;
 pub(crate) use verifier::MyVerify;
 pub(crate) use webfinger::RelayResolver;
diff --git a/src/middleware/timings.rs b/src/middleware/timings.rs
new file mode 100644
index 0000000..4a0f2d5
--- /dev/null
+++ b/src/middleware/timings.rs
@@ -0,0 +1,72 @@
+use actix_web::dev::{Service, ServiceRequest, Transform};
+use futures_util::future::LocalBoxFuture;
+use std::{
+    future::{ready, Ready},
+    time::Instant,
+};
+
+pub(crate) struct Timings;
+pub(crate) struct TimingsMiddleware<S>(S);
+
+struct LogOnDrop {
+    begin: Instant,
+    path: String,
+    method: String,
+}
+
+impl Drop for LogOnDrop {
+    fn drop(&mut self) {
+        let duration = self.begin.elapsed();
+        metrics::histogram!("relay.request.complete", duration, "path" => self.path.clone(), "method" => self.method.clone());
+    }
+}
+
+impl<S> Transform<S, ServiceRequest> for Timings
+where
+    S: Service<ServiceRequest>,
+    S::Future: 'static,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type InitError = ();
+    type Transform = TimingsMiddleware<S>;
+    type Future = Ready<Result<Self::Transform, Self::InitError>>;
+
+    fn new_transform(&self, service: S) -> Self::Future {
+        ready(Ok(TimingsMiddleware(service)))
+    }
+}
+
+impl<S> Service<ServiceRequest> for TimingsMiddleware<S>
+where
+    S: Service<ServiceRequest>,
+    S::Future: 'static,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = LocalBoxFuture<'static, Result<S::Response, S::Error>>;
+
+    fn poll_ready(
+        &self,
+        ctx: &mut core::task::Context<'_>,
+    ) -> std::task::Poll<Result<(), Self::Error>> {
+        self.0.poll_ready(ctx)
+    }
+
+    fn call(&self, req: ServiceRequest) -> Self::Future {
+        let logger = LogOnDrop {
+            begin: Instant::now(),
+            path: req.path().to_string(),
+            method: req.method().to_string(),
+        };
+        let fut = self.0.call(req);
+
+        Box::pin(async move {
+            let res = fut.await;
+
+            drop(logger);
+
+            res
+        })
+    }
+}