feat: Avoid spawning async-io thread if possible

This commit adds conditions that prevent the async-io thread from being
spawned in certain cases. This is important because sometimes async-io
is used in environments where spawning additional threads can interfere
with other functionality in the program; for example, in DPDK poll
loops.

If there is only one async-io::block_on invocation, there is no need to
spawn the "async-io" thread; we can just modify block_on to not give up
the reactor if it can detect that it is the only block_on call running
in the program. From there, we also need to avoid spawning the async-io
thread if there are no resources left when block_on exits. This should
allow async-io to not spawn additional threads in single-threaded
programs where it is not needed.

Signed-off-by: John Nunley <dev@notgull.net>

Author: notgull

Cargo.toml

@@ -45,8 +45,10 @@ autocfg = "1"
 [dev-dependencies]
 async-channel = "2.0.0"
 async-net = "2.0.0"
+async-task = "4.7.1"
 blocking = "1"
 criterion = { version = "0.7", default-features = false, features = ["cargo_bench_support"] }
+fastrand = "2.3.0"
 getrandom = "0.3"
 signal-hook = "0.3"
 tempfile = "3"

src/driver.rs

@@ -1,7 +1,7 @@
 use std::cell::{Cell, RefCell};
 use std::future::Future;
 use std::pin::pin;
-use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+use std::sync::atomic::{fence, AtomicBool, AtomicUsize, Ordering};
 use std::sync::{Arc, OnceLock};
 use std::task::{Context, Poll, Waker};
 use std::thread;
@@ -14,10 +14,14 @@ use crate::reactor::Reactor;
 /// Number of currently active `block_on()` invocations.
 static BLOCK_ON_COUNT: AtomicUsize = AtomicUsize::new(0);
 
+/// Number of currently alive `Async` or `Timer` instances.
+static ALIVE_RESOURCE_COUNT: AtomicUsize = AtomicUsize::new(0);
+
+/// Raw unparker for the "async-io" thread.
+static UNPARKER: OnceLock<parking::Unparker> = OnceLock::new();
+
 /// Unparker for the "async-io" thread.
 fn unparker() -> &'static parking::Unparker {
-    static UNPARKER: OnceLock<parking::Unparker> = OnceLock::new();
-
     UNPARKER.get_or_init(|| {
         let (parker, unparker) = parking::pair();
 
@@ -35,9 +39,35 @@ fn unparker() -> &'static parking::Unparker {
     })
 }
 
-/// Initializes the "async-io" thread.
-pub(crate) fn init() {
-    let _ = unparker();
+/// Tell if the "async-io" thread is spawned.
+pub fn is_async_io_thread_spawned() -> bool {
+    UNPARKER.get().is_some()
+}
+
+/// Decrement the number of currently alive `Async` or `Timer` instances.
+#[inline]
+pub(crate) fn decrement_alive() {
+    ALIVE_RESOURCE_COUNT.fetch_sub(1, Ordering::Relaxed);
+}
+
+/// Increments the number of currently alive `Async` or `Timer` instances.
+#[inline]
+pub(crate) fn increment_alive() {
+    // If this is the first resource, spawn the `async-io` thread if necessary.
+    if ALIVE_RESOURCE_COUNT.fetch_add(1, Ordering::Relaxed) == 0 {
+        init();
+    }
+
+    #[cold]
+    fn init() {
+        // Emit a fence to ensure everything is in order.
+        fence(Ordering::SeqCst);
+
+        // If there are no `block_on()` calls, spawn the `async-io` thread.
+        if BLOCK_ON_COUNT.load(Ordering::SeqCst) == 0 {
+            unparker().unpark();
+        }
+    }
 }
 
 /// The main loop for the "async-io" thread.
@@ -118,12 +148,20 @@ pub fn block_on<T>(future: impl Future<Output = T>) -> T {
     let _enter = span.enter();
 
     // Increment `BLOCK_ON_COUNT` so that the "async-io" thread becomes less aggressive.
-    BLOCK_ON_COUNT.fetch_add(1, Ordering::SeqCst);
+    if BLOCK_ON_COUNT.fetch_add(1, Ordering::SeqCst) > 0 {
+        // There are other `block_on()` calls active, "async-io" must be spawned
+        // if it isn't already.
+        let _ = unparker();
+    }
 
     // Make sure to decrement `BLOCK_ON_COUNT` at the end and wake the "async-io" thread.
     let _guard = CallOnDrop(|| {
         BLOCK_ON_COUNT.fetch_sub(1, Ordering::SeqCst);
-        unparker().unpark();
+
+        // Only wake the `async-io` thread iff there are live resources.
+        if ALIVE_RESOURCE_COUNT.load(Ordering::SeqCst) > 0 {
+            unparker().unpark();
+        }
     });
 
     // Creates a parker and an associated waker that unparks it.
@@ -263,8 +301,11 @@ pub fn block_on<T>(future: impl Future<Output = T>) -> T {
                         break;
                     }
 
-                    // Check if this thread been handling I/O events for a long time.
-                    if start.elapsed() > Duration::from_micros(500) {
+                    // Check if this thread been handling I/O events for a long time
+                    // and if there are other threads waiting for I/O events.
+                    if start.elapsed() > Duration::from_micros(500)
+                        && BLOCK_ON_COUNT.load(Ordering::SeqCst) > 1
+                    {
                         #[cfg(feature = "tracing")]
                         tracing::trace!("stops hogging the reactor");
 

src/lib.rs

@@ -25,11 +25,31 @@
 //! [`block_on()`] function. The "async-io" thread is therefore just a fallback mechanism
 //! processing I/O events in case no other threads are.
 //!
+//! This thread is only spawned under either of the following conditions:
+//!
+//! - [`Async`]s or [`Timer`]s are created without an active [`block_on`] call.
+//! - More than one thread is calling [`block_on`] at a time.
+//!
+//! Therefore, if you would like to prevent this thread from being spawned, ensure
+//! that there is only one thread calling [`block_on`] at a time. This thread
+//! will handle all I/O events in lieu of the "async-io" thread.
+//!
+//! **Note:** If you are using [`async-process`], keep in mind that [`async-process`]'s
+//! reaper thread calls [`block_on`] on Unix platforms in order to reap child
+//! processes. If you are using [`async-process`] and are trying to avoid spawning
+//! the "async-io" thread, ensure that your executor is running the [`async_process::driver()`]
+//! future. In addition, the executor backing [`smol::spawn()`] as well as the
+//! multithreaded version of [`smol_macros::main!`] use [`block_on`] as well.
+//!
 //! [epoll]: https://en.wikipedia.org/wiki/Epoll
 //! [kqueue]: https://en.wikipedia.org/wiki/Kqueue
 //! [event ports]: https://illumos.org/man/port_create
 //! [IOCP]: https://learn.microsoft.com/en-us/windows/win32/fileio/i-o-completion-ports
 //! [`polling`]: https://docs.rs/polling
+//! [`async-process`]: https://docs.rs/async-process
+//! [`async_process::driver()`]: https://docs.rs/async-process/latest/async_process/fn.driver.html
+//! [`smol::spawn()`]: https://docs.rs/smol/latest/smol/fn.spawn.html
+//! [`smol_macros::main!`]: https://docs.rs/smol-macros/latest/smol_macros/macro.main.html
 //!
 //! # Examples
 //!
@@ -79,6 +99,10 @@ use std::{
 #[cfg(windows)]
 use std::os::windows::io::{AsRawSocket, AsSocket, BorrowedSocket, OwnedSocket, RawSocket};
 
+// Not public API.
+#[doc(hidden)]
+pub use driver::is_async_io_thread_spawned;
+
 use futures_io::{AsyncRead, AsyncWrite};
 use futures_lite::stream::{self, Stream};
 
@@ -192,6 +216,8 @@ impl Timer {
     /// # });
     /// ```
     pub fn never() -> Timer {
+        crate::driver::increment_alive();
+
         Timer {
             id_and_waker: None,
             when: None,
@@ -271,6 +297,8 @@ impl Timer {
     /// # });
     /// ```
     pub fn interval_at(start: Instant, period: Duration) -> Timer {
+        crate::driver::increment_alive();
+
         Timer {
             id_and_waker: None,
             when: Some(start),
@@ -458,6 +486,8 @@ impl Drop for Timer {
             // Deregister the timer from the reactor.
             Reactor::get().remove_timer(when, id);
         }
+
+        crate::driver::decrement_alive();
     }
 }
 
@@ -680,6 +710,8 @@ impl<T: AsFd> Async<T> {
     /// it is not set. If not set to non-blocking mode, I/O operations may block the current thread
     /// and cause a deadlock in an asynchronous context.
     pub fn new_nonblocking(io: T) -> io::Result<Async<T>> {
+        crate::driver::increment_alive();
+
         // SAFETY: It is impossible to drop the I/O source while it is registered through
         // this type.
         let registration = unsafe { Registration::new(io.as_fd()) };
@@ -774,6 +806,8 @@ impl<T: AsSocket> Async<T> {
     /// it is not set. If not set to non-blocking mode, I/O operations may block the current thread
     /// and cause a deadlock in an asynchronous context.
     pub fn new_nonblocking(io: T) -> io::Result<Async<T>> {
+        crate::driver::increment_alive();
+
         // Create the registration.
         //
         // SAFETY: It is impossible to drop the I/O source while it is registered through
@@ -1165,6 +1199,8 @@ impl<T> Drop for Async<T> {
             // Drop the I/O handle to close it.
             self.io.take();
         }
+
+        crate::driver::decrement_alive();
     }
 }
 

src/reactor.rs

@@ -93,16 +93,13 @@ impl Reactor {
     pub(crate) fn get() -> &'static Reactor {
         static REACTOR: OnceLock<Reactor> = OnceLock::new();
 
-        REACTOR.get_or_init(|| {
-            crate::driver::init();
-            Reactor {
-                poller: Poller::new().expect("cannot initialize I/O event notification"),
-                ticker: AtomicUsize::new(0),
-                sources: Mutex::new(Slab::new()),
-                events: Mutex::new(Events::new()),
-                timers: Mutex::new(BTreeMap::new()),
-                timer_ops: ConcurrentQueue::bounded(TIMER_QUEUE_SIZE),
-            }
+        REACTOR.get_or_init(|| Reactor {
+            poller: Poller::new().expect("cannot initialize I/O event notification"),
+            ticker: AtomicUsize::new(0),
+            sources: Mutex::new(Slab::new()),
+            events: Mutex::new(Events::new()),
+            timers: Mutex::new(BTreeMap::new()),
+            timer_ops: ConcurrentQueue::bounded(TIMER_QUEUE_SIZE),
         })
     }
 

tests/async.rs

@@ -6,7 +6,7 @@ use std::sync::Arc;
 use std::thread;
 use std::time::Duration;
 
-use async_io::{Async, Timer};
+use async_io::{is_async_io_thread_spawned, Async, Timer};
 use futures_lite::{future, prelude::*};
 #[cfg(unix)]
 use tempfile::tempdir;
@@ -36,6 +36,7 @@ fn spawn<T: Send + 'static>(
 fn tcp_connect() -> io::Result<()> {
     future::block_on(async {
         let listener = Async::<TcpListener>::bind(([127, 0, 0, 1], 0))?;
+        assert!(is_async_io_thread_spawned());
         let addr = listener.get_ref().local_addr()?;
         let task = spawn(async move { listener.accept().await });
 

tests/driver_not_spawned.rs

@@ -0,0 +1,114 @@
+use async_channel::bounded;
+use async_io::{block_on, is_async_io_thread_spawned, Async, Timer};
+use futures_lite::prelude::*;
+use std::net::{TcpListener, TcpStream};
+use std::sync::atomic;
+use std::thread;
+use std::time::Duration;
+
+const TIMER_COUNT: u64 = 10_000;
+const IO_COUNT: usize = 200;
+
+static TOTAL_IO: atomic::AtomicU64 = atomic::AtomicU64::new(0);
+
+#[test]
+fn test_one_block_on() {
+    // The driver thread should not be spawned.
+    assert!(!is_async_io_thread_spawned());
+
+    // Running smol::block_on() should not spawn the driver thread.
+    block_on(async {
+        assert!(!is_async_io_thread_spawned());
+
+        // We should be able to handle a lot of timers and sources.
+        let mut timers = Vec::new();
+        for i in 1..TIMER_COUNT {
+            timers.push(Timer::after(Duration::from_millis(i / 5)));
+        }
+
+        let (spawner, executor) = async_channel::unbounded();
+        let mut tasks = Vec::new();
+
+        for _ in 0..IO_COUNT {
+            let (runnable, task) = async_task::Builder::new().propagate_panic(true).spawn(
+                move |_| async move {
+                    let mut rng = fastrand::Rng::new();
+
+                    // Create a TCP pipe and send bytes to and from.
+                    let listener = Async::<TcpListener>::bind(([127, 0, 0, 1], 0))?;
+                    let stream1 =
+                        Async::<TcpStream>::connect(listener.get_ref().local_addr()?).await?;
+                    let stream2 = listener.accept().await?.0;
+
+                    let mut bytes = [0u8; 64];
+                    let mut read_buffer = [0u8; 64];
+
+                    loop {
+                        rng.fill(&mut bytes);
+
+                        Timer::after(Duration::from_micros(rng.u64(..1_000))).await;
+                        (&stream1).write_all(&bytes).await?;
+                        Timer::after(Duration::from_micros(rng.u64(..1_000))).await;
+                        (&stream2).read_exact(&mut read_buffer).await?;
+
+                        assert_eq!(bytes, read_buffer);
+                        TOTAL_IO.fetch_add(bytes.len() as u64, atomic::Ordering::Relaxed);
+                        futures_lite::future::yield_now().await;
+                    }
+
+                    #[allow(unreachable_code)]
+                    std::io::Result::Ok(())
+                },
+                {
+                    let spawner = spawner.clone();
+                    move |task| {
+                        spawner.try_send(task).ok();
+                    }
+                },
+            );
+            runnable.schedule();
+            tasks.push(task);
+        }
+
+        // Future to process timers.
+        let process_timers = async move {
+            for timer in timers {
+                timer.await;
+            }
+
+            // After the timer is done, cancel every task.
+            for task in tasks {
+                if let Some(res) = task.cancel().await {
+                    res.unwrap();
+                }
+            }
+        };
+
+        // Future to process sources.
+        let process_sources = async move {
+            while let Ok(task) = executor.recv().await {
+                task.run();
+                futures_lite::future::yield_now().await;
+            }
+        };
+
+        process_timers.or(process_sources).await;
+
+        // Spawning another thread should spawn the driver thread.
+        let (thread_spawned_signal, thread_spawned) = bounded::<()>(1);
+        let (signal, shutdown) = bounded::<()>(1);
+        thread::spawn(move || {
+            block_on(async move {
+                thread_spawned_signal.send(()).await.unwrap();
+                shutdown.recv().await.unwrap();
+            })
+        });
+
+        thread_spawned.recv().await.unwrap();
+        assert!(is_async_io_thread_spawned());
+
+        // Stopping the other thread should not stop the driver thread.
+        signal.send(()).await.unwrap();
+        assert!(is_async_io_thread_spawned());
+    });
+}

tests/timer.rs

@@ -4,7 +4,7 @@ use std::sync::{Arc, Mutex};
 use std::thread;
 use std::time::{Duration, Instant};
 
-use async_io::Timer;
+use async_io::{is_async_io_thread_spawned, Timer};
 use futures_lite::{future, FutureExt, StreamExt};
 
 fn spawn<T: Send + 'static>(
@@ -26,6 +26,7 @@ fn smoke() {
     future::block_on(async {
         let start = Instant::now();
         Timer::after(Duration::from_secs(1)).await;
+        assert!(is_async_io_thread_spawned());
         assert!(start.elapsed() >= Duration::from_secs(1));
     });
 }