apply limits

1 files changed, 9 insertions, 19 deletions

diff --git a/src/main.rs b/src/main.rs
index f55771c..f51cd67 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,6 +4,7 @@
 pub mod config;
 pub mod error;
 pub mod files;
+pub mod limiter;
 pub mod proxy;
 
 use crate::{
@@ -22,18 +23,9 @@ use hyper::{
     service::service_fn,
     Request, Response, StatusCode,
 };
+use limiter::Limiter;
 use log::{debug, error, info, warn};
-use std::{
-    fs::File,
-    io::BufReader,
-    net::SocketAddr,
-    path::Path,
-    process::exit,
-    sync::{
-        atomic::{AtomicUsize, Ordering::Relaxed},
-        Arc,
-    },
-};
+use std::{fs::File, io::BufReader, net::SocketAddr, path::Path, process::exit, sync::Arc};
 use tokio::{
     io::{AsyncRead, AsyncWrite},
     net::TcpListener,
@@ -43,7 +35,8 @@ use tokio_rustls::TlsAcceptor;
 
 pub struct State {
     pub config: Config,
-    pub total_connection: AtomicUsize,
+    pub l_incoming: Limiter,
+    pub l_outgoing: Limiter,
 }
 
 #[tokio::main]
@@ -62,8 +55,9 @@ async fn main() -> anyhow::Result<()> {
         }
     };
     let state = Arc::new(State {
+        l_incoming: Limiter::new(config.limits.max_incoming_connections),
+        l_outgoing: Limiter::new(config.limits.max_outgoing_connections),
         config,
-        total_connection: AtomicUsize::new(0),
     });
 
     {
@@ -145,9 +139,7 @@ pub async fn serve_stream<T: AsyncRead + AsyncWrite + Unpin + Send + 'static>(
     stream: T,
     addr: SocketAddr,
 ) {
-    let conns = state.total_connection.fetch_add(1, Relaxed);
-
-    if conns >= state.config.limits.max_incoming_connections {
+    if let Some(_limit_guard) = state.l_incoming.obtain() {
         let conn = http1::Builder::new()
             .serve_connection(
                 stream,
@@ -177,8 +169,6 @@ pub async fn serve_stream<T: AsyncRead + AsyncWrite + Unpin + Send + 'static>(
     } else {
         warn!("connection dropped: too many incoming");
     }
-
-    state.total_connection.fetch_sub(1, Relaxed);
 }
 
 fn load_certs(path: &Path) -> anyhow::Result<Vec<rustls::Certificate>> {
@@ -216,7 +206,7 @@ async fn service(
         .ok_or(ServiceError::NoHost)?;
 
     let mut resp = match route {
-        HostConfig::Backend { backend } => proxy_request(req, addr, backend).await,
+        HostConfig::Backend { backend } => proxy_request(&state, req, addr, backend).await,
         HostConfig::Files { files } => serve_files(req, files).await,
     }?;