forked from mirrors/relay
Improve Timings middleware
This commit is contained in:
parent
e987149757
commit
ed0ea6521e
3 changed files with 82 additions and 25 deletions
1
Cargo.lock
generated
1
Cargo.lock
generated
|
@ -308,6 +308,7 @@ dependencies = [
|
|||
"mime",
|
||||
"opentelemetry",
|
||||
"opentelemetry-otlp",
|
||||
"pin-project-lite",
|
||||
"quanta",
|
||||
"rand",
|
||||
"rsa",
|
||||
|
|
|
@ -46,6 +46,7 @@ metrics-util = "0.14.0"
|
|||
mime = "0.3.16"
|
||||
opentelemetry = { version = "0.18", features = ["rt-tokio"] }
|
||||
opentelemetry-otlp = "0.11"
|
||||
pin-project-lite = "0.2.9"
|
||||
quanta = "0.10.1"
|
||||
rand = "0.8"
|
||||
rsa = "0.7"
|
||||
|
|
|
@ -1,10 +1,10 @@
|
|||
use actix_web::{
|
||||
body::MessageBody,
|
||||
dev::{Service, ServiceRequest, ServiceResponse, Transform},
|
||||
http::StatusCode,
|
||||
};
|
||||
use futures_util::future::LocalBoxFuture;
|
||||
use std::{
|
||||
future::{ready, Ready},
|
||||
future::{ready, Future, Ready},
|
||||
time::Instant,
|
||||
};
|
||||
|
||||
|
@ -15,12 +15,30 @@ struct LogOnDrop {
|
|||
begin: Instant,
|
||||
path: String,
|
||||
method: String,
|
||||
disarm: bool,
|
||||
arm: bool,
|
||||
}
|
||||
|
||||
pin_project_lite::pin_project! {
|
||||
pub(crate) struct TimingsFuture<F> {
|
||||
#[pin]
|
||||
future: F,
|
||||
|
||||
log_on_drop: Option<LogOnDrop>,
|
||||
}
|
||||
}
|
||||
|
||||
pin_project_lite::pin_project! {
|
||||
pub(crate) struct TimingsBody<B> {
|
||||
#[pin]
|
||||
body: B,
|
||||
|
||||
log_on_drop: LogOnDrop,
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for LogOnDrop {
|
||||
fn drop(&mut self) {
|
||||
if !self.disarm {
|
||||
if self.arm {
|
||||
let duration = self.begin.elapsed();
|
||||
metrics::histogram!("relay.request.complete", duration, "path" => self.path.clone(), "method" => self.method.clone());
|
||||
}
|
||||
|
@ -32,7 +50,7 @@ where
|
|||
S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error>,
|
||||
S::Future: 'static,
|
||||
{
|
||||
type Response = S::Response;
|
||||
type Response = ServiceResponse<TimingsBody<B>>;
|
||||
type Error = S::Error;
|
||||
type InitError = ();
|
||||
type Transform = TimingsMiddleware<S>;
|
||||
|
@ -48,9 +66,9 @@ where
|
|||
S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error>,
|
||||
S::Future: 'static,
|
||||
{
|
||||
type Response = S::Response;
|
||||
type Response = ServiceResponse<TimingsBody<B>>;
|
||||
type Error = S::Error;
|
||||
type Future = LocalBoxFuture<'static, Result<S::Response, S::Error>>;
|
||||
type Future = TimingsFuture<S::Future>;
|
||||
|
||||
fn poll_ready(
|
||||
&self,
|
||||
|
@ -60,29 +78,66 @@ where
|
|||
}
|
||||
|
||||
fn call(&self, req: ServiceRequest) -> Self::Future {
|
||||
let mut logger = LogOnDrop {
|
||||
let log_on_drop = LogOnDrop {
|
||||
begin: Instant::now(),
|
||||
path: req.path().to_string(),
|
||||
method: req.method().to_string(),
|
||||
disarm: false,
|
||||
arm: false,
|
||||
};
|
||||
let fut = self.0.call(req);
|
||||
|
||||
Box::pin(async move {
|
||||
let res = fut.await;
|
||||
let future = self.0.call(req);
|
||||
|
||||
let status = match &res {
|
||||
Ok(res) => res.status(),
|
||||
Err(e) => e.as_response_error().status_code(),
|
||||
};
|
||||
if status == StatusCode::NOT_FOUND || status == StatusCode::METHOD_NOT_ALLOWED {
|
||||
logger.disarm = true;
|
||||
}
|
||||
|
||||
// TODO: Drop after body write
|
||||
drop(logger);
|
||||
|
||||
res
|
||||
})
|
||||
TimingsFuture {
|
||||
future,
|
||||
log_on_drop: Some(log_on_drop),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<F, B> Future for TimingsFuture<F>
|
||||
where
|
||||
F: Future<Output = Result<ServiceResponse<B>, actix_web::Error>>,
|
||||
{
|
||||
type Output = Result<ServiceResponse<TimingsBody<B>>, actix_web::Error>;
|
||||
|
||||
fn poll(
|
||||
self: std::pin::Pin<&mut Self>,
|
||||
cx: &mut std::task::Context<'_>,
|
||||
) -> std::task::Poll<Self::Output> {
|
||||
let this = self.project();
|
||||
|
||||
let res = std::task::ready!(this.future.poll(cx));
|
||||
|
||||
let mut log_on_drop = this
|
||||
.log_on_drop
|
||||
.take()
|
||||
.expect("TimingsFuture polled after completion");
|
||||
|
||||
let status = match &res {
|
||||
Ok(res) => res.status(),
|
||||
Err(e) => e.as_response_error().status_code(),
|
||||
};
|
||||
|
||||
log_on_drop.arm =
|
||||
status != StatusCode::NOT_FOUND && status != StatusCode::METHOD_NOT_ALLOWED;
|
||||
|
||||
let res = res.map(|r| r.map_body(|_, body| TimingsBody { body, log_on_drop }));
|
||||
|
||||
std::task::Poll::Ready(res)
|
||||
}
|
||||
}
|
||||
|
||||
impl<B: MessageBody> MessageBody for TimingsBody<B> {
|
||||
type Error = B::Error;
|
||||
|
||||
fn size(&self) -> actix_web::body::BodySize {
|
||||
self.body.size()
|
||||
}
|
||||
|
||||
fn poll_next(
|
||||
self: std::pin::Pin<&mut Self>,
|
||||
cx: &mut std::task::Context<'_>,
|
||||
) -> std::task::Poll<Option<Result<actix_web::web::Bytes, Self::Error>>> {
|
||||
self.project().body.poll_next(cx)
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Reference in a new issue