vmcore/
vm_task.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Infrastructure for spawning tasks and issuing async IO related to VM
5//! activity.
6
7// UNSAFETY: Needed to implement the unsafe new_dyn_overlapped_file method.
8#![cfg_attr(windows, expect(unsafe_code))]
9
10use inspect::Inspect;
11use pal_async::driver::Driver;
12use pal_async::task::Spawn;
13use pal_async::task::TaskMetadata;
14use std::fmt::Debug;
15use std::pin::Pin;
16use std::sync::Arc;
17
18/// A source for [`VmTaskDriver`]s.
19///
20/// This is used to create device-specific drivers that implement [`Driver`] and
21/// [`Spawn`]. These drivers' behavior can be customized based on the needs of
22/// the device.
23///
24/// The backend for these drivers is customizable for different environments.
25#[derive(Clone)]
26pub struct VmTaskDriverSource {
27    backend: Arc<dyn DynVmBackend>,
28}
29
30impl VmTaskDriverSource {
31    /// Returns a new task driver source backed by `backend`.
32    pub fn new(backend: impl 'static + BuildVmTaskDriver) -> Self {
33        Self {
34            backend: Arc::new(backend),
35        }
36    }
37
38    /// Returns a VM task driver with default parameters.
39    ///
40    /// Use this when you don't care where your task runs.
41    pub fn simple(&self) -> VmTaskDriver {
42        // Don't provide a name, since backends won't do anything with it for
43        // default settings.
44        self.builder().build("")
45    }
46
47    /// Returns a builder for a custom VM task driver.
48    pub fn builder(&self) -> VmTaskDriverBuilder<'_> {
49        VmTaskDriverBuilder {
50            backend: self.backend.as_ref(),
51            run_on_target: false,
52            target_vp: None,
53        }
54    }
55}
56
57/// Trait implemented by backends for [`VmTaskDriverSource`].
58pub trait BuildVmTaskDriver: Send + Sync {
59    /// The associated driver type.
60    type Driver: TargetedDriver;
61
62    /// Builds a new driver that can drive IO and spawn tasks.
63    fn build(&self, name: String, target_vp: Option<u32>, run_on_target: bool) -> Self::Driver;
64}
65
66/// Trait implemented by drivers built with [`BuildVmTaskDriver`].
67pub trait TargetedDriver: 'static + Send + Sync + Inspect {
68    /// Returns the implementation to use for spawning tasks.
69    fn spawner(&self) -> &dyn Spawn;
70    /// Returns the implementation to use for driving IO.
71    fn driver(&self) -> &dyn Driver;
72    /// Retargets the driver to the specified virtual processor.
73    fn retarget_vp(&self, target_vp: u32);
74    /// Returns whether a driver's target VP is ready for tasks and IO.
75    ///
76    /// A driver must be operable even if this is false, but the tasks and IO
77    /// may run on a different target VP.
78    fn is_target_vp_ready(&self) -> bool {
79        true
80    }
81    /// Waits for this driver's target VP to be ready for tasks and IO.
82    fn wait_target_vp_ready(&self) -> impl std::future::Future<Output = ()> + Send {
83        std::future::ready(())
84    }
85}
86
87trait DynTargetedDriver: 'static + Send + Sync + Inspect {
88    fn spawner(&self) -> &dyn Spawn;
89    fn driver(&self) -> &dyn Driver;
90    fn retarget_vp(&self, target_vp: u32);
91    fn is_ready(&self) -> bool;
92    fn wait_ready(&self) -> Pin<Box<dyn '_ + std::future::Future<Output = ()> + Send>>;
93}
94
95impl<T: TargetedDriver> DynTargetedDriver for T {
96    fn spawner(&self) -> &dyn Spawn {
97        self.spawner()
98    }
99
100    fn driver(&self) -> &dyn Driver {
101        self.driver()
102    }
103
104    fn retarget_vp(&self, target_vp: u32) {
105        self.retarget_vp(target_vp)
106    }
107
108    fn is_ready(&self) -> bool {
109        self.is_target_vp_ready()
110    }
111
112    fn wait_ready(&self) -> Pin<Box<dyn '_ + std::future::Future<Output = ()> + Send>> {
113        Box::pin(self.wait_target_vp_ready())
114    }
115}
116
117trait DynVmBackend: Send + Sync {
118    fn build(
119        &self,
120        name: String,
121        target_vp: Option<u32>,
122        run_on_target: bool,
123    ) -> Arc<dyn DynTargetedDriver>;
124}
125
126impl<T: BuildVmTaskDriver> DynVmBackend for T {
127    fn build(
128        &self,
129        name: String,
130        target_vp: Option<u32>,
131        run_on_target: bool,
132    ) -> Arc<dyn DynTargetedDriver> {
133        Arc::new(self.build(name, target_vp, run_on_target))
134    }
135}
136
137/// A builder returned by [`VmTaskDriverSource::builder`].
138pub struct VmTaskDriverBuilder<'a> {
139    backend: &'a dyn DynVmBackend,
140    run_on_target: bool,
141    target_vp: Option<u32>,
142}
143
144impl VmTaskDriverBuilder<'_> {
145    /// A hint to the backend specifies whether the driver should spawned tasks
146    /// that always on a thread handling the target VP.
147    ///
148    /// If `false` (the default), then when spawned tasks are awoken, they may
149    /// run on any executor (such as the current one). If `true`, the backend
150    /// will run them on the same thread that would drive async IO.
151    ///
152    /// Some devices will want to override the default to reduce jitter or
153    /// ensure that IO is issued from the correct processor.
154    pub fn run_on_target(&mut self, run_on_target: bool) -> &mut Self {
155        self.run_on_target = run_on_target;
156        self
157    }
158
159    /// A hint to the backend specifying the guest VP associated with spawned
160    /// tasks and IO.
161    ///
162    /// Backends can use this to ensure that spawned tasks and async IO will run
163    /// near or on the target VP.
164    pub fn target_vp(&mut self, target_vp: u32) -> &mut Self {
165        self.target_vp = Some(target_vp);
166        self
167    }
168
169    /// Builds a VM task driver.
170    ///
171    /// `name` is used by some backends to identify a spawned thread. It is
172    /// ignored by other backends.
173    pub fn build(&self, name: impl Into<String>) -> VmTaskDriver {
174        VmTaskDriver {
175            inner: self
176                .backend
177                .build(name.into(), self.target_vp, self.run_on_target),
178        }
179    }
180}
181
182/// A driver returned by [`VmTaskDriverSource`].
183///
184/// This can be used to spawn tasks (via [`Spawn`]) and issue async IO (via [`Driver`]).
185#[derive(Clone, Inspect)]
186pub struct VmTaskDriver {
187    #[inspect(flatten)]
188    inner: Arc<dyn DynTargetedDriver>,
189}
190
191impl VmTaskDriver {
192    /// Updates the target VP for the task.
193    pub fn retarget_vp(&self, target_vp: u32) {
194        self.inner.retarget_vp(target_vp)
195    }
196
197    /// Returns whether the target VP is ready for tasks and IO.
198    ///
199    /// A driver must be operable even if this is false, but the tasks and IO
200    /// may run on a different target VP.
201    pub fn is_target_vp_ready(&self) -> bool {
202        self.inner.is_ready()
203    }
204
205    /// Waits for the target VP to be ready for tasks and IO.
206    pub async fn wait_target_vp_ready(&self) {
207        self.inner.wait_ready().await
208    }
209}
210
211impl Driver for VmTaskDriver {
212    fn new_dyn_timer(&self) -> pal_async::driver::PollImpl<dyn pal_async::timer::PollTimer> {
213        self.inner.driver().new_dyn_timer()
214    }
215
216    #[cfg(unix)]
217    fn new_dyn_fd_ready(
218        &self,
219        fd: std::os::fd::RawFd,
220    ) -> std::io::Result<pal_async::driver::PollImpl<dyn pal_async::fd::PollFdReady>> {
221        self.inner.driver().new_dyn_fd_ready(fd)
222    }
223
224    #[cfg(unix)]
225    fn new_dyn_socket_ready(
226        &self,
227        socket: std::os::fd::RawFd,
228    ) -> std::io::Result<pal_async::driver::PollImpl<dyn pal_async::socket::PollSocketReady>> {
229        self.inner.driver().new_dyn_socket_ready(socket)
230    }
231
232    #[cfg(windows)]
233    fn new_dyn_socket_ready(
234        &self,
235        socket: std::os::windows::io::RawSocket,
236    ) -> std::io::Result<pal_async::driver::PollImpl<dyn pal_async::socket::PollSocketReady>> {
237        self.inner.driver().new_dyn_socket_ready(socket)
238    }
239
240    #[cfg(unix)]
241    fn new_dyn_wait(
242        &self,
243        fd: std::os::fd::RawFd,
244        read_size: usize,
245    ) -> std::io::Result<pal_async::driver::PollImpl<dyn pal_async::wait::PollWait>> {
246        self.inner.driver().new_dyn_wait(fd, read_size)
247    }
248
249    #[cfg(windows)]
250    fn new_dyn_wait(
251        &self,
252        handle: std::os::windows::io::RawHandle,
253    ) -> std::io::Result<pal_async::driver::PollImpl<dyn pal_async::wait::PollWait>> {
254        self.inner.driver().new_dyn_wait(handle)
255    }
256
257    #[cfg(windows)]
258    unsafe fn new_dyn_overlapped_file(
259        &self,
260        handle: std::os::windows::io::RawHandle,
261    ) -> std::io::Result<
262        pal_async::driver::PollImpl<dyn pal_async::windows::overlapped::IoOverlapped>,
263    > {
264        // SAFETY: passthru from caller
265        unsafe { self.inner.driver().new_dyn_overlapped_file(handle) }
266    }
267}
268
269impl Spawn for VmTaskDriver {
270    fn scheduler(&self, metadata: &TaskMetadata) -> Arc<dyn pal_async::task::Schedule> {
271        self.inner.spawner().scheduler(metadata)
272    }
273}
274
275/// A backend that spawns all tasks and IO on a single driver.
276#[derive(Debug)]
277pub struct SingleDriverBackend<T>(T);
278
279impl<T: Driver + Spawn + Clone> SingleDriverBackend<T> {
280    /// Returns a new driver backend that spawns all tasks and IO on `driver`,
281    /// regardless of policy.
282    pub fn new(driver: T) -> Self {
283        Self(driver)
284    }
285}
286
287/// The driver for [`SingleDriverBackend`].
288#[derive(Debug)]
289pub struct SingleDriver<T>(T);
290
291impl<T> Inspect for SingleDriver<T> {
292    fn inspect(&self, req: inspect::Request<'_>) {
293        req.ignore();
294    }
295}
296
297impl<T: Driver + Spawn + Clone> BuildVmTaskDriver for SingleDriverBackend<T> {
298    type Driver = SingleDriver<T>;
299
300    fn build(&self, _name: String, _target_vp: Option<u32>, _run_on_target: bool) -> Self::Driver {
301        SingleDriver(self.0.clone())
302    }
303}
304
305impl<T: Driver + Spawn> TargetedDriver for SingleDriver<T> {
306    fn spawner(&self) -> &dyn Spawn {
307        &self.0
308    }
309
310    fn driver(&self) -> &dyn Driver {
311        &self.0
312    }
313
314    fn retarget_vp(&self, _target_vp: u32) {}
315}
316
317pub mod thread {
318    //! Provides a thread-based task VM task driver backend
319    //! [`ThreadDriverBackend`].
320
321    use super::BuildVmTaskDriver;
322    use super::TargetedDriver;
323    use inspect::Inspect;
324    use pal_async::DefaultDriver;
325    use pal_async::DefaultPool;
326    use pal_async::driver::Driver;
327    use pal_async::task::Spawn;
328
329    /// A backend for [`VmTaskDriverSource`](super::VmTaskDriverSource) based on
330    /// individual threads.
331    ///
332    /// If no target VP is specified, this backend will spawn tasks and IO a
333    /// default single-threaded IO driver. If a target VP is specified, the
334    /// backend will spawn a separate thread and spawn tasks and IOs there.
335    #[derive(Debug)]
336    pub struct ThreadDriverBackend {
337        default_driver: DefaultDriver,
338    }
339
340    impl ThreadDriverBackend {
341        /// Returns a new backend, using `default_driver` to back task drivers
342        /// that did not specify a target VP.
343        pub fn new(default_driver: DefaultDriver) -> Self {
344            Self { default_driver }
345        }
346    }
347
348    impl BuildVmTaskDriver for ThreadDriverBackend {
349        type Driver = ThreadDriver;
350
351        fn build(
352            &self,
353            name: String,
354            target_vp: Option<u32>,
355            _run_on_target: bool,
356        ) -> Self::Driver {
357            // Build a standalone thread for this device if a target VP was specified.
358            if target_vp.is_some() {
359                let (_, driver) = DefaultPool::spawn_on_thread(name);
360                ThreadDriver {
361                    inner: driver,
362                    has_dedicated_thread: true,
363                }
364            } else {
365                ThreadDriver {
366                    inner: self.default_driver.clone(),
367                    has_dedicated_thread: false,
368                }
369            }
370        }
371    }
372
373    /// The driver for [`ThreadDriverBackend`].
374    #[derive(Debug, Inspect)]
375    pub struct ThreadDriver {
376        #[inspect(skip)]
377        inner: DefaultDriver,
378        has_dedicated_thread: bool,
379    }
380
381    impl TargetedDriver for ThreadDriver {
382        fn spawner(&self) -> &dyn Spawn {
383            &self.inner
384        }
385
386        fn driver(&self) -> &dyn Driver {
387            &self.inner
388        }
389
390        fn retarget_vp(&self, _target_vp: u32) {}
391    }
392}