use crate::{host::internal_request, *};
use http::Method;
use pin_project_lite::pin_project;
use std::{collections::HashMap, task::ready, time::Instant};
use tracing::Span;
#[derive(Debug, Table, Serialize, Deserialize)]
pub(crate) struct RouteStat {
pub path: String,
pub method_hits_and_latency: HashMap<String, (u64, f64)>,
pub is_asset: bool,
}
impl RouteStat {
pub async fn record(req_method: Method, path: String, latency: f64) {
let req_method = req_method.to_string();
if let Ok(Some(mut stats)) = RouteStat::select_by_path(&path).await {
let entry = stats.method_hits_and_latency.entry(req_method).or_default();
let updated_hits = entry.0 + 1;
let updated_avg_latency = (entry.0 as f64 * entry.1 + latency) / (updated_hits as f64);
*entry = (updated_hits, updated_avg_latency);
if let Err(e) = stats.save().await {
warn!(target:"analytics", "Failed to update stats: {e}");
}
} else {
let is_asset =
mime_guess::from_path(&path).first().is_some() || path.ends_with(".webmanifest");
let mut mhal = HashMap::new();
mhal.insert(req_method, (1, latency));
let stats = RouteStat {
path,
method_hits_and_latency: mhal,
is_asset,
};
if let Err(e) = stats.save().await {
warn!(target:"analytics", "Failed to save new stats: {e}");
}
}
}
}
fn record_response_metrics(
resp: &Response,
latency: std::time::Duration,
_span: &Span,
req_method: Method,
req_path: String,
) {
let latency = latency.as_secs_f64() * 1000.0; let short_latency = format!("{latency:.3}");
let status = resp.status();
let boring_resp = matches!(
status,
StatusCode::NOT_MODIFIED | StatusCode::METHOD_NOT_ALLOWED | StatusCode::NOT_FOUND
);
match boring_resp {
true => trace!(target: "response", latency = %short_latency, code = %status.as_u16()),
false => {
debug!(target: "response", latency = %short_latency, code = %status.as_u16());
RT.spawn(RouteStat::record(req_method, req_path, latency));
}
}
}
#[derive(Clone)]
pub(crate) struct AnalyticsLayer;
impl AnalyticsLayer {
pub fn init() -> Self {
Self
}
}
impl<S> tower::Layer<S> for AnalyticsLayer {
type Service = AnalyticsMiddleware<S>;
fn layer(&self, inner: S) -> Self::Service {
AnalyticsMiddleware { inner }
}
}
#[doc(hidden)]
#[derive(Clone)]
pub(crate) struct AnalyticsMiddleware<S> {
inner: S,
}
use core::{
future::Future,
pin::Pin,
task::{Context, Poll},
};
use std::boxed::Box;
impl<S> tower::Service<Request<Body>> for AnalyticsMiddleware<S>
where
S: tower::Service<Request<Body>, Response = Response> + Send + 'static,
S::Future: Send + 'static,
S::Error: std::fmt::Display + 'static,
{
type Response = S::Response;
type Error = S::Error;
type Future = ResponseFuture<S::Future>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, req: Request<Body>) -> Self::Future {
let start = Instant::now();
let internal_req = internal_request(&req);
let span = make_span(&req);
let req_method = req.method().clone();
let req_path = req.uri().path().to_owned();
let future = {
let _guard = span.enter();
self.inner.call(req)
};
ResponseFuture {
inner: future,
span,
start,
req_method: Some(req_method),
req_path: Some(req_path),
internal_req,
}
}
}
fn make_span(request: &Request) -> Span {
let method = request.method().as_str();
let uri = request.uri();
let path = uri.path();
let uri = if path.starts_with("/auth/") && uri.query().is_some() {
path.to_owned() + "?[redacted]"
} else {
uri.to_string()
};
tracing::debug_span!("http", method, uri)
}
pin_project! {
pub(crate) struct ResponseFuture<F> {
#[pin]
pub(crate) inner: F,
pub(crate) span: Span,
pub(crate) start: Instant,
pub(crate) req_method: Option<Method>,
pub(crate) req_path: Option<String>,
pub(crate) internal_req: bool,
}
}
impl<Fut, E> Future for ResponseFuture<Fut>
where
Fut: Future<Output = Result<Response, E>>,
E: std::fmt::Display + 'static,
{
type Output = Result<Response, E>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();
let _guard = this.span.enter();
let result = ready!(this.inner.poll(cx));
match result {
Ok(res) => {
if !*this.internal_req {
record_response_metrics(
&res,
this.start.elapsed(),
this.span,
this.req_method.take().expect("Request must have a method"),
this.req_path.take().expect("Request must have a path"),
);
}
Poll::Ready(Ok(res))
}
Err(err) => Poll::Ready(Err(err)),
}
}
}