1use rust_webx_core::error::Result;
4use rust_webx_core::http::IHttpContext;
5use rust_webx_core::middleware::IMiddleware;
6use std::ops::ControlFlow;
7use std::sync::atomic::{AtomicU64, Ordering};
8
9static REQUEST_COUNT: AtomicU64 = AtomicU64::new(0);
10
11pub struct TimingMiddleware;
16
17impl TimingMiddleware {
18 pub fn new() -> Self {
19 Self
20 }
21}
22
23impl Default for TimingMiddleware {
24 fn default() -> Self {
25 Self::new()
26 }
27}
28
29#[async_trait::async_trait]
30impl IMiddleware for TimingMiddleware {
31 async fn invoke(&self, _ctx: &mut dyn IHttpContext) -> Result<ControlFlow<()>> {
32 Ok(ControlFlow::Continue(()))
33 }
34
35 async fn after(&self, ctx: &mut dyn IHttpContext) -> Result<()> {
36 let count = REQUEST_COUNT.fetch_add(1, Ordering::Relaxed);
37 ctx.response_mut()
38 .set_header("x-request-count", &count.to_string());
39 Ok(())
40 }
41}