ntex_util/channel/
pool.rs

1//! A one-shot pool, futures-aware channel.
2use slab::Slab;
3use std::{fmt, future::Future, pin::Pin, task::Context, task::Poll};
4
5use super::{Canceled, cell::Cell};
6use crate::task::LocalWaker;
7
8/// Creates a new futures-aware pool of one-shot channels.
9pub fn new<T>() -> Pool<T> {
10    Pool(Cell::new(Slab::new()))
11}
12
13/// A futures-aware pool of one-shot channels.
14pub struct Pool<T>(Cell<Slab<Inner<T>>>);
15
16impl<T> fmt::Debug for Pool<T> {
17    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
18        f.debug_struct("Pool")
19            .field("size", &self.0.get_ref().len())
20            .finish()
21    }
22}
23
24bitflags::bitflags! {
25    #[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
26    struct Flags: u8 {
27        const SENDER = 0b0000_0001;
28        const RECEIVER = 0b0000_0010;
29    }
30}
31
32#[derive(Debug)]
33struct Inner<T> {
34    flags: Flags,
35    value: Option<T>,
36    tx_waker: LocalWaker,
37    rx_waker: LocalWaker,
38}
39
40impl<T> Default for Pool<T> {
41    fn default() -> Pool<T> {
42        new()
43    }
44}
45
46impl<T> Pool<T> {
47    /// Create a new one-shot channel.
48    pub fn channel(&self) -> (Sender<T>, Receiver<T>) {
49        let token = self.0.get_mut().insert(Inner {
50            flags: Flags::all(),
51            value: None,
52            tx_waker: LocalWaker::default(),
53            rx_waker: LocalWaker::default(),
54        });
55
56        (
57            Sender {
58                token,
59                inner: self.0.clone(),
60            },
61            Receiver {
62                token,
63                inner: self.0.clone(),
64            },
65        )
66    }
67
68    /// Shrinks the capacity of the pool as much as possible.
69    pub fn shrink_to_fit(&self) {
70        self.0.get_mut().shrink_to_fit();
71    }
72}
73
74impl<T> Clone for Pool<T> {
75    fn clone(&self) -> Self {
76        Pool(self.0.clone())
77    }
78}
79
80/// Represents the completion half of a oneshot through which the result of a
81/// computation is signaled.
82#[derive(Debug)]
83pub struct Sender<T> {
84    token: usize,
85    inner: Cell<Slab<Inner<T>>>,
86}
87
88/// A future representing the completion of a computation happening elsewhere in
89/// memory.
90#[derive(Debug)]
91#[must_use = "futures do nothing unless polled"]
92pub struct Receiver<T> {
93    token: usize,
94    inner: Cell<Slab<Inner<T>>>,
95}
96
97#[allow(clippy::mut_from_ref)]
98fn get_inner<T>(inner: &Cell<Slab<Inner<T>>>, token: usize) -> &mut Inner<T> {
99    unsafe { inner.get_mut().get_unchecked_mut(token) }
100}
101
102// The oneshots do not ever project Pin to the inner T
103impl<T> Unpin for Receiver<T> {}
104impl<T> Unpin for Sender<T> {}
105
106impl<T> Sender<T> {
107    /// Completes this oneshot with a successful result.
108    ///
109    /// This function will consume `self` and indicate to the other end, the
110    /// `Receiver`, that the value provided is the result of the computation this
111    /// represents.
112    ///
113    /// If the value is successfully enqueued for the remote end to receive,
114    /// then `Ok(())` is returned. If the receiving end was dropped before
115    /// this function was called, however, then `Err` is returned with the value
116    /// provided.
117    pub fn send(self, val: T) -> Result<(), T> {
118        let inner = get_inner(&self.inner, self.token);
119        if inner.flags.contains(Flags::RECEIVER) {
120            inner.value = Some(val);
121            inner.rx_waker.wake();
122            Ok(())
123        } else {
124            Err(val)
125        }
126    }
127
128    /// Tests to see whether this `Sender`'s corresponding `Receiver`
129    /// has gone away.
130    pub fn is_canceled(&self) -> bool {
131        !get_inner(&self.inner, self.token)
132            .flags
133            .contains(Flags::RECEIVER)
134    }
135
136    /// Polls the channel to determine if receiving path is dropped
137    pub fn poll_canceled(&self, cx: &mut Context<'_>) -> Poll<()> {
138        let inner = get_inner(&self.inner, self.token);
139
140        if inner.flags.contains(Flags::RECEIVER) {
141            inner.tx_waker.register(cx.waker());
142            Poll::Pending
143        } else {
144            Poll::Ready(())
145        }
146    }
147}
148
149impl<T> Drop for Sender<T> {
150    fn drop(&mut self) {
151        let inner = get_inner(&self.inner, self.token);
152        if inner.flags.contains(Flags::RECEIVER) {
153            inner.rx_waker.wake();
154            inner.flags.remove(Flags::SENDER);
155        } else {
156            self.inner.get_mut().remove(self.token);
157        }
158    }
159}
160
161impl<T> Receiver<T> {
162    /// Polls for the value.
163    ///
164    /// Returns [`Canceled`] if the sender is dropped without sending a value.
165    pub fn poll_recv(&self, cx: &mut Context<'_>) -> Poll<Result<T, Canceled>> {
166        let inner = get_inner(&self.inner, self.token);
167
168        // If we've got a value, then skip the logic below as we're done.
169        if let Some(val) = inner.value.take() {
170            return Poll::Ready(Ok(val));
171        }
172
173        // Check if sender is dropped and return error if it is.
174        if inner.flags.contains(Flags::SENDER) {
175            inner.rx_waker.register(cx.waker());
176            Poll::Pending
177        } else {
178            Poll::Ready(Err(Canceled))
179        }
180    }
181}
182
183impl<T> Drop for Receiver<T> {
184    fn drop(&mut self) {
185        let inner = get_inner(&self.inner, self.token);
186        if inner.flags.contains(Flags::SENDER) {
187            inner.tx_waker.wake();
188            inner.flags.remove(Flags::RECEIVER);
189        } else {
190            self.inner.get_mut().remove(self.token);
191        }
192    }
193}
194
195impl<T> Future for Receiver<T> {
196    type Output = Result<T, Canceled>;
197
198    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
199        self.poll_recv(cx)
200    }
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206    use crate::future::lazy;
207
208    #[ntex::test]
209    async fn test_pool() {
210        let p = new();
211        p.shrink_to_fit();
212        assert!(format!("{p:?}").contains("Pool"));
213
214        let (tx, rx) = p.channel();
215        assert!(format!("{tx:?}").contains("Sender"));
216        assert!(format!("{rx:?}").contains("Receiver"));
217
218        tx.send("test").unwrap();
219        assert_eq!(rx.await.unwrap(), "test");
220        assert!(format!("{Canceled}").contains("canceled"));
221        assert!(format!("{Canceled:?}").contains("Canceled"));
222
223        let p2 = p.clone();
224        let (tx, rx) = p2.channel();
225        assert!(!tx.is_canceled());
226        drop(rx);
227        assert!(tx.is_canceled());
228        assert!(tx.send("test").is_err());
229
230        let (tx, rx) = new::<&'static str>().channel();
231        drop(tx);
232        assert!(rx.await.is_err());
233
234        let (tx, mut rx) = new::<&'static str>().channel();
235        assert_eq!(lazy(|cx| Pin::new(&mut rx).poll(cx)).await, Poll::Pending);
236        tx.send("test").unwrap();
237        assert_eq!(rx.await.unwrap(), "test");
238
239        let (tx, mut rx) = new::<&'static str>().channel();
240        assert_eq!(lazy(|cx| Pin::new(&mut rx).poll(cx)).await, Poll::Pending);
241        drop(tx);
242        assert!(rx.await.is_err());
243
244        let (mut tx, rx) = new::<&'static str>().channel();
245        assert!(!tx.is_canceled());
246        assert_eq!(
247            lazy(|cx| Pin::new(&mut tx).poll_canceled(cx)).await,
248            Poll::Pending
249        );
250        drop(rx);
251        assert!(tx.is_canceled());
252        assert_eq!(
253            lazy(|cx| Pin::new(&mut tx).poll_canceled(cx)).await,
254            Poll::Ready(())
255        );
256
257        let p = Pool::default();
258        let (tx, rx) = p.channel();
259        tx.send("test").unwrap();
260        assert_eq!(rx.await.unwrap(), "test");
261    }
262}
ntex_util/channel/pool.rs

ntex_util/channel/
pool.rs
