pal_async/
io_pool.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Single-threaded task pools backed by platform-specific IO backends.
5
6use crate::task::Schedule;
7use crate::task::Scheduler;
8use crate::task::Spawn;
9use crate::task::TaskMetadata;
10use crate::task::TaskQueue;
11use crate::task::task_queue;
12use std::future::Future;
13use std::future::poll_fn;
14use std::pin::pin;
15use std::sync::Arc;
16use std::sync::Weak;
17use std::task::Poll;
18
19/// A single-threaded task pool backed by IO backend `T`.
20#[derive(Debug)]
21pub struct IoPool<T> {
22    driver: IoDriver<T>,
23    tasks: TaskQueue,
24}
25
26/// A driver to spawn tasks and IO objects on [`IoPool`].
27#[derive(Debug)]
28pub struct IoDriver<T> {
29    pub(crate) inner: Arc<T>,
30    scheduler: Arc<Scheduler>,
31}
32
33impl<T> Clone for IoDriver<T> {
34    fn clone(&self) -> Self {
35        Self {
36            inner: self.inner.clone(),
37            scheduler: self.scheduler.clone(),
38        }
39    }
40}
41
42impl<T> IoDriver<T> {
43    /// Returns a weak reference to this driver.
44    ///
45    /// A [`WeakIoDriver`] does not keep the backing [`IoPool`] alive: [`IoPool::run`]
46    /// returns once the last strong `IoDriver` is dropped. Use this to publish a
47    /// reference to a pool's own driver (e.g. into thread-local storage on the
48    /// pool's thread) without preventing the pool from ever shutting down. A
49    /// strong self-reference held for the lifetime of `run` would instead keep the
50    /// pool's task queue open forever (the pool never sees its last driver drop),
51    /// leaking the thread and its IO backend.
52    pub fn downgrade(&self) -> WeakIoDriver<T> {
53        WeakIoDriver {
54            inner: Arc::downgrade(&self.inner),
55            scheduler: Arc::downgrade(&self.scheduler),
56        }
57    }
58}
59
60/// A weak reference to an [`IoDriver`], obtained via [`IoDriver::downgrade`].
61#[derive(Debug)]
62pub struct WeakIoDriver<T> {
63    inner: Weak<T>,
64    scheduler: Weak<Scheduler>,
65}
66
67impl<T> Clone for WeakIoDriver<T> {
68    fn clone(&self) -> Self {
69        Self {
70            inner: self.inner.clone(),
71            scheduler: self.scheduler.clone(),
72        }
73    }
74}
75
76impl<T> WeakIoDriver<T> {
77    /// Attempts to upgrade to a strong [`IoDriver`].
78    ///
79    /// Returns `None` once the last strong `IoDriver` has been dropped and no
80    /// task on the pool is still alive.
81    pub fn upgrade(&self) -> Option<IoDriver<T>> {
82        Some(IoDriver {
83            inner: self.inner.upgrade()?,
84            scheduler: self.scheduler.upgrade()?,
85        })
86    }
87}
88
89/// Trait implemented by IO backends.
90pub trait IoBackend: Send + Sync {
91    /// The name of the backend.
92    fn name() -> &'static str;
93    /// Run the
94    fn run<Fut: Future>(self: &Arc<Self>, fut: Fut) -> Fut::Output;
95}
96
97impl<T: IoBackend + Default> IoPool<T> {
98    /// Creates a new task pool.
99    pub fn new() -> Self {
100        Self::named(T::name().to_owned())
101    }
102
103    /// Creates a new task pool with the given name, used to identify the
104    /// executor in traces.
105    pub fn named(name: impl Into<Arc<str>>) -> Self {
106        let (tasks, scheduler) = task_queue(name);
107        Self {
108            driver: IoDriver {
109                inner: Arc::new(T::default()),
110                scheduler: Arc::new(scheduler),
111            },
112            tasks,
113        }
114    }
115
116    /// Creates and runs a task pool, seeding it with an initial future
117    /// `f(driver)`, until all tasks have completed.
118    pub fn run_with<F, R>(f: F) -> R
119    where
120        F: AsyncFnOnce(IoDriver<T>) -> R,
121    {
122        let mut pool = Self::named(std::thread::current().name().unwrap_or_else(|| T::name()));
123        let fut = f(pool.driver.clone());
124        drop(pool.driver.scheduler);
125        pool.driver
126            .inner
127            .run(async { futures::future::join(fut, pool.tasks.run()).await.0 })
128    }
129
130    /// Creates a new pool and runs it on a newly spawned thread with the given
131    /// name. Returns the thread handle and the pool's driver.
132    pub fn spawn_on_thread(name: impl Into<String>) -> (std::thread::JoinHandle<()>, IoDriver<T>)
133    where
134        T: 'static,
135    {
136        let pool = Self::new();
137        let driver = pool.driver.clone();
138        let thread = std::thread::Builder::new()
139            .name(name.into())
140            .spawn(move || pool.run())
141            .unwrap();
142        (thread, driver)
143    }
144}
145
146impl<T: IoBackend> IoPool<T> {
147    /// Returns the IO driver.
148    pub fn driver(&self) -> IoDriver<T> {
149        self.driver.clone()
150    }
151
152    /// Runs `f` and the task pool until `f` completes.
153    pub fn run_until<Fut: Future>(&mut self, f: Fut) -> Fut::Output {
154        let mut tasks = pin!(self.tasks.run());
155        let mut f = pin!(f);
156        self.driver.inner.run(poll_fn(|cx| {
157            if let Poll::Ready(r) = f.as_mut().poll(cx) {
158                Poll::Ready(r)
159            } else {
160                assert!(tasks.as_mut().poll(cx).is_pending());
161                Poll::Pending
162            }
163        }))
164    }
165
166    /// Runs the task pool until all tasks are completed.
167    pub fn run(mut self) {
168        // Update the executor name with the current thread's name.
169        if let Some(name) = std::thread::current().name() {
170            self.driver.scheduler.set_name(name);
171        }
172        drop(self.driver.scheduler);
173        self.driver.inner.run(self.tasks.run())
174    }
175}
176
177impl<T: IoBackend> Spawn for IoDriver<T> {
178    fn scheduler(&self, _metadata: &TaskMetadata) -> Arc<dyn Schedule> {
179        self.scheduler.clone()
180    }
181}
pal_async/io_pool.rs

pal_async/
io_pool.rs
