Skip to content

feat(proxy): bound accepted client connections #191

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gakonst wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(proxy): bound accepted client connections #191

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ rand = { version = "0.10.2", optional = true }
rcgen = { version = "0.14.0", features = ["x509-parser"], optional = true }
thiserror = "2.0.7"
time = { version = "0.3.35", optional = true }
tokio = { version = "1.24.2", features = ["macros", "rt"] }
tokio = { version = "1.24.2", features = ["macros", "rt", "sync"] }
tokio-graceful = "0.2.0"
tokio-native-tls = { version = "0.3.1", optional = true }
tokio-rustls = { version = "0.26.3", features = ["logging", "tls12"] }
Expand Down
23 changes: 23 additions & 0 deletions src/proxy/builder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ use hyper_util::{
use std::{
future::{Pending, pending},
net::SocketAddr,
num::NonZeroUsize,
sync::Arc,
};
use thiserror::Error;
Expand Down Expand Up @@ -141,6 +142,7 @@ impl<CA> ProxyBuilder<WantsClient<CA>> {
return ProxyBuilder(WantsHandlers {
al: self.0.al,
ca: self.0.ca,
max_concurrent_connections: None,
http_connector: Err(Error::from(e)),
client: None,
http_handler: NoopHandler::new(),
Expand All @@ -165,6 +167,7 @@ impl<CA> ProxyBuilder<WantsClient<CA>> {
ProxyBuilder(WantsHandlers {
al: self.0.al,
ca: self.0.ca,
max_concurrent_connections: None,
http_connector: Ok(https),
client: None,
http_handler: NoopHandler::new(),
Expand All @@ -189,6 +192,7 @@ impl<CA> ProxyBuilder<WantsClient<CA>> {
return ProxyBuilder(WantsHandlers {
al: self.0.al,
ca: self.0.ca,
max_concurrent_connections: None,
http_connector: Err(Error::from(e)),
client: None,
http_handler: NoopHandler::new(),
Expand All @@ -206,6 +210,7 @@ impl<CA> ProxyBuilder<WantsClient<CA>> {
ProxyBuilder(WantsHandlers {
al: self.0.al,
ca: self.0.ca,
max_concurrent_connections: None,
http_connector: Ok(https),
client: None,
http_handler: NoopHandler::new(),
Expand All @@ -227,6 +232,7 @@ impl<CA> ProxyBuilder<WantsClient<CA>> {
ProxyBuilder(WantsHandlers {
al: self.0.al,
ca: self.0.ca,
max_concurrent_connections: None,
http_connector: Ok(connector),
client: None,
http_handler: NoopHandler::new(),
Expand All @@ -242,6 +248,7 @@ impl<CA> ProxyBuilder<WantsClient<CA>> {
pub struct WantsHandlers<CA, C, H, W, F> {
al: AddrOrListener,
ca: CA,
max_concurrent_connections: Option<NonZeroUsize>,
http_connector: Result<C, Error>,
client: Option<ClientBuilder>,
http_handler: H,
Expand All @@ -260,6 +267,7 @@ impl<CA, C, H, W, F> ProxyBuilder<WantsHandlers<CA, C, H, W, F>> {
ProxyBuilder(WantsHandlers {
al: self.0.al,
ca: self.0.ca,
max_concurrent_connections: self.0.max_concurrent_connections,
http_connector: self.0.http_connector,
client: self.0.client,
http_handler,
Expand All @@ -278,6 +286,7 @@ impl<CA, C, H, W, F> ProxyBuilder<WantsHandlers<CA, C, H, W, F>> {
ProxyBuilder(WantsHandlers {
al: self.0.al,
ca: self.0.ca,
max_concurrent_connections: self.0.max_concurrent_connections,
http_connector: self.0.http_connector,
client: self.0.client,
http_handler: self.0.http_handler,
Expand Down Expand Up @@ -312,6 +321,18 @@ impl<CA, C, H, W, F> ProxyBuilder<WantsHandlers<CA, C, H, W, F>> {
})
}

/// Set the maximum number of concurrently accepted client connections.
///
/// The proxy waits for a connection slot before accepting another socket.
/// This bounds the file descriptors held by the process and applies
/// backpressure through the listener backlog.
pub fn with_max_concurrent_connections(self, max_concurrent_connections: NonZeroUsize) -> Self {
ProxyBuilder(WantsHandlers {
max_concurrent_connections: Some(max_concurrent_connections),
..self.0
})
}

/// Set a future that when ready will gracefully shutdown the proxy server.
pub fn with_graceful_shutdown<F2: Future<Output = ()> + Send + 'static>(
self,
Expand All @@ -320,6 +341,7 @@ impl<CA, C, H, W, F> ProxyBuilder<WantsHandlers<CA, C, H, W, F>> {
ProxyBuilder(WantsHandlers {
al: self.0.al,
ca: self.0.ca,
max_concurrent_connections: self.0.max_concurrent_connections,
http_connector: self.0.http_connector,
client: self.0.client,
http_handler: self.0.http_handler,
Expand All @@ -338,6 +360,7 @@ impl<CA, C, H, W, F> ProxyBuilder<WantsHandlers<CA, C, H, W, F>> {
Ok(Proxy {
al: self.0.al,
ca: Arc::new(self.0.ca),
max_concurrent_connections: self.0.max_concurrent_connections,
http_connector: self.0.http_connector?,
client: self.0.client,
http_handler: self.0.http_handler,
Expand Down
21 changes: 19 additions & 2 deletions src/proxy/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,8 @@ use hyper_util::{
server::conn::auto::Builder as ServerBuilder,
};
use internal::InternalProxy;
use std::{error::Error as StdError, sync::Arc};
use tokio::net::TcpListener;
use std::{error::Error as StdError, num::NonZeroUsize, sync::Arc};
use tokio::{net::TcpListener, sync::Semaphore};
use tokio_graceful::Shutdown;
use tokio_tungstenite::Connector;
use tracing::error;
Expand Down Expand Up @@ -73,6 +73,7 @@ use tracing::error;
pub struct Proxy<C, CA, H, W, F> {
al: AddrOrListener,
ca: Arc<CA>,
max_concurrent_connections: Option<NonZeroUsize>,
http_connector: C,
client: Option<ClientBuilder>,
http_handler: H,
Expand Down Expand Up @@ -130,8 +131,23 @@ where

let shutdown = Shutdown::new(self.graceful_shutdown);
let guard = shutdown.guard_weak();
let connection_permits = self
.max_concurrent_connections
.map(|maximum| Arc::new(Semaphore::new(maximum.get())));

loop {
let connection_permit = if let Some(permits) = &connection_permits {
let permit = tokio::select! {
permit = Arc::clone(permits).acquire_owned() => permit,
_ = guard.cancelled() => break,
};
let Ok(permit) = permit else {
break;
};
Some(permit)
} else {
None
};
tokio::select! {
res = listener.accept() => {
let (tcp, client_addr) = match res {
Expand All @@ -150,6 +166,7 @@ where
let websocket_connector = self.websocket_connector.clone();

shutdown.spawn_task_fn(move |guard| async move {
let _connection_permit = connection_permit;
let conn = server.serve_connection_with_upgrades(
TokioIo::new(tcp),
service_fn(|req| {
Expand Down