futures_test/future/
assert_unmoved.rs

1use futures_core::future::Future;
2use futures_core::task::{Context, Poll};
3use pin_utils::{unsafe_pinned, unsafe_unpinned};
4use std::marker::PhantomPinned;
5use std::pin::Pin;
6use std::ptr;
7use std::thread::panicking;
8
9/// Combinator for the
10/// [`FutureTestExt::assert_unmoved`](super::FutureTestExt::assert_unmoved)
11/// method.
12#[derive(Debug, Clone)]
13#[must_use = "futures do nothing unless you `.await` or poll them"]
14pub struct AssertUnmoved<Fut> {
15    future: Fut,
16    this_ptr: *const AssertUnmoved<Fut>,
17    _pinned: PhantomPinned,
18}
19
20impl<Fut> AssertUnmoved<Fut> {
21    unsafe_pinned!(future: Fut);
22    unsafe_unpinned!(this_ptr: *const Self);
23
24    pub(super) fn new(future: Fut) -> Self {
25        Self {
26            future,
27            this_ptr: ptr::null(),
28            _pinned: PhantomPinned,
29        }
30    }
31}
32
33impl<Fut: Future> Future for AssertUnmoved<Fut> {
34    type Output = Fut::Output;
35
36    fn poll(
37        mut self: Pin<&mut Self>,
38        cx: &mut Context<'_>,
39    ) -> Poll<Self::Output> {
40        let cur_this = &*self as *const Self;
41        if self.this_ptr.is_null() {
42            // First time being polled
43            *self.as_mut().this_ptr() = cur_this;
44        } else {
45            assert_eq!(self.this_ptr, cur_this, "Future moved between poll calls");
46        }
47        self.as_mut().future().poll(cx)
48    }
49}
50
51impl<Fut> Drop for AssertUnmoved<Fut> {
52    fn drop(&mut self) {
53        // If the thread is panicking then we can't panic again as that will
54        // cause the process to be aborted.
55        if !panicking() && !self.this_ptr.is_null() {
56            let cur_this = &*self as *const Self;
57            assert_eq!(self.this_ptr, cur_this, "Future moved before drop");
58        }
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use futures_core::future::Future;
65    use futures_core::task::{Context, Poll};
66    use futures_util::future::pending;
67    use futures_util::task::noop_waker;
68    use std::pin::Pin;
69
70    use super::AssertUnmoved;
71
72    #[test]
73    fn dont_panic_when_not_polled() {
74        // This shouldn't panic.
75        let future = AssertUnmoved::new(pending::<()>());
76        drop(future);
77    }
78
79    #[test]
80    #[should_panic(expected = "Future moved between poll calls")]
81    fn dont_double_panic() {
82        // This test should only panic, not abort the process.
83        let waker = noop_waker();
84        let mut cx = Context::from_waker(&waker);
85
86        // First we allocate the future on the stack and poll it.
87        let mut future = AssertUnmoved::new(pending::<()>());
88        let pinned_future = unsafe { Pin::new_unchecked(&mut future) };
89        assert_eq!(pinned_future.poll(&mut cx), Poll::Pending);
90
91        // Next we move it back to the heap and poll it again. This second call
92        // should panic (as the future is moved), but we shouldn't panic again
93        // whilst dropping `AssertUnmoved`.
94        let mut future = Box::new(future);
95        let pinned_boxed_future = unsafe { Pin::new_unchecked(&mut *future) };
96        assert_eq!(pinned_boxed_future.poll(&mut cx), Poll::Pending);
97    }
98}