use std::convert::Infallible;
use std::future::Future;
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::Arc;
use bytes::Bytes;
use http_body_util::combinators::BoxBody;
use hyper::body::Incoming;
pub type DispatchResponse = hyper::Response<BoxBody<Bytes, Infallible>>;
pub trait Dispatcher: Send + Sync + 'static {
fn dispatch(
&self,
req: hyper::Request<Incoming>,
) -> Pin<Box<dyn Future<Output = DispatchResponse> + Send + '_>>;
}
pub async fn serve(
addr: impl tokio::net::ToSocketAddrs,
dispatcher: impl Dispatcher,
) -> std::io::Result<()> {
let dispatcher = Arc::new(dispatcher);
let listener = tokio::net::TcpListener::bind(addr).await?;
trace_info!(
addr = %listener.local_addr().unwrap_or_else(|_| SocketAddr::from(([0, 0, 0, 0], 0))),
"A2A server listening"
);
loop {
let (stream, _peer) = match listener.accept().await {
Ok(pair) => pair,
Err(e) => {
trace_warn!(error = %e, "accept() failed; retrying");
pause_after_accept_error(&e).await;
continue;
}
};
let _ = stream.set_nodelay(true);
let io = hyper_util::rt::TokioIo::new(stream);
let dispatcher = Arc::clone(&dispatcher);
tokio::spawn(async move {
let service = hyper::service::service_fn(move |req| {
let d = Arc::clone(&dispatcher);
async move { Ok::<_, Infallible>(d.dispatch(req).await) }
});
let _ =
hyper_util::server::conn::auto::Builder::new(hyper_util::rt::TokioExecutor::new())
.serve_connection(io, service)
.await;
});
}
}
pub async fn serve_with_addr(
addr: impl tokio::net::ToSocketAddrs,
dispatcher: impl Dispatcher,
) -> std::io::Result<SocketAddr> {
let dispatcher = Arc::new(dispatcher);
let listener = tokio::net::TcpListener::bind(addr).await?;
let local_addr = listener.local_addr()?;
trace_info!(%local_addr, "A2A server listening");
tokio::spawn(async move {
loop {
let (stream, _peer) = match listener.accept().await {
Ok(pair) => pair,
Err(e) => {
trace_warn!(error = %e, "accept() failed; retrying");
pause_after_accept_error(&e).await;
continue;
}
};
let _ = stream.set_nodelay(true);
let io = hyper_util::rt::TokioIo::new(stream);
let dispatcher = Arc::clone(&dispatcher);
tokio::spawn(async move {
let service = hyper::service::service_fn(move |req| {
let d = Arc::clone(&dispatcher);
async move { Ok::<_, Infallible>(d.dispatch(req).await) }
});
let _ = hyper_util::server::conn::auto::Builder::new(
hyper_util::rt::TokioExecutor::new(),
)
.serve_connection(io, service)
.await;
});
}
});
Ok(local_addr)
}
pub(crate) async fn pause_after_accept_error(err: &std::io::Error) {
let backoff = accept_retry_backoff(err);
if !backoff.is_zero() {
tokio::time::sleep(backoff).await;
}
}
pub(crate) fn accept_retry_backoff(err: &std::io::Error) -> std::time::Duration {
match err.raw_os_error() {
Some(23 | 24) => std::time::Duration::from_millis(20),
_ => std::time::Duration::ZERO,
}
}
#[cfg(test)]
mod tests {
use super::*;
use http_body_util::{BodyExt, Empty};
use hyper_util::client::legacy::Client;
use hyper_util::rt::TokioExecutor;
#[test]
fn accept_retry_backoff_pauses_only_on_fd_exhaustion() {
use std::io::{Error, ErrorKind};
assert!(!accept_retry_backoff(&Error::from_raw_os_error(24)).is_zero());
assert!(!accept_retry_backoff(&Error::from_raw_os_error(23)).is_zero());
assert!(accept_retry_backoff(&Error::from_raw_os_error(103)).is_zero());
assert!(
accept_retry_backoff(&Error::new(ErrorKind::ConnectionAborted, "aborted")).is_zero()
);
}
struct MockDispatcher;
impl Dispatcher for MockDispatcher {
fn dispatch(
&self,
_req: hyper::Request<Incoming>,
) -> Pin<Box<dyn Future<Output = DispatchResponse> + Send + '_>> {
Box::pin(async {
let body = http_body_util::Full::new(Bytes::from("ok"));
hyper::Response::new(BoxBody::new(body.map_err(|e| match e {})))
})
}
}
#[tokio::test]
async fn serve_with_addr_returns_bound_address() {
let addr = serve_with_addr("127.0.0.1:0", MockDispatcher)
.await
.expect("server should bind");
assert_ne!(addr.port(), 0, "should bind to a real port");
assert!(addr.ip().is_loopback());
let client = Client::builder(TokioExecutor::new()).build_http::<Empty<Bytes>>();
let resp = client
.get(format!("http://{addr}/").parse().unwrap())
.await
.unwrap();
assert_eq!(resp.status(), 200);
let body = resp.into_body().collect().await.unwrap().to_bytes();
assert_eq!(&body[..], b"ok");
}
#[tokio::test]
async fn serve_with_addr_handles_multiple_connections() {
let addr = serve_with_addr("127.0.0.1:0", MockDispatcher)
.await
.expect("server should bind");
let client = Client::builder(TokioExecutor::new()).build_http::<Empty<Bytes>>();
for i in 0..3 {
let resp = client
.get(format!("http://{addr}/").parse().unwrap())
.await
.unwrap_or_else(|e| panic!("request {i} failed: {e}"));
let body = resp.into_body().collect().await.unwrap().to_bytes();
assert_eq!(&body[..], b"ok", "request {i} returned unexpected body");
}
}
#[tokio::test(start_paused = true)]
async fn backoff_pauses_only_on_fd_exhaustion() {
use std::io::Error;
use std::time::Duration;
let start = tokio::time::Instant::now();
pause_after_accept_error(&Error::from_raw_os_error(24)).await;
assert!(
start.elapsed() >= Duration::from_millis(20),
"fd exhaustion (EMFILE) must pause the accept loop; no elapsed \
time means it would busy-spin while the descriptor table is full"
);
let start = tokio::time::Instant::now();
pause_after_accept_error(&Error::from_raw_os_error(23)).await;
assert!(
start.elapsed() >= Duration::from_millis(20),
"ENFILE must pause for the same reason"
);
let start = tokio::time::Instant::now();
pause_after_accept_error(&Error::from_raw_os_error(103)).await;
assert_eq!(
start.elapsed(),
Duration::ZERO,
"a transient per-connection error must retry at once"
);
}
#[tokio::test]
async fn serve_binds_and_answers_requests() {
let probe = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("reserve a port");
let addr = probe.local_addr().expect("addr");
drop(probe);
let server = tokio::spawn(async move { serve(addr, MockDispatcher).await });
let client = Client::builder(TokioExecutor::new()).build_http::<Empty<Bytes>>();
let mut last_err = None;
let mut response = None;
for _ in 0..50 {
match client
.get(format!("http://{addr}/").parse().expect("uri"))
.await
{
Ok(resp) => {
response = Some(resp);
break;
}
Err(e) => {
last_err = Some(e);
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
}
}
}
let resp = response.unwrap_or_else(|| {
panic!(
"nothing served on {addr} after ~1s; `serve` never bound it. \
last error: {last_err:?}"
)
});
assert_eq!(
resp.status(),
200,
"the dispatcher's response must come back"
);
let body = resp.into_body().collect().await.expect("body").to_bytes();
assert_eq!(&body[..], b"ok", "the body must come from MockDispatcher");
server.abort();
}
}