salvo_extra 0.91.0

Salvo is a powerful web framework that can make your work easier.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

//! Middleware for limiting concurrency.
//! 
//! This middleware limits the maximum number of requests being processed concurrently,
//! which helps prevent server overload during traffic spikes.
//!
//! # Example
//! 
//! ```no_run
//! use std::fs::create_dir_all;
//! use std::path::Path;
//! 
//! use salvo_core::prelude::*;
//! use salvo_extra::concurrency_limiter::*;
//! 
//! #[handler]
//! async fn index(res: &mut Response) {
//!     res.render(Text::Html(INDEX_HTML));
//! }
//! #[handler]
//! async fn upload(req: &mut Request, res: &mut Response) {
//!     let file = req.file("file").await;
//!     tokio::time::sleep(tokio::time::Duration::from_secs(10)).await;
//!     if let Some(file) = file {
//!         let dest = format!("temp/{}", file.name().unwrap_or("file"));
//!         tracing::debug!(dest = %dest, "upload file");
//!         if let Err(e) = std::fs::copy(file.path(), Path::new(&dest)) {
//!             res.status_code(StatusCode::INTERNAL_SERVER_ERROR);
//!             res.render(Text::Plain(format!("file not found in request: {e}")));
//!         } else {
//!             res.render(Text::Plain(format!("File uploaded to {dest}")));
//!         }
//!     } else {
//!         res.status_code(StatusCode::BAD_REQUEST);
//!         res.render(Text::Plain("file not found in request"));
//!     }
//! }
//! 
//! #[tokio::main]
//! async fn main() {
//!     create_dir_all("temp").unwrap();
//!     let router = Router::new()
//!         .get(index)
//!         .push(Router::new().hoop(max_concurrency(1)).path("limited").post(upload))
//!         .push(Router::with_path("unlimit").post(upload));
//! 
//!     let acceptor = TcpListener::new("0.0.0.0:8698").bind().await;
//!     Server::new(acceptor).serve(router).await;
//! }
//! 
//! static INDEX_HTML: &str = r#"<!DOCTYPE html>
//! <html>
//!     <head>
//!         <title>Upload file</title>
//!     </head>
//!     <body>
//!         <h1>Upload file</h1>
//!         <form action="/unlimit" method="post" enctype="multipart/form-data">
//!             <h3>Unlimit</h3>
//!             <input type="file" name="file" />
//!             <input type="submit" value="upload" />
//!         </form>
//!         <form action="/limited" method="post" enctype="multipart/form-data">
//!             <h3>Limited</h3>
//!             <input type="file" name="file" />
//!             <input type="submit" value="upload" />
//!         </form>
//!     </body>
//! </html>
//! "#;
//! ```

use tokio::sync::{Semaphore, TryAcquireError};

use salvo_core::http::StatusError;
use salvo_core::http::{Request, Response};
use salvo_core::{async_trait, Depot, FlowCtrl, Handler};

/// MaxConcurrency
#[derive(Debug)]
pub struct MaxConcurrency {
    semaphore: Semaphore,
}
#[async_trait]
impl Handler for MaxConcurrency {
    #[inline]
    async fn handle(
        &self,
        req: &mut Request,
        depot: &mut Depot,
        res: &mut Response,
        ctrl: &mut FlowCtrl,
    ) {
        match self.semaphore.try_acquire() {
            Ok(_) => {
                ctrl.call_next(req, depot, res).await;
            }
            Err(e) => match e {
                TryAcquireError::Closed => {
                    tracing::error!(
                        "Max concurrency semaphore is never closed, acquire should never fail: {}",
                        e
                    );
                    res.render(StatusError::payload_too_large().brief("Max concurrency reached."));
                }
                TryAcquireError::NoPermits => {
                    tracing::error!("NoPermits : {}", e);
                    res.render(StatusError::too_many_requests().brief("Max concurrency reached."));
                }
            },
        }
    }
}
/// Create a new `MaxConcurrency`.
#[inline]
#[must_use] pub fn max_concurrency(size: usize) -> MaxConcurrency {
    MaxConcurrency {
        semaphore: Semaphore::new(size),
    }
}