diff --git a/axum/CHANGELOG.md b/axum/CHANGELOG.md index 5ccf1f65..167637aa 100644 --- a/axum/CHANGELOG.md +++ b/axum/CHANGELOG.md @@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 rejection type instead of `axum::response::Response` ([#3261]) - **breaking:** `axum::serve` now applies hyper's default `header_read_timeout` ([#3478]) - **added:** Implement `OptionalFromRequest` for `Multipart` ([#3220]) +- **added:** New `ListenerExt::limit_connections` allows limiting concurrent `axum::serve` connections ([#3489]) - **changed:** `serve` has an additional generic argument and can now work with any response body type, not just `axum::body::Body` ([#3205]) - **change:** Update minimum rust version to 1.78 ([#3412]) @@ -22,6 +23,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 [#3220]: https://github.com/tokio-rs/axum/pull/3220 [#3412]: https://github.com/tokio-rs/axum/pull/3412 [#3478]: https://github.com/tokio-rs/axum/pull/3478 +[#3489]: https://github.com/tokio-rs/axum/pull/3489 # 0.8.4 diff --git a/axum/src/serve/listener.rs b/axum/src/serve/listener.rs index 91effae5..94e58fe4 100644 --- a/axum/src/serve/listener.rs +++ b/axum/src/serve/listener.rs @@ -1,8 +1,17 @@ -use std::{fmt, future::Future, time::Duration}; +use std::{ + fmt, + future::Future, + pin::Pin, + sync::Arc, + task::{Context, Poll}, + time::Duration, +}; +use pin_project_lite::pin_project; use tokio::{ io::{self, AsyncRead, AsyncWrite}, net::{TcpListener, TcpStream}, + sync::{OwnedSemaphorePermit, Semaphore}, }; /// Types that can listen for connections. @@ -64,6 +73,24 @@ impl Listener for tokio::net::UnixListener { /// Extensions to [`Listener`]. pub trait ListenerExt: Listener + Sized { + /// Limit the number of concurrent connections. Once the limit has + /// been reached, no additional connections will be accepted until + /// an existing connection is closed. Listener implementations will + /// typically continue to queue incoming connections, up to an OS + /// and implementation-specific listener backlog limit. + /// + /// Compare [`tower::limit::concurrency`], which provides ways to + /// limit concurrent in-flight requests, but does not limit connections + /// that are idle or in the process of sending request headers. + /// + /// [`tower::limit::concurrency`]: https://docs.rs/tower/latest/tower/limit/concurrency/ + fn limit_connections(self, limit: usize) -> ConnLimiter { + ConnLimiter { + listener: self, + sem: Arc::new(Semaphore::new(limit)), + } + } + /// Run a mutable closure on every accepted `Io`. /// /// # Example @@ -99,6 +126,84 @@ pub trait ListenerExt: Listener + Sized { impl ListenerExt for L {} +/// Return type of [`ListenerExt::limit_connections`]. +/// +/// See that method for details. +#[derive(Debug)] +pub struct ConnLimiter { + listener: T, + sem: Arc, +} + +impl Listener for ConnLimiter { + type Io = ConnLimiterIo; + type Addr = T::Addr; + + async fn accept(&mut self) -> (Self::Io, Self::Addr) { + let permit = self.sem.clone().acquire_owned().await.unwrap(); + let (io, addr) = self.listener.accept().await; + (ConnLimiterIo { io, permit }, addr) + } + + fn local_addr(&self) -> tokio::io::Result { + self.listener.local_addr() + } +} + +pin_project! { + /// A connection counted by [`ConnLimiter`]. + /// + /// See [`ListenerExt::limit_connections`] for details. + #[derive(Debug)] + pub struct ConnLimiterIo { + #[pin] + io: T, + permit: OwnedSemaphorePermit, + } +} + +// Simply forward implementation to `io` field. +impl AsyncRead for ConnLimiterIo { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut io::ReadBuf<'_>, + ) -> Poll> { + self.project().io.poll_read(cx, buf) + } +} + +// Simply forward implementation to `io` field. +impl AsyncWrite for ConnLimiterIo { + fn is_write_vectored(&self) -> bool { + self.io.is_write_vectored() + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.project().io.poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.project().io.poll_shutdown(cx) + } + + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + self.project().io.poll_write(cx, buf) + } + + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[std::io::IoSlice<'_>], + ) -> Poll> { + self.project().io.poll_write_vectored(cx, bufs) + } +} + /// Return type of [`ListenerExt::tap_io`]. /// /// See that method for details. @@ -165,3 +270,54 @@ fn is_connection_error(e: &io::Error) -> bool { | io::ErrorKind::ConnectionReset ) } + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use tokio::{io, time}; + + use super::{Listener, ListenerExt}; + + #[tokio::test(start_paused = true)] + async fn limit_connections() { + static COUNT: AtomicUsize = AtomicUsize::new(0); + + struct MyListener; + + impl Listener for MyListener { + type Io = io::DuplexStream; + type Addr = (); + + async fn accept(&mut self) -> (Self::Io, Self::Addr) { + COUNT.fetch_add(1, Ordering::SeqCst); + (io::duplex(0).0, ()) // dummy connection + } + + fn local_addr(&self) -> io::Result { + Ok(()) + } + } + + let mut listener = MyListener.limit_connections(1); + + assert_eq!(COUNT.load(Ordering::SeqCst), 0); + + // First 'accept' succeeds immediately. + let conn1 = listener.accept().await; + assert_eq!(COUNT.load(Ordering::SeqCst), 1); + + time::timeout(time::Duration::from_secs(1), listener.accept()) + .await + .expect_err("Second 'accept' should time out."); + // It never reaches MyListener::accept to be counted. + assert_eq!(COUNT.load(Ordering::SeqCst), 1); + + // Close the first connection. + drop(conn1); + + // Now 'accept' again succeeds immediately. + let _conn2 = listener.accept().await; + assert_eq!(COUNT.load(Ordering::SeqCst), 2); + } +} diff --git a/axum/src/serve/mod.rs b/axum/src/serve/mod.rs index 412ab829..4149b2bf 100644 --- a/axum/src/serve/mod.rs +++ b/axum/src/serve/mod.rs @@ -23,7 +23,7 @@ use tower_service::Service; mod listener; -pub use self::listener::{Listener, ListenerExt, TapIo}; +pub use self::listener::{ConnLimiter, ConnLimiterIo, Listener, ListenerExt, TapIo}; /// Serve the service with the supplied listener. ///