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;
10use futures_util::Stream;
11pub use personality::*;
12
13use crate::CancelToken;
14
15#[non_exhaustive]
16#[derive(Debug, Default)]
17pub(crate) struct Ext<'a> {
18    personality: Option<u16>,
19    cancel: Option<Cow<'a, CancelToken>>,
20}
21
22impl<'a> Ext<'a> {
23    pub fn to_owned(&self) -> Ext<'static> {
24        Ext {
25            personality: self.personality,
26            cancel: self
27                .cancel
28                .as_ref()
29                .map(|x| Cow::Owned(x.clone().into_owned())),
30        }
31    }
32}
33
34impl<'a> Ext<'a> {
35    pub fn with_personality(&self, personality: u16) -> Self {
36        Self {
37            personality: Some(personality),
38            cancel: self.cancel.clone(),
39        }
40    }
41
42    pub fn with_cancel(&self, token: &'a CancelToken) -> Self {
43        Self {
44            personality: self.personality,
45            cancel: Some(Cow::Borrowed(token)),
46        }
47    }
48
49    pub fn get_cancel(&self) -> Option<&CancelToken> {
50        self.cancel.as_deref()
51    }
52
53    pub fn set_extra(&self, extra: &mut Extra) -> bool {
54        let mut changed = false;
55        if let Some(personality) = self.personality {
56            extra.set_personality(personality);
57            changed = true;
58        }
59        changed
60    }
61}
62
63/// Extension trait for futures.
64///
65/// # Implementation
66///
67/// Extra data are passed down to runtime when the combinators are polled using
68/// a custom [`Waker`], and those data are single-threaded. This means
69/// - when [`Waker`]s are sent to other threads, the data will be lost.
70/// - when using a "sub-executor" like `FuturesUnordered`, which also creates
71///   its own waker, the data will be lost.
72///
73/// So try to keep the path from the wrapped future to runtime clean, something
74/// like this will generally work:
75///
76/// ```rust,ignore
77/// use std::vec::Vec;
78///
79/// use compio::runtime::{FutureExt, CancelToken};
80/// use compio::fs::File;
81///
82/// let file = File::open("/tmp/file");
83/// let cancel = CancelToken::new();
84/// file.read(Vec::with_capacity(1024)).with_cancel(cancel.clone()).await;
85/// ```
86///
87/// [`Waker`]: std::task::Waker
88pub trait FutureExt {
89    /// Sets the personality for this future.
90    ///
91    /// This only takes effect on io-uring drivers and will be ignored on other
92    /// ones.
93    fn with_personality(self, personality: u16) -> WithPersonality<Self>
94    where
95        Self: Sized,
96    {
97        WithPersonality::new(self, personality)
98    }
99
100    /// Sets the cancel token for this future.
101    ///
102    /// If multiple [`CancelToken`]s are set, the innermost one (the one being
103    /// polled last) will take precedence.
104    fn with_cancel(self, token: CancelToken) -> WithCancel<Self>
105    where
106        Self: Sized,
107    {
108        WithCancel::new(self, token)
109    }
110}
111
112impl<F: Future + ?Sized> FutureExt for F {}
113
114/// Extension trait for streams.
115///
116/// # Implementation
117///
118/// Extra data are passed down to runtime when the combinators are polled using
119/// a custom [`Waker`], and those data are single-threaded. This means
120/// - when [`Waker`]s are sent to other threads, the data will be lost.
121/// - when using a "sub-executor" like `FuturesUnordered`, which also creates
122///   its own waker, the data will be lost.
123///
124/// So try to keep the path from the wrapped stream to runtime clean, something
125/// like this will generally work:
126///
127/// ```rust,ignore
128/// use std::vec::Vec;
129///
130/// use compio::runtime::{StreamExt, CancelToken};
131/// use compio::net::TcpStream;
132///
133/// let input = TcpStream::connect("127.0.0.1:8000").await.unwrap();
134/// let cancel = CancelToken::new();
135/// let mut input_stream = input.read_multi(0).with_cancel(cancel.clone());
136/// while let Some(buf) = input_stream.next().await {
137///     let _ = buf.unwrap();
138/// }
139/// ```
140///
141/// [`Waker`]: std::task::Waker
142pub trait StreamExt {
143    /// Sets the personality for this stream.
144    ///
145    /// This only takes effect on io-uring drivers and will be ignored on other
146    /// ones.
147    fn with_personality(self, personality: u16) -> WithPersonality<Self>
148    where
149        Self: Sized,
150    {
151        WithPersonality::new(self, personality)
152    }
153
154    /// Sets the cancel token for this stream.
155    ///
156    /// If multiple [`CancelToken`]s are set, the innermost one (the one being
157    /// polled last) will take precedence.
158    fn with_cancel(self, token: CancelToken) -> WithCancel<Self>
159    where
160        Self: Sized,
161    {
162        WithCancel::new(self, token)
163    }
164}
165
166impl<F: Stream + ?Sized> StreamExt for F {}
compio_runtime/future/combinator/mod.rs

compio_runtime/future/combinator/
mod.rs
