use std::sync::Arc;
use std::time::Duration;
use tokio::io::BufReader;
use tokio::net::TcpListener;
use tokio::sync::{watch, Semaphore};
use tokio::task::JoinHandle;
use crate::channel::AgiChannel;
use crate::error::{AgiError, Result};
use crate::handler::AgiHandler;
use crate::request::AgiRequest;
#[derive(Clone)]
pub struct ShutdownHandle {
tx: watch::Sender<bool>,
}
impl ShutdownHandle {
pub fn shutdown(&self) {
let _ = self.tx.send(true);
}
}
pub struct AgiServer<H: AgiHandler> {
listener: TcpListener,
handler: Arc<H>,
max_connections: Option<usize>,
shutdown_rx: watch::Receiver<bool>,
}
#[must_use]
pub struct AgiServerBuilder<H> {
bind_addr: String,
handler: Option<H>,
max_connections: Option<usize>,
}
impl<H: AgiHandler> AgiServer<H> {
pub fn builder() -> AgiServerBuilder<H> {
AgiServerBuilder {
bind_addr: "127.0.0.1:4573".to_owned(),
handler: None,
max_connections: None,
}
}
pub async fn run(mut self) -> Result<()> {
let semaphore = self.max_connections.map(|n| Arc::new(Semaphore::new(n)));
let mut handles: Vec<JoinHandle<()>> = Vec::new();
loop {
tokio::select! {
result = self.listener.accept() => {
let (stream, peer) = match result {
Ok(conn) => conn,
Err(err) => {
tracing::warn!(%err, "failed to accept connection");
tokio::time::sleep(Duration::from_millis(100)).await;
continue;
}
};
tracing::debug!(%peer, "new AGI connection");
let handler = Arc::clone(&self.handler);
let permit = if let Some(sem) = &semaphore {
let acquire = sem.clone().acquire_owned();
tokio::select! {
result = acquire => match result {
Ok(p) => Some(p),
Err(_) => {
tracing::error!("connection semaphore closed unexpectedly");
return Err(AgiError::Io(std::io::Error::other(
"connection semaphore closed",
)));
}
},
_ = self.shutdown_rx.changed() => {
tracing::info!("AGI server shutting down");
return Ok(());
}
}
} else {
None
};
handles.retain(|h| !h.is_finished());
handles.push(tokio::spawn(async move {
let _permit = permit;
if let Err(err) = handle_connection(handler, stream).await {
tracing::warn!(%peer, %err, "AGI session error");
}
}));
}
result = self.shutdown_rx.changed() => {
if result.is_err() || *self.shutdown_rx.borrow() {
tracing::info!("AGI server shutting down");
return Ok(());
}
}
}
}
}
}
async fn handle_connection<H: AgiHandler>(
handler: Arc<H>,
stream: tokio::net::TcpStream,
) -> Result<()> {
let (read_half, write_half) = stream.into_split();
let mut reader = BufReader::new(read_half);
let request = match tokio::time::timeout(
Duration::from_secs(30),
AgiRequest::parse_from_reader(&mut reader),
)
.await
{
Ok(result) => result?,
Err(_elapsed) => {
tracing::warn!("AGI prelude read timed out after 30s");
return Ok(());
}
};
let channel = AgiChannel::new(reader, write_half);
handler.handle(request, channel).await
}
impl<H: AgiHandler> AgiServerBuilder<H> {
pub fn bind(mut self, addr: impl Into<String>) -> Self {
self.bind_addr = addr.into();
self
}
pub fn handler(mut self, handler: H) -> Self {
self.handler = Some(handler);
self
}
pub fn max_connections(mut self, n: usize) -> Self {
self.max_connections = Some(n);
self
}
pub async fn build(self) -> Result<(AgiServer<H>, ShutdownHandle)> {
let handler = self.handler.ok_or_else(|| AgiError::InvalidConfig {
details: "handler is required".to_owned(),
})?;
let listener = TcpListener::bind(&self.bind_addr).await?;
let (shutdown_tx, shutdown_rx) = watch::channel(false);
tracing::info!(addr = %self.bind_addr, "FastAGI server bound");
let server = AgiServer {
listener,
handler: Arc::new(handler),
max_connections: self.max_connections,
shutdown_rx,
};
let handle = ShutdownHandle { tx: shutdown_tx };
Ok((server, handle))
}
}