use crate::RuntimeError;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::Poll;
pub(crate) trait Acceptor {
type Accepted;
fn accept(&self) -> impl Future<Output = Result<Self::Accepted, std::io::Error>> + Send + '_;
}
impl Acceptor for tokio::net::TcpListener {
type Accepted = (tokio::net::TcpStream, std::net::SocketAddr);
fn accept(&self) -> impl Future<Output = Result<Self::Accepted, std::io::Error>> + Send + '_ {
tokio::net::TcpListener::accept(self)
}
}
impl Acceptor for tokio::net::UnixListener {
type Accepted = tokio::net::UnixStream;
async fn accept(&self) -> Result<Self::Accepted, std::io::Error> {
let (stream, _addr) = tokio::net::UnixListener::accept(self).await?;
Ok(stream)
}
}
pub(crate) async fn accept_loop<L, F, Fut>(
listener: &L,
shutdown_notify: &tokio::sync::Notify,
conn_limit: Option<&Arc<tokio::sync::Semaphore>>,
on_accept: F,
) -> Result<(), RuntimeError>
where
L: Acceptor,
F: Fn(L::Accepted) -> Fut,
Fut: Future<Output = ()> + Send + 'static,
{
loop {
tokio::select! {
result = listener.accept() => {
match result {
Ok(accepted) => {
spawn_with_limit(conn_limit, on_accept(accepted)).await;
}
Err(e) if crate::error::is_transient_accept_error(&e) => {
tracing::warn!("accept: fd limit reached, backing off");
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
}
Err(e) => return Err(e.into()),
}
}
() = shutdown_notify.notified() => {
return Ok(());
}
}
}
}
pub(crate) async fn accept_loop_with_permit<L, F, Fut>(
listener: &L,
shutdown_notify: &tokio::sync::Notify,
conn_limit: Option<&Arc<tokio::sync::Semaphore>>,
script: Option<&Arc<crate::http::mock::LifecycleScript>>,
on_accept: F,
) -> Result<(), RuntimeError>
where
L: Acceptor,
F: Fn(L::Accepted, Option<tokio::sync::OwnedSemaphorePermit>) -> Fut,
Fut: Future<Output = ()> + Send + 'static,
{
loop {
tokio::select! {
result = listener.accept() => {
match result {
Ok(accepted) => {
let permit = match conn_limit {
Some(_) => acquire_connection_permit(conn_limit, script).await.ok(),
None => None,
};
if conn_limit.is_none() || permit.is_some() {
tokio::spawn(on_accept(accepted, permit));
}
}
Err(error) if crate::error::is_transient_accept_error(&error) => {
tracing::warn!("accept: fd limit reached, backing off");
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
}
Err(error) => return Err(error.into()),
}
}
() = shutdown_notify.notified() => return Ok(()),
}
}
}
pub(crate) async fn acquire_connection_permit(
conn_limit: Option<&Arc<tokio::sync::Semaphore>>,
script: Option<&Arc<crate::http::mock::LifecycleScript>>,
) -> Result<tokio::sync::OwnedSemaphorePermit, tokio::sync::AcquireError> {
let semaphore = match conn_limit {
Some(semaphore) => Arc::clone(semaphore),
None => return std::future::pending().await,
};
let future = semaphore.acquire_owned();
tokio::pin!(future);
let immediate =
std::future::poll_fn(
|context| match Future::poll(Pin::new(&mut future), context) {
Poll::Ready(result) => Poll::Ready(Some(result)),
Poll::Pending => Poll::Ready(None),
},
)
.await;
match (immediate, script) {
(Some(result), _) => result,
(None, Some(script)) => {
script
.pause(crate::http::mock::LifecycleCheckpoint::ConnectionPermitWaitPending)
.await;
future.await
}
(None, None) => future.await,
}
}
async fn spawn_with_limit<Fut>(conn_limit: Option<&Arc<tokio::sync::Semaphore>>, fut: Fut)
where
Fut: Future<Output = ()> + Send + 'static,
{
let permit = match conn_limit {
None => {
tokio::spawn(fut);
return;
}
Some(sem) => Arc::clone(sem).acquire_owned().await,
};
if let Ok(permit) = permit {
tokio::spawn(async move {
fut.await;
drop(permit);
});
}
}
pub(crate) async fn tls_handshake(
stream: tokio::net::TcpStream,
acceptor: &tokio_rustls::TlsAcceptor,
) -> Option<tokio_rustls::server::TlsStream<tokio::net::TcpStream>> {
let result =
tokio::time::timeout(std::time::Duration::from_secs(10), acceptor.accept(stream)).await;
match result {
Ok(Ok(s)) => Some(s),
Ok(Err(e)) if crate::error::is_benign_io(&e) => None,
Ok(Err(e)) => {
tracing::warn!("TLS handshake error: {e}");
None
}
Err(_) => None,
}
}