forked from asonix/pict-rs
239 lines
6.2 KiB
Rust
239 lines
6.2 KiB
Rust
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<String>,
|
|
armed: bool,
|
|
}
|
|
|
|
struct MetricsGuardWithStatus {
|
|
start: Instant,
|
|
matched_path: Option<String>,
|
|
status: StatusCode,
|
|
}
|
|
|
|
impl MetricsGuard {
|
|
fn new(matched_path: Option<String>) -> 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<S> {
|
|
inner: S,
|
|
}
|
|
|
|
pub(crate) struct MetricsError {
|
|
guard: RefCell<Option<MetricsGuard>>,
|
|
inner: actix_web::Error,
|
|
}
|
|
|
|
pin_project_lite::pin_project! {
|
|
pub(crate) struct MetricsFuture<F> {
|
|
guard: Option<MetricsGuard>,
|
|
|
|
#[pin]
|
|
inner: F,
|
|
}
|
|
}
|
|
|
|
pin_project_lite::pin_project! {
|
|
pub(crate) struct MetricsBody<B> {
|
|
guard: Option<MetricsGuardWithStatus>,
|
|
|
|
#[pin]
|
|
inner: B,
|
|
}
|
|
}
|
|
|
|
impl<S, B> Transform<S, ServiceRequest> for Metrics
|
|
where
|
|
S: Service<ServiceRequest, Response = ServiceResponse<B>>,
|
|
S::Future: 'static,
|
|
S::Error: Into<actix_web::Error>,
|
|
{
|
|
type Response = ServiceResponse<MetricsBody<B>>;
|
|
type Error = actix_web::Error;
|
|
type InitError = ();
|
|
type Transform = MetricsMiddleware<S>;
|
|
type Future = Ready<Result<Self::Transform, Self::InitError>>;
|
|
|
|
fn new_transform(&self, service: S) -> Self::Future {
|
|
ready(Ok(MetricsMiddleware { inner: service }))
|
|
}
|
|
}
|
|
|
|
impl<S, B> Service<ServiceRequest> for MetricsMiddleware<S>
|
|
where
|
|
S: Service<ServiceRequest, Response = ServiceResponse<B>>,
|
|
S::Future: 'static,
|
|
S::Error: Into<actix_web::Error>,
|
|
{
|
|
type Response = ServiceResponse<MetricsBody<B>>;
|
|
type Error = actix_web::Error;
|
|
type Future = MetricsFuture<S::Future>;
|
|
|
|
fn poll_ready(&self, ctx: &mut core::task::Context<'_>) -> Poll<Result<(), Self::Error>> {
|
|
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<F, B, E> Future for MetricsFuture<F>
|
|
where
|
|
F: Future<Output = Result<ServiceResponse<B>, E>>,
|
|
E: Into<actix_web::Error>,
|
|
{
|
|
type Output = Result<ServiceResponse<MetricsBody<B>>, actix_web::Error>;
|
|
|
|
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
|
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<B> MessageBody for MetricsBody<B>
|
|
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<Option<Result<actix_web::web::Bytes, Self::Error>>> {
|
|
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<actix_web::body::BoxBody> {
|
|
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()
|
|
})
|
|
}
|
|
}
|