Skip to main content

compio_runtime/future/combinator/
mod.rs

1//! Future combinators.
2
3mod cancel;
4mod personality;
5
6use std::borrow::Cow;
7
8pub use cancel::*;
9use compio_driver::Extra;
10pub use personality::*;
11
12use crate::CancelToken;
13
14#[non_exhaustive]
15#[derive(Debug, Default)]
16pub(crate) struct Ext<'a> {
17    personality: Option<u16>,
18    cancel: Option<Cow<'a, CancelToken>>,
19}
20
21impl<'a> Ext<'a> {
22    pub fn to_owned(&self) -> Ext<'static> {
23        Ext {
24            personality: self.personality,
25            cancel: self
26                .cancel
27                .as_ref()
28                .map(|x| Cow::Owned(x.clone().into_owned())),
29        }
30    }
31}
32
33impl<'a> Ext<'a> {
34    pub fn with_personality(&self, personality: u16) -> Self {
35        Self {
36            personality: Some(personality),
37            cancel: self.cancel.clone(),
38        }
39    }
40
41    pub fn with_cancel(&self, token: &'a CancelToken) -> Self {
42        Self {
43            personality: self.personality,
44            cancel: Some(Cow::Borrowed(token)),
45        }
46    }
47
48    pub fn get_cancel(&self) -> Option<&CancelToken> {
49        self.cancel.as_deref()
50    }
51
52    pub fn set_extra(&self, extra: &mut Extra) -> bool {
53        let mut changed = false;
54        if let Some(personality) = self.personality {
55            extra.set_personality(personality);
56            changed = true;
57        }
58        changed
59    }
60}
61
62/// Extension trait for futures.
63///
64/// # Implementation
65///
66/// Extra data are passed down to runtime when the combinators are polled using
67/// a custom [`Waker`], and those data are single-threaded. This means
68/// - when [`Waker`]s are sent to other threads, the data will be lost.
69/// - when using a "sub-executor" like `FuturesUnordered`, which also creates
70///   its own waker, the data will be lost.
71///
72/// So try to keep the path from the wrapped future to runtime clean, something
73/// like this will generally work:
74///
75/// ```rust,ignore
76/// use std::vec::Vec;
77///
78/// use compio::runtime::{FutureExt, CancelToken};
79/// use compio::fs::File;
80///
81/// let file = File::open("/tmp/file");
82/// let cancel = CancelToken::new();
83/// file.read(Vec::with_capacity(1024)).with_cancel(cancel.clone()).await;
84/// ```
85///
86/// [`Waker`]: std::task::Waker
87pub trait FutureExt {
88    /// Sets the personality for this future.
89    ///
90    /// This only takes effect on io-uring drivers and will be ignored on other
91    /// ones.
92    fn with_personality(self, personality: u16) -> WithPersonality<Self>
93    where
94        Self: Sized,
95    {
96        WithPersonality::new(self, personality)
97    }
98
99    /// Sets the cancel token for this future.
100    ///
101    /// If multiple [`CancelToken`]s are set, the innermost one (the one being
102    /// polled last) will take precedence.
103    fn with_cancel(self, token: CancelToken) -> WithCancel<Self>
104    where
105        Self: Sized,
106    {
107        WithCancel::new(self, token)
108    }
109}
110
111impl<F: Future + ?Sized> FutureExt for F {}