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::task::Poll;
17
18/// An single-threaded task pool backed by IO backend `T`.
19#[derive(Debug)]
20pub struct IoPool<T> {
21    driver: IoDriver<T>,
22    tasks: TaskQueue,
23}
24
25/// A driver to spawn tasks and IO objects on [`IoPool`].
26#[derive(Debug)]
27pub struct IoDriver<T> {
28    pub(crate) inner: Arc<T>,
29    scheduler: Arc<Scheduler>,
30}
31
32impl<T> Clone for IoDriver<T> {
33    fn clone(&self) -> Self {
34        Self {
35            inner: self.inner.clone(),
36            scheduler: self.scheduler.clone(),
37        }
38    }
39}
40
41/// Trait implemented by IO backends.
42pub trait IoBackend: Send + Sync {
43    /// The name of the backend.
44    fn name() -> &'static str;
45    /// Run the
46    fn run<Fut: Future>(self: &Arc<Self>, fut: Fut) -> Fut::Output;
47}
48
49impl<T: IoBackend + Default> IoPool<T> {
50    /// Creates a new task pool.
51    pub fn new() -> Self {
52        Self::named(T::name().to_owned())
53    }
54
55    /// Creates a new task pool with the given name, used to identify the
56    /// executor in traces.
57    pub fn named(name: impl Into<Arc<str>>) -> Self {
58        let (tasks, scheduler) = task_queue(name);
59        Self {
60            driver: IoDriver {
61                inner: Arc::new(T::default()),
62                scheduler: Arc::new(scheduler),
63            },
64            tasks,
65        }
66    }
67
68    /// Creates and runs a task pool, seeding it with an initial future
69    /// `f(driver)`, until all tasks have completed.
70    pub fn run_with<F, R>(f: F) -> R
71    where
72        F: AsyncFnOnce(IoDriver<T>) -> R,
73    {
74        let mut pool = Self::named(std::thread::current().name().unwrap_or_else(|| T::name()));
75        let fut = f(pool.driver.clone());
76        drop(pool.driver.scheduler);
77        pool.driver
78            .inner
79            .run(async { futures::future::join(fut, pool.tasks.run()).await.0 })
80    }
81
82    /// Creates a new pool and runs it on a newly spawned thread with the given
83    /// name. Returns the thread handle and the pool's driver.
84    pub fn spawn_on_thread(name: impl Into<String>) -> (std::thread::JoinHandle<()>, IoDriver<T>)
85    where
86        T: 'static,
87    {
88        let pool = Self::new();
89        let driver = pool.driver.clone();
90        let thread = std::thread::Builder::new()
91            .name(name.into())
92            .spawn(move || pool.run())
93            .unwrap();
94        (thread, driver)
95    }
96}
97
98impl<T: IoBackend> IoPool<T> {
99    /// Returns the IO driver.
100    pub fn driver(&self) -> IoDriver<T> {
101        self.driver.clone()
102    }
103
104    /// Runs `f` and the task pool until `f` completes.
105    pub fn run_until<Fut: Future>(&mut self, f: Fut) -> Fut::Output {
106        let mut tasks = pin!(self.tasks.run());
107        let mut f = pin!(f);
108        self.driver.inner.run(poll_fn(|cx| {
109            if let Poll::Ready(r) = f.as_mut().poll(cx) {
110                Poll::Ready(r)
111            } else {
112                assert!(tasks.as_mut().poll(cx).is_pending());
113                Poll::Pending
114            }
115        }))
116    }
117
118    /// Runs the task pool until all tasks are completed.
119    pub fn run(mut self) {
120        // Update the executor name with the current thread's name.
121        if let Some(name) = std::thread::current().name() {
122            self.driver.scheduler.set_name(name);
123        }
124        drop(self.driver.scheduler);
125        self.driver.inner.run(self.tasks.run())
126    }
127}
128
129impl<T: IoBackend> Spawn for IoDriver<T> {
130    fn scheduler(&self, _metadata: &TaskMetadata) -> Arc<dyn Schedule> {
131        self.scheduler.clone()
132    }
133}
pal_async/io_pool.rs

pal_async/
io_pool.rs
