use actix_web::{ body::MessageBody, dev::{Service, ServiceRequest, ServiceResponse, Transform}, http::StatusCode, HttpResponse, ResponseError, }; use std::{ cell::RefCell, future::{ready, Future, Ready}, pin::Pin, task::{Context, Poll}, time::Instant, }; struct MetricsGuard { start: Instant, matched_path: Option, armed: bool, } struct MetricsGuardWithStatus { start: Instant, matched_path: Option, status: StatusCode, } impl MetricsGuard { fn new(matched_path: Option) -> Self { metrics::counter!("pict-rs.request.start", "path" => format!("{matched_path:?}")) .increment(1); Self { start: Instant::now(), matched_path, armed: true, } } fn with_status(mut self, status: StatusCode) -> MetricsGuardWithStatus { self.armed = false; MetricsGuardWithStatus { start: self.start, matched_path: self.matched_path.clone(), status, } } } impl Drop for MetricsGuard { fn drop(&mut self) { if self.armed { metrics::counter!("pict-rs.request.complete", "path" => format!("{:?}", self.matched_path)).increment(1); metrics::histogram!("pict-rs.request.timings", "path" => format!("{:?}", self.matched_path)).record(self.start.elapsed().as_secs_f64()); } } } impl Drop for MetricsGuardWithStatus { fn drop(&mut self) { metrics::counter!("pict-rs.request.complete", "path" => format!("{:?}", self.matched_path), "status" => self.status.to_string()).increment(1); metrics::histogram!("pict-rs.request.timings", "path" => format!("{:?}", self.matched_path), "status" => self.status.to_string()).record(self.start.elapsed().as_secs_f64()); } } pub(crate) struct Metrics; pub(crate) struct MetricsMiddleware { inner: S, } pub(crate) struct MetricsError { guard: RefCell>, inner: actix_web::Error, } pin_project_lite::pin_project! { pub(crate) struct MetricsFuture { guard: Option, #[pin] inner: F, } } pin_project_lite::pin_project! { pub(crate) struct MetricsBody { guard: Option, #[pin] inner: B, } } impl Transform for Metrics where S: Service>, S::Future: 'static, S::Error: Into, { type Response = ServiceResponse>; type Error = actix_web::Error; type InitError = (); type Transform = MetricsMiddleware; type Future = Ready>; fn new_transform(&self, service: S) -> Self::Future { ready(Ok(MetricsMiddleware { inner: service })) } } impl Service for MetricsMiddleware where S: Service>, S::Future: 'static, S::Error: Into, { type Response = ServiceResponse>; type Error = actix_web::Error; type Future = MetricsFuture; fn poll_ready(&self, ctx: &mut core::task::Context<'_>) -> Poll> { let res = std::task::ready!(self.inner.poll_ready(ctx)); Poll::Ready(res.map_err(|e| { MetricsError { guard: RefCell::new(None), inner: e.into(), } .into() })) } fn call(&self, req: ServiceRequest) -> Self::Future { let matched_path = req.match_pattern(); MetricsFuture { guard: Some(MetricsGuard::new(matched_path)), inner: self.inner.call(req), } } } impl Future for MetricsFuture where F: Future, E>>, E: Into, { type Output = Result>, actix_web::Error>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.project(); match std::task::ready!(this.inner.poll(cx)) { Ok(response) => { let guard = this.guard.take(); Poll::Ready(Ok(response.map_body(|head, inner| MetricsBody { guard: guard.map(|guard| guard.with_status(head.status)), inner, }))) } Err(e) => { let guard = this.guard.take(); Poll::Ready(Err(MetricsError { guard: RefCell::new(guard), inner: e.into(), } .into())) } } } } impl MessageBody for MetricsBody where B: MessageBody, { type Error = B::Error; fn size(&self) -> actix_web::body::BodySize { self.inner.size() } fn poll_next( self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll>> { let this = self.project(); let opt = std::task::ready!(this.inner.poll_next(cx)); if opt.is_none() { this.guard.take(); } Poll::Ready(opt) } } impl std::fmt::Debug for MetricsError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("MetricsError") .field("guard", &"Guard") .field("inner", &self.inner) .finish() } } impl std::fmt::Display for MetricsError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { self.inner.fmt(f) } } impl std::error::Error for MetricsError { fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { self.inner.source() } } impl ResponseError for MetricsError { fn status_code(&self) -> StatusCode { self.inner.as_response_error().status_code() } fn error_response(&self) -> HttpResponse { let guard = self.guard.borrow_mut().take(); self.inner.error_response().map_body(|head, inner| { MetricsBody { guard: guard.map(|guard| guard.with_status(head.status)), inner, } .boxed() }) } }