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}