Skip to main content

blocking/
lib.rs

1//! A thread pool for isolating blocking I/O in async programs.
2//!
3//! Sometimes there's no way to avoid blocking I/O. Consider files or stdin, which have weak async
4//! support on modern operating systems. While [IOCP], [AIO], and [io_uring] are possible
5//! solutions, they're not always available or ideal.
6//!
7//! Since blocking is not allowed inside futures, we must move blocking I/O onto a special thread
8//! pool provided by this crate. The pool dynamically spawns and stops threads depending on the
9//! current number of running I/O jobs.
10//!
11//! Note that there is a limit on the number of active threads. Once that limit is hit, a running
12//! job has to finish before others get a chance to run. When a thread is idle, it waits for the
13//! next job or shuts down after a certain timeout.
14//!
15//! The default number of threads (set to 500) can be altered by setting `BLOCKING_MAX_THREADS` environment
16//! variable with value between 1 and 10000. This can also be set, at runtime, via the
17//! [`set_max_blocking_threads`] function.
18//!
19//! [IOCP]: https://en.wikipedia.org/wiki/Input/output_completion_port
20//! [AIO]: http://man7.org/linux/man-pages/man2/io_submit.2.html
21//! [io_uring]: https://lwn.net/Articles/776703
22//!
23//! # Examples
24//!
25//! Read the contents of a file:
26//!
27//! ```no_run
28//! use blocking::unblock;
29//! use std::fs;
30//!
31//! # futures_lite::future::block_on(async {
32//! let contents = unblock(|| fs::read_to_string("file.txt")).await?;
33//! println!("{}", contents);
34//! # std::io::Result::Ok(()) });
35//! ```
36//!
37//! Read a file and pipe its contents to stdout:
38//!
39//! ```no_run
40//! use blocking::{unblock, Unblock};
41//! use futures_lite::io;
42//! use std::fs::File;
43//!
44//! # futures_lite::future::block_on(async {
45//! let input = unblock(|| File::open("file.txt")).await?;
46//! let input = Unblock::new(input);
47//! let mut output = Unblock::new(std::io::stdout());
48//!
49//! io::copy(input, &mut output).await?;
50//! # std::io::Result::Ok(()) });
51//! ```
52//!
53//! Iterate over the contents of a directory:
54//!
55//! ```no_run
56//! use blocking::Unblock;
57//! use futures_lite::prelude::*;
58//! use std::fs;
59//!
60//! # futures_lite::future::block_on(async {
61//! let mut dir = Unblock::new(fs::read_dir(".")?);
62//! while let Some(item) = dir.next().await {
63//!     println!("{}", item?.file_name().to_string_lossy());
64//! }
65//! # std::io::Result::Ok(()) });
66//! ```
67//!
68//! Spawn a process:
69//!
70//! ```no_run
71//! use blocking::unblock;
72//! use std::process::Command;
73//!
74//! # futures_lite::future::block_on(async {
75//! let out = unblock(|| Command::new("dir").output()).await?;
76//! # std::io::Result::Ok(()) });
77//! ```
78
79#![warn(missing_docs, missing_debug_implementations, rust_2018_idioms)]
80#![forbid(unsafe_code)]
81#![doc(
82    html_favicon_url = "https://raw.githubusercontent.com/smol-rs/smol/master/assets/images/logo_fullsize_transparent.png"
83)]
84#![doc(
85    html_logo_url = "https://raw.githubusercontent.com/smol-rs/smol/master/assets/images/logo_fullsize_transparent.png"
86)]
87
88use std::any::Any;
89use std::collections::VecDeque;
90use std::fmt;
91use std::io::{self, Read, Seek, SeekFrom, Write};
92use std::num::NonZeroUsize;
93use std::panic;
94use std::pin::Pin;
95use std::sync::atomic::{AtomicUsize, Ordering};
96use std::sync::{Condvar, Mutex, MutexGuard, PoisonError};
97use std::task::{Context, Poll};
98use std::thread;
99use std::time::Duration;
100
101#[cfg(not(target_family = "wasm"))]
102use std::env;
103
104use async_channel::{bounded, Receiver};
105use async_task::Runnable;
106use futures_io::{AsyncRead, AsyncSeek, AsyncWrite};
107use futures_lite::{
108    future::{self, Future},
109    ready,
110    stream::Stream,
111};
112use piper::{pipe, Reader, Writer};
113
114#[doc(no_inline)]
115pub use async_task::Task;
116
117/// Default value for max threads that Executor can grow to
118#[cfg(not(target_family = "wasm"))]
119const DEFAULT_MAX_THREADS: NonZeroUsize = {
120    if let Some(size) = NonZeroUsize::new(500) {
121        size
122    } else {
123        panic!("DEFAULT_MAX_THREADS is non-zero");
124    }
125};
126
127/// Minimum value for max threads config
128#[cfg(not(target_family = "wasm"))]
129const MIN_MAX_THREADS: usize = 1;
130
131/// Maximum value for max threads config
132#[cfg(not(target_family = "wasm"))]
133const MAX_MAX_THREADS: usize = 10000;
134
135/// Env variable that allows to override default value for max threads.
136#[cfg(not(target_family = "wasm"))]
137const MAX_THREADS_ENV: &str = "BLOCKING_MAX_THREADS";
138
139/// Set the maximum number of threads used by the backing thread pool.
140///
141/// # Example
142///
143/// ```no_run
144/// use blocking::unblock;
145/// use std::fs::{read_dir, File};
146/// use std::io::prelude::*;
147/// # use std::num::NonZeroUsize;
148///
149/// blocking::set_max_blocking_threads(NonZeroUsize::new(100).unwrap());
150///
151/// # fn test() -> std::io::Result<()> {
152/// let mut files = Vec::new();
153/// for entry in read_dir("/path/to/large/directory").unwrap() {
154///     files.push(unblock(move || -> std::io::Result<String> {
155///         let mut contents = String::new();
156///         let mut file = File::open(entry?.path())?;
157///         file.read_to_string(&mut contents)?;
158///         Ok(contents)
159///     }));
160/// }
161/// # Ok(())
162/// # }
163/// ```
164pub fn set_max_blocking_threads(threads: NonZeroUsize) {
165    let executor = Executor::get();
166    let mut inner = executor
167        .inner
168        .lock()
169        .unwrap_or_else(PoisonError::into_inner);
170    let old_limit = inner.thread_limit;
171    inner.thread_limit = Some(threads);
172    if let Some(old_limit) = old_limit {
173        // If the limit has decreased, wake up all threads to terminate those over
174        // the new limit.
175        if old_limit > threads {
176            executor.cvar.notify_all();
177        }
178    }
179}
180
181/// The blocking executor.
182struct Executor {
183    /// Inner state of the executor.
184    inner: Mutex<Inner>,
185
186    /// Used to put idle threads to sleep and wake them up when new work comes in.
187    cvar: Condvar,
188}
189
190/// Inner state of the blocking executor.
191struct Inner {
192    /// Number of idle threads in the pool.
193    ///
194    /// Idle threads are sleeping, waiting to get a task to run.
195    idle_count: usize,
196
197    /// Total number of threads in the pool.
198    ///
199    /// This is the number of idle threads + the number of active threads.
200    thread_count: usize,
201
202    /// The queue of blocking tasks.
203    queue: VecDeque<Runnable>,
204
205    /// Maximum number of threads in the pool
206    thread_limit: Option<NonZeroUsize>,
207}
208
209impl Executor {
210    #[cfg(not(target_family = "wasm"))]
211    fn max_threads() -> NonZeroUsize {
212        match env::var(MAX_THREADS_ENV) {
213            Ok(v) => v
214                .parse::<usize>()
215                .ok()
216                .and_then(|v| NonZeroUsize::new(v.clamp(MIN_MAX_THREADS, MAX_MAX_THREADS)))
217                .unwrap_or(DEFAULT_MAX_THREADS),
218            Err(_) => DEFAULT_MAX_THREADS,
219        }
220    }
221
222    #[cfg(target_family = "wasm")]
223    fn max_threads() -> NonZeroUsize {
224        NonZeroUsize::new(1).unwrap()
225    }
226
227    /// Get a reference to the global executor.
228    #[inline]
229    fn get() -> &'static Self {
230        #[cfg(not(target_family = "wasm"))]
231        {
232            static EXECUTOR: Executor = Executor {
233                inner: Mutex::new(Inner {
234                    idle_count: 0,
235                    thread_count: 0,
236                    queue: VecDeque::new(),
237                    thread_limit: None,
238                }),
239                cvar: Condvar::new(),
240            };
241
242            &EXECUTOR
243        }
244
245        #[cfg(target_family = "wasm")]
246        panic!("cannot spawn a blocking task on WASM")
247    }
248
249    /// Spawns a future onto this executor.
250    ///
251    /// Returns a [`Task`] handle for the spawned task.
252    fn spawn<T: Send + 'static>(future: impl Future<Output = T> + Send + 'static) -> Task<T> {
253        let (runnable, task) = async_task::Builder::new().propagate_panic(true).spawn(
254            move |()| future,
255            |r| {
256                // Initialize the executor if we haven't already.
257                let executor = Self::get();
258
259                // Schedule the task on our executor.
260                executor.schedule(r)
261            },
262        );
263        runnable.schedule();
264        task
265    }
266
267    /// Runs the main loop on the current thread.
268    ///
269    /// This function runs blocking tasks until it becomes idle and times out.
270    fn main_loop(&'static self) {
271        #[cfg(feature = "tracing")]
272        let _span = tracing::trace_span!("blocking::main_loop").entered();
273
274        let mut inner = self.inner.lock().unwrap_or_else(PoisonError::into_inner);
275        loop {
276            // This thread is not idle anymore because it's going to run tasks.
277            inner.idle_count -= 1;
278
279            // Run tasks in the queue.
280            while let Some(runnable) = inner.queue.pop_front() {
281                // We have found a task - grow the pool if needed.
282                self.grow_pool(inner);
283
284                // Run the task.
285                panic::catch_unwind(|| runnable.run()).ok();
286
287                // Re-lock the inner state and continue.
288                inner = self.inner.lock().unwrap_or_else(PoisonError::into_inner);
289            }
290
291            // This thread is now becoming idle.
292            inner.idle_count += 1;
293
294            // Put the thread to sleep until another task is scheduled.
295            let timeout = Duration::from_millis(500);
296            #[cfg(feature = "tracing")]
297            tracing::trace!(?timeout, "going to sleep");
298            let (lock, res) = self.cvar.wait_timeout(inner, timeout).unwrap();
299            inner = lock;
300
301            // If there are too many threads active in the pool, stop this thread.
302            if (Some(inner.thread_count) > inner.thread_limit.map(NonZeroUsize::get))
303                // If there are no tasks after a while, stop this thread.
304                && (res.timed_out() && inner.queue.is_empty())
305            {
306                inner.idle_count -= 1;
307                inner.thread_count -= 1;
308                break;
309            }
310
311            #[cfg(feature = "tracing")]
312            tracing::trace!("notified");
313        }
314
315        #[cfg(feature = "tracing")]
316        tracing::trace!("shutting down due to lack of tasks");
317    }
318
319    /// Schedules a runnable task for execution.
320    fn schedule(&'static self, runnable: Runnable) {
321        let mut inner = self.inner.lock().unwrap_or_else(PoisonError::into_inner);
322        inner.queue.push_back(runnable);
323
324        // Notify a sleeping thread and spawn more threads if needed.
325        self.cvar.notify_one();
326        self.grow_pool(inner);
327    }
328
329    /// Spawns more blocking threads if the pool is overloaded with work.
330    fn grow_pool(&'static self, mut inner: MutexGuard<'static, Inner>) {
331        #[cfg(feature = "tracing")]
332        let _span = tracing::trace_span!(
333            "grow_pool",
334            queue_len = inner.queue.len(),
335            idle_count = inner.idle_count,
336            thread_count = inner.thread_count,
337        )
338        .entered();
339
340        let thread_limit = inner
341            .thread_limit
342            .get_or_insert_with(Self::max_threads)
343            .get();
344
345        // If runnable tasks greatly outnumber idle threads and there aren't too many threads
346        // already, then be aggressive: wake all idle threads and spawn one more thread.
347        while inner.queue.len() > inner.idle_count * 5 && inner.thread_count < thread_limit {
348            #[cfg(feature = "tracing")]
349            tracing::trace!("spawning a new thread to handle blocking tasks");
350
351            // The new thread starts in idle state.
352            inner.idle_count += 1;
353            inner.thread_count += 1;
354
355            // Notify all existing idle threads because we need to hurry up.
356            self.cvar.notify_all();
357
358            // Generate a new thread ID.
359            static ID: AtomicUsize = AtomicUsize::new(1);
360            let id = ID.fetch_add(1, Ordering::Relaxed);
361
362            // Spawn the new thread.
363            if let Err(_e) = thread::Builder::new()
364                .name(format!("blocking-{id}"))
365                .spawn(move || self.main_loop())
366            {
367                // We were unable to spawn the thread, so we need to undo the state changes.
368                #[cfg(feature = "tracing")]
369                tracing::error!("failed to spawn a blocking thread: {}", _e);
370                inner.idle_count -= 1;
371                inner.thread_count -= 1;
372
373                // The current number of threads is likely to be the system's upper limit, so update
374                // thread_limit accordingly.
375                inner.thread_limit = {
376                    let new_limit = inner.thread_count;
377
378                    // If the limit is about to be set to zero, set it to one instead so that if,
379                    // in the future, we are able to spawn more threads, we will be able to do so.
380                    Some(NonZeroUsize::new(new_limit).unwrap_or_else(|| {
381                        #[cfg(feature = "tracing")]
382                        tracing::warn!(
383                            "attempted to lower thread_limit to zero; setting to one instead"
384                        );
385                        NonZeroUsize::new(1).unwrap()
386                    }))
387                };
388            }
389        }
390    }
391}
392
393/// Runs blocking code on a thread pool.
394///
395/// # Examples
396///
397/// Read the contents of a file:
398///
399/// ```no_run
400/// use blocking::unblock;
401/// use std::fs;
402///
403/// # futures_lite::future::block_on(async {
404/// let contents = unblock(|| fs::read_to_string("file.txt")).await?;
405/// # std::io::Result::Ok(()) });
406/// ```
407///
408/// Spawn a process:
409///
410/// ```no_run
411/// use blocking::unblock;
412/// use std::process::Command;
413///
414/// # futures_lite::future::block_on(async {
415/// let out = unblock(|| Command::new("dir").output()).await?;
416/// # std::io::Result::Ok(()) });
417/// ```
418pub fn unblock<T, F>(f: F) -> Task<T>
419where
420    F: FnOnce() -> T + Send + 'static,
421    T: Send + 'static,
422{
423    Executor::spawn(async move { f() })
424}
425
426/// Runs blocking I/O on a thread pool.
427///
428/// Blocking I/O must be isolated from async code. This type moves blocking I/O operations onto a
429/// special thread pool while exposing a familiar async interface.
430///
431/// This type implements traits [`Stream`], [`AsyncRead`], [`AsyncWrite`], or [`AsyncSeek`] if the
432/// inner type implements [`Iterator`], [`Read`], [`Write`], or [`Seek`], respectively.
433///
434/// # Caveats
435///
436/// [`Unblock`] is a low-level primitive, and as such it comes with some caveats.
437///
438/// For higher-level primitives built on top of [`Unblock`], look into [`async-fs`] or
439/// [`async-process`] (on Windows).
440///
441/// [`async-fs`]: https://github.com/smol-rs/async-fs
442/// [`async-process`]: https://github.com/smol-rs/async-process
443///
444/// [`Unblock`] communicates with I/O operations on the thread pool through a pipe. That means an
445/// async read/write operation simply receives/sends some bytes from/into the pipe. When in reading
446/// mode, the thread pool reads bytes from the I/O handle and forwards them into the pipe until it
447/// becomes full. When in writing mode, the thread pool reads bytes from the pipe and forwards them
448/// into the I/O handle.
449///
450/// Use [`Unblock::with_capacity()`] to configure the capacity of the pipe.
451///
452/// ### Reading
453///
454/// If you create an [`Unblock`]`<`[`Stdin`][`std::io::Stdin`]`>`, read some bytes from it,
455/// and then drop it, a blocked read operation may keep hanging on the thread pool. The next
456/// attempt to read from stdin will lose bytes read by the hanging operation. This is a difficult
457/// problem to solve, so make sure you only use a single stdin handle for the duration of the
458/// entire program.
459///
460/// ### Writing
461///
462/// If writing data through the [`AsyncWrite`] trait, make sure to flush before dropping the
463/// [`Unblock`] handle or some buffered data might get lost.
464///
465/// ### Seeking
466///
467/// Because of buffering in the pipe, if [`Unblock`] wraps a [`File`][`std::fs::File`], a single
468/// read operation may move the file cursor farther than is the span of the operation. In fact,
469/// reading just keeps going in the background until the pipe gets full. Keep this mind when
470/// using [`AsyncSeek`] with [relative][`SeekFrom::Current`] offsets.
471///
472/// # Examples
473///
474/// ```
475/// use blocking::Unblock;
476/// use futures_lite::prelude::*;
477///
478/// # futures_lite::future::block_on(async {
479/// let mut stdout = Unblock::new(std::io::stdout());
480/// stdout.write_all(b"Hello world!").await?;
481/// stdout.flush().await?;
482/// # std::io::Result::Ok(()) });
483/// ```
484pub struct Unblock<T> {
485    state: State<T>,
486    cap: Option<usize>,
487}
488
489impl<T> Unblock<T> {
490    /// Wraps a blocking I/O handle into the async [`Unblock`] interface.
491    ///
492    /// # Examples
493    ///
494    /// ```no_run
495    /// use blocking::Unblock;
496    ///
497    /// let stdin = Unblock::new(std::io::stdin());
498    /// ```
499    pub fn new(io: T) -> Unblock<T> {
500        Unblock {
501            state: State::Idle(Some(Box::new(io))),
502            cap: None,
503        }
504    }
505
506    /// Wraps a blocking I/O handle into the async [`Unblock`] interface with a custom buffer
507    /// capacity.
508    ///
509    /// When communicating with the inner [`Stream`]/[`Read`]/[`Write`] type from async code, data
510    /// transferred between blocking and async code goes through a buffer of limited capacity. This
511    /// constructor configures that capacity.
512    ///
513    /// The default capacity is:
514    ///
515    /// * For [`Iterator`] types: 8192 items.
516    /// * For [`Read`]/[`Write`] types: 8 MB.
517    ///
518    /// # Examples
519    ///
520    /// ```no_run
521    /// use blocking::Unblock;
522    ///
523    /// let stdout = Unblock::with_capacity(64 * 1024, std::io::stdout());
524    /// ```
525    pub fn with_capacity(cap: usize, io: T) -> Unblock<T> {
526        Unblock {
527            state: State::Idle(Some(Box::new(io))),
528            cap: Some(cap),
529        }
530    }
531
532    /// Gets a mutable reference to the blocking I/O handle.
533    ///
534    /// This is an async method because the I/O handle might be on the thread pool and needs to
535    /// be moved onto the current thread before we can get a reference to it.
536    ///
537    /// # Examples
538    ///
539    /// ```no_run
540    /// use blocking::{unblock, Unblock};
541    /// use std::fs::File;
542    ///
543    /// # futures_lite::future::block_on(async {
544    /// let file = unblock(|| File::create("file.txt")).await?;
545    /// let mut file = Unblock::new(file);
546    ///
547    /// let metadata = file.get_mut().await.metadata()?;
548    /// # std::io::Result::Ok(()) });
549    /// ```
550    pub async fn get_mut(&mut self) -> &mut T {
551        // Wait for the running task to stop and ignore I/O errors if there are any.
552        future::poll_fn(|cx| self.poll_stop(cx)).await.ok();
553
554        // Assume idle state and get a reference to the inner value.
555        match &mut self.state {
556            State::Idle(t) => t.as_mut().expect("inner value was taken out"),
557            State::WithMut(..)
558            | State::Streaming(..)
559            | State::Reading(..)
560            | State::Writing(..)
561            | State::Seeking(..) => {
562                unreachable!("when stopped, the state machine must be in idle state");
563            }
564        }
565    }
566
567    /// Performs a blocking operation on the I/O handle.
568    ///
569    /// # Examples
570    ///
571    /// ```no_run
572    /// use blocking::{unblock, Unblock};
573    /// use std::fs::File;
574    ///
575    /// # futures_lite::future::block_on(async {
576    /// let file = unblock(|| File::create("file.txt")).await?;
577    /// let mut file = Unblock::new(file);
578    ///
579    /// let metadata = file.with_mut(|f| f.metadata()).await?;
580    /// # std::io::Result::Ok(()) });
581    /// ```
582    pub async fn with_mut<R, F>(&mut self, op: F) -> R
583    where
584        F: FnOnce(&mut T) -> R + Send + 'static,
585        R: Send + 'static,
586        T: Send + 'static,
587    {
588        // Wait for the running task to stop and ignore I/O errors if there are any.
589        future::poll_fn(|cx| self.poll_stop(cx)).await.ok();
590
591        // Assume idle state and take out the inner value.
592        let mut t = match &mut self.state {
593            State::Idle(t) => t.take().expect("inner value was taken out"),
594            State::WithMut(..)
595            | State::Streaming(..)
596            | State::Reading(..)
597            | State::Writing(..)
598            | State::Seeking(..) => {
599                unreachable!("when stopped, the state machine must be in idle state");
600            }
601        };
602
603        let (sender, receiver) = bounded(1);
604        let task = Executor::spawn(async move {
605            sender.try_send(op(&mut t)).ok();
606            t
607        });
608        self.state = State::WithMut(task);
609
610        receiver
611            .recv()
612            .await
613            .expect("`Unblock::with_mut()` operation has panicked")
614    }
615
616    /// Extracts the inner blocking I/O handle.
617    ///
618    /// This is an async method because the I/O handle might be on the thread pool and needs to
619    /// be moved onto the current thread before we can extract it.
620    ///
621    /// # Examples
622    ///
623    /// ```no_run
624    /// use blocking::{unblock, Unblock};
625    /// use futures_lite::prelude::*;
626    /// use std::fs::File;
627    ///
628    /// # futures_lite::future::block_on(async {
629    /// let file = unblock(|| File::create("file.txt")).await?;
630    /// let file = Unblock::new(file);
631    ///
632    /// let file = file.into_inner().await;
633    /// # std::io::Result::Ok(()) });
634    /// ```
635    pub async fn into_inner(self) -> T {
636        // There's a bug in rustdoc causing it to render `mut self` as `__arg0: Self`, so we just
637        // bind `self` to a local mutable variable.
638        let mut this = self;
639
640        // Wait for the running task to stop and ignore I/O errors if there are any.
641        future::poll_fn(|cx| this.poll_stop(cx)).await.ok();
642
643        // Assume idle state and extract the inner value.
644        match &mut this.state {
645            State::Idle(t) => *t.take().expect("inner value was taken out"),
646            State::WithMut(..)
647            | State::Streaming(..)
648            | State::Reading(..)
649            | State::Writing(..)
650            | State::Seeking(..) => {
651                unreachable!("when stopped, the state machine must be in idle state");
652            }
653        }
654    }
655
656    /// Waits for the running task to stop.
657    ///
658    /// On success, the state machine is moved into the idle state.
659    fn poll_stop(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
660        loop {
661            match &mut self.state {
662                State::Idle(_) => return Poll::Ready(Ok(())),
663
664                State::WithMut(task) => {
665                    // Poll the task to wait for it to finish.
666                    let io = ready!(Pin::new(task).poll(cx));
667                    self.state = State::Idle(Some(io));
668                }
669
670                State::Streaming(any, task) => {
671                    // Drop the receiver to close the channel. This stops the `send()` operation in
672                    // the task, after which the task returns the iterator back.
673                    any.take();
674
675                    // Poll the task to retrieve the iterator.
676                    let iter = ready!(Pin::new(task).poll(cx));
677                    self.state = State::Idle(Some(iter));
678                }
679
680                State::Reading(reader, task) => {
681                    // Drop the reader to close the pipe. This stops copying inside the task, after
682                    // which the task returns the I/O handle back.
683                    reader.take();
684
685                    // Poll the task to retrieve the I/O handle.
686                    let (res, io) = ready!(Pin::new(task).poll(cx));
687                    // Make sure to move into the idle state before reporting errors.
688                    self.state = State::Idle(Some(io));
689                    res?;
690                }
691
692                State::Writing(writer, task) => {
693                    // Drop the writer to close the pipe. This stops copying inside the task, after
694                    // which the task flushes the I/O handle and
695                    writer.take();
696
697                    // Poll the task to retrieve the I/O handle.
698                    let (res, io) = ready!(Pin::new(task).poll(cx));
699                    // Make sure to move into the idle state before reporting errors.
700                    self.state = State::Idle(Some(io));
701                    res?;
702                }
703
704                State::Seeking(task) => {
705                    // Poll the task to wait for it to finish.
706                    let (_, res, io) = ready!(Pin::new(task).poll(cx));
707                    // Make sure to move into the idle state before reporting errors.
708                    self.state = State::Idle(Some(io));
709                    res?;
710                }
711            }
712        }
713    }
714}
715
716impl<T: fmt::Debug> fmt::Debug for Unblock<T> {
717    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
718        struct Closed;
719        impl fmt::Debug for Closed {
720            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
721                f.write_str("<closed>")
722            }
723        }
724
725        struct Blocked;
726        impl fmt::Debug for Blocked {
727            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
728                f.write_str("<blocked>")
729            }
730        }
731
732        match &self.state {
733            State::Idle(None) => f.debug_struct("Unblock").field("io", &Closed).finish(),
734            State::Idle(Some(io)) => {
735                let io: &T = io;
736                f.debug_struct("Unblock").field("io", io).finish()
737            }
738            State::WithMut(..)
739            | State::Streaming(..)
740            | State::Reading(..)
741            | State::Writing(..)
742            | State::Seeking(..) => f.debug_struct("Unblock").field("io", &Blocked).finish(),
743        }
744    }
745}
746
747/// Current state of a blocking task.
748enum State<T> {
749    /// There is no blocking task.
750    ///
751    /// The inner value is readily available, unless it has already been extracted. The value is
752    /// extracted out by [`Unblock::into_inner()`], [`AsyncWrite::poll_close()`], or by awaiting
753    /// [`Unblock`].
754    Idle(Option<Box<T>>),
755
756    /// A [`Unblock::with_mut()`] closure was spawned and is still running.
757    WithMut(Task<Box<T>>),
758
759    /// The inner value is an [`Iterator`] currently iterating in a task.
760    ///
761    /// The `dyn Any` value here is a `Pin<Box<Receiver<<T as Iterator>::Item>>>`.
762    Streaming(Option<Box<dyn Any + Send + Sync>>, Task<Box<T>>),
763
764    /// The inner value is a [`Read`] currently reading in a task.
765    Reading(Option<Reader>, Task<(io::Result<()>, Box<T>)>),
766
767    /// The inner value is a [`Write`] currently writing in a task.
768    Writing(Option<Writer>, Task<(io::Result<()>, Box<T>)>),
769
770    /// The inner value is a [`Seek`] currently seeking in a task.
771    Seeking(Task<(SeekFrom, io::Result<u64>, Box<T>)>),
772}
773
774impl<T: Iterator + Send + 'static> Stream for Unblock<T>
775where
776    T::Item: Send + 'static,
777{
778    type Item = T::Item;
779
780    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T::Item>> {
781        loop {
782            match &mut self.state {
783                // If not in idle or active streaming state, stop the running task.
784                State::WithMut(..)
785                | State::Streaming(None, _)
786                | State::Reading(..)
787                | State::Writing(..)
788                | State::Seeking(..) => {
789                    // Wait for the running task to stop.
790                    ready!(self.poll_stop(cx)).ok();
791                }
792
793                // If idle, start a streaming task.
794                State::Idle(iter) => {
795                    // Take the iterator out to run it on a blocking task.
796                    let mut iter = iter.take().expect("inner iterator was taken out");
797
798                    // This channel capacity seems to work well in practice. If it's too low, there
799                    // will be too much synchronization between tasks. If too high, memory
800                    // consumption increases.
801                    let (sender, receiver) = bounded(self.cap.unwrap_or(8 * 1024)); // 8192 items
802
803                    // Spawn a blocking task that runs the iterator and returns it when done.
804                    let task = Executor::spawn(async move {
805                        for item in &mut iter {
806                            if sender.send(item).await.is_err() {
807                                break;
808                            }
809                        }
810                        iter
811                    });
812
813                    // Move into the busy state and poll again.
814                    self.state = State::Streaming(Some(Box::new(Box::pin(receiver))), task);
815                }
816
817                // If streaming, receive an item.
818                State::Streaming(Some(any), task) => {
819                    let receiver = any.downcast_mut::<Pin<Box<Receiver<T::Item>>>>().unwrap();
820
821                    // Poll the channel.
822                    let opt = ready!(receiver.as_mut().poll_next(cx));
823
824                    // If the channel is closed, retrieve the iterator back from the blocking task.
825                    // This is not really a required step, but it's cleaner to drop the iterator on
826                    // the same thread that created it.
827                    if opt.is_none() {
828                        // Poll the task to retrieve the iterator.
829                        let iter = ready!(Pin::new(task).poll(cx));
830                        self.state = State::Idle(Some(iter));
831                    }
832
833                    return Poll::Ready(opt);
834                }
835            }
836        }
837    }
838}
839
840impl<T: Read + Send + 'static> AsyncRead for Unblock<T> {
841    fn poll_read(
842        mut self: Pin<&mut Self>,
843        cx: &mut Context<'_>,
844        buf: &mut [u8],
845    ) -> Poll<io::Result<usize>> {
846        loop {
847            match &mut self.state {
848                // If not in idle or active reading state, stop the running task.
849                State::WithMut(..)
850                | State::Reading(None, _)
851                | State::Streaming(..)
852                | State::Writing(..)
853                | State::Seeking(..) => {
854                    // Wait for the running task to stop.
855                    ready!(self.poll_stop(cx))?;
856                }
857
858                // If idle, start a reading task.
859                State::Idle(io) => {
860                    // Take the I/O handle out to read it on a blocking task.
861                    let mut io = io.take().expect("inner value was taken out");
862
863                    // This pipe capacity seems to work well in practice. If it's too low, there
864                    // will be too much synchronization between tasks. If too high, memory
865                    // consumption increases.
866                    let (reader, mut writer) = pipe(self.cap.unwrap_or(8 * 1024 * 1024)); // 8 MB
867
868                    // Spawn a blocking task that reads and returns the I/O handle when done.
869                    let task = Executor::spawn(async move {
870                        // Copy bytes from the I/O handle into the pipe until the pipe is closed or
871                        // an error occurs.
872                        loop {
873                            match future::poll_fn(|cx| writer.poll_fill(cx, &mut io)).await {
874                                Ok(0) => return (Ok(()), io),
875                                Ok(_) => {}
876                                Err(err) => return (Err(err), io),
877                            }
878                        }
879                    });
880
881                    // Move into the busy state and poll again.
882                    self.state = State::Reading(Some(reader), task);
883                }
884
885                // If reading, read bytes from the pipe.
886                State::Reading(Some(reader), task) => {
887                    // Poll the pipe.
888                    let n = ready!(reader.poll_drain(cx, buf))?;
889
890                    // If the pipe is closed, retrieve the I/O handle back from the blocking task.
891                    // This is not really a required step, but it's cleaner to drop the handle on
892                    // the same thread that created it.
893                    if n == 0 {
894                        // Poll the task to retrieve the I/O handle.
895                        let (res, io) = ready!(Pin::new(task).poll(cx));
896                        // Make sure to move into the idle state before reporting errors.
897                        self.state = State::Idle(Some(io));
898                        res?;
899                    }
900
901                    return Poll::Ready(Ok(n));
902                }
903            }
904        }
905    }
906}
907
908impl<T: Write + Send + 'static> AsyncWrite for Unblock<T> {
909    fn poll_write(
910        mut self: Pin<&mut Self>,
911        cx: &mut Context<'_>,
912        buf: &[u8],
913    ) -> Poll<io::Result<usize>> {
914        loop {
915            match &mut self.state {
916                // If not in idle or active writing state, stop the running task.
917                State::WithMut(..)
918                | State::Writing(None, _)
919                | State::Streaming(..)
920                | State::Reading(..)
921                | State::Seeking(..) => {
922                    // Wait for the running task to stop.
923                    ready!(self.poll_stop(cx))?;
924                }
925
926                // If idle, start the writing task.
927                State::Idle(io) => {
928                    // Take the I/O handle out to write on a blocking task.
929                    let mut io = io.take().expect("inner value was taken out");
930
931                    // This pipe capacity seems to work well in practice. If it's too low, there will
932                    // be too much synchronization between tasks. If too high, memory consumption
933                    // increases.
934                    let (mut reader, writer) = pipe(self.cap.unwrap_or(8 * 1024 * 1024)); // 8 MB
935
936                    // Spawn a blocking task that writes and returns the I/O handle when done.
937                    let task = Executor::spawn(async move {
938                        // Copy bytes from the pipe into the I/O handle until the pipe is closed or an
939                        // error occurs. Flush the I/O handle at the end.
940                        loop {
941                            match future::poll_fn(|cx| reader.poll_drain(cx, &mut io)).await {
942                                Ok(0) => return (io.flush(), io),
943                                Ok(_) => {}
944                                Err(err) => {
945                                    io.flush().ok();
946                                    return (Err(err), io);
947                                }
948                            }
949                        }
950                    });
951
952                    // Move into the busy state and poll again.
953                    self.state = State::Writing(Some(writer), task);
954                }
955
956                // If writing, write more bytes into the pipe.
957                State::Writing(Some(writer), _) => return writer.poll_fill(cx, buf),
958            }
959        }
960    }
961
962    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
963        loop {
964            match &mut self.state {
965                // If not in idle state, stop the running task.
966                State::WithMut(..)
967                | State::Streaming(..)
968                | State::Writing(..)
969                | State::Reading(..)
970                | State::Seeking(..) => {
971                    // Wait for the running task to stop.
972                    ready!(self.poll_stop(cx))?;
973                }
974
975                // Idle implies flushed.
976                State::Idle(_) => return Poll::Ready(Ok(())),
977            }
978        }
979    }
980
981    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
982        // First, make sure the I/O handle is flushed.
983        ready!(Pin::new(&mut self).poll_flush(cx))?;
984
985        // Then move into the idle state with no I/O handle, thus dropping it.
986        self.state = State::Idle(None);
987        Poll::Ready(Ok(()))
988    }
989}
990
991impl<T: Seek + Send + 'static> AsyncSeek for Unblock<T> {
992    fn poll_seek(
993        mut self: Pin<&mut Self>,
994        cx: &mut Context<'_>,
995        pos: SeekFrom,
996    ) -> Poll<io::Result<u64>> {
997        loop {
998            match &mut self.state {
999                // If not in idle state, stop the running task.
1000                State::WithMut(..)
1001                | State::Streaming(..)
1002                | State::Reading(..)
1003                | State::Writing(..) => {
1004                    // Wait for the running task to stop.
1005                    ready!(self.poll_stop(cx))?;
1006                }
1007
1008                State::Idle(io) => {
1009                    // Take the I/O handle out to seek on a blocking task.
1010                    let mut io = io.take().expect("inner value was taken out");
1011
1012                    let task = Executor::spawn(async move {
1013                        let res = io.seek(pos);
1014                        (pos, res, io)
1015                    });
1016                    self.state = State::Seeking(task);
1017                }
1018
1019                State::Seeking(task) => {
1020                    // Poll the task to wait for it to finish.
1021                    let (original_pos, res, io) = ready!(Pin::new(task).poll(cx));
1022                    // Make sure to move into the idle state before reporting errors.
1023                    self.state = State::Idle(Some(io));
1024                    let current = res?;
1025
1026                    // If the `pos` argument matches the original one, return the result.
1027                    if original_pos == pos {
1028                        return Poll::Ready(Ok(current));
1029                    }
1030                }
1031            }
1032        }
1033    }
1034}
1035
1036#[cfg(all(test, not(target_family = "wasm")))]
1037mod tests {
1038    use super::*;
1039
1040    #[test]
1041    fn test_max_threads() {
1042        // properly set env var
1043        env::set_var(MAX_THREADS_ENV, "100");
1044        assert_eq!(100, Executor::max_threads().get());
1045
1046        // passed value below minimum, so we set it to minimum
1047        env::set_var(MAX_THREADS_ENV, "0");
1048        assert_eq!(1, Executor::max_threads().get());
1049
1050        // passed value above maximum, so we set to allowed maximum
1051        env::set_var(MAX_THREADS_ENV, "50000");
1052        assert_eq!(10000, Executor::max_threads().get());
1053
1054        // no env var, use default
1055        env::set_var(MAX_THREADS_ENV, "");
1056        assert_eq!(500, Executor::max_threads().get());
1057
1058        // not a number, use default
1059        env::set_var(MAX_THREADS_ENV, "NOTINT");
1060        assert_eq!(500, Executor::max_threads().get());
1061    }
1062}