tokio/signal/
reusable_box.rs

1use std::alloc::Layout;
2use std::future::Future;
3use std::panic::AssertUnwindSafe;
4use std::pin::Pin;
5use std::ptr::{self, NonNull};
6use std::task::{Context, Poll};
7use std::{fmt, panic};
8
9/// A reusable `Pin<Box<dyn Future<Output = T> + Send>>`.
10///
11/// This type lets you replace the future stored in the box without
12/// reallocating when the size and alignment permits this.
13pub(crate) struct ReusableBoxFuture<T> {
14    boxed: NonNull<dyn Future<Output = T> + Send>,
15}
16
17impl<T> ReusableBoxFuture<T> {
18    /// Create a new `ReusableBoxFuture<T>` containing the provided future.
19    pub(crate) fn new<F>(future: F) -> Self
20    where
21        F: Future<Output = T> + Send + 'static,
22    {
23        let boxed: Box<dyn Future<Output = T> + Send> = Box::new(future);
24
25        let boxed = Box::into_raw(boxed);
26
27        // SAFETY: Box::into_raw does not return null pointers.
28        let boxed = unsafe { NonNull::new_unchecked(boxed) };
29
30        Self { boxed }
31    }
32
33    /// Replaces the future currently stored in this box.
34    ///
35    /// This reallocates if and only if the layout of the provided future is
36    /// different from the layout of the currently stored future.
37    pub(crate) fn set<F>(&mut self, future: F)
38    where
39        F: Future<Output = T> + Send + 'static,
40    {
41        if let Err(future) = self.try_set(future) {
42            *self = Self::new(future);
43        }
44    }
45
46    /// Replaces the future currently stored in this box.
47    ///
48    /// This function never reallocates, but returns an error if the provided
49    /// future has a different size or alignment from the currently stored
50    /// future.
51    pub(crate) fn try_set<F>(&mut self, future: F) -> Result<(), F>
52    where
53        F: Future<Output = T> + Send + 'static,
54    {
55        // SAFETY: The pointer is not dangling.
56        let self_layout = {
57            let dyn_future: &(dyn Future<Output = T> + Send) = unsafe { self.boxed.as_ref() };
58            Layout::for_value(dyn_future)
59        };
60
61        if Layout::new::<F>() == self_layout {
62            // SAFETY: We just checked that the layout of F is correct.
63            unsafe {
64                self.set_same_layout(future);
65            }
66
67            Ok(())
68        } else {
69            Err(future)
70        }
71    }
72
73    /// Sets the current future.
74    ///
75    /// # Safety
76    ///
77    /// This function requires that the layout of the provided future is the
78    /// same as `self.layout`.
79    unsafe fn set_same_layout<F>(&mut self, future: F)
80    where
81        F: Future<Output = T> + Send + 'static,
82    {
83        // Drop the existing future, catching any panics.
84        let result = panic::catch_unwind(AssertUnwindSafe(|| unsafe {
85            ptr::drop_in_place(self.boxed.as_ptr());
86        }));
87
88        // Overwrite the future behind the pointer. This is safe because the
89        // allocation was allocated with the same size and alignment as the type F.
90        let self_ptr: *mut F = self.boxed.as_ptr() as *mut F;
91        // SAFETY: The pointer is valid and the layout is exactly same.
92        unsafe {
93            ptr::write(self_ptr, future);
94        }
95
96        // Update the vtable of self.boxed. The pointer is not null because we
97        // just got it from self.boxed, which is not null.
98        self.boxed = unsafe { NonNull::new_unchecked(self_ptr) };
99
100        // If the old future's destructor panicked, resume unwinding.
101        match result {
102            Ok(()) => {}
103            Err(payload) => {
104                panic::resume_unwind(payload);
105            }
106        }
107    }
108
109    /// Gets a pinned reference to the underlying future.
110    pub(crate) fn get_pin(&mut self) -> Pin<&mut (dyn Future<Output = T> + Send)> {
111        // SAFETY: The user of this box cannot move the box, and we do not move it
112        // either.
113        unsafe { Pin::new_unchecked(self.boxed.as_mut()) }
114    }
115
116    /// Polls the future stored inside this box.
117    pub(crate) fn poll(&mut self, cx: &mut Context<'_>) -> Poll<T> {
118        self.get_pin().poll(cx)
119    }
120}
121
122impl<T> Future for ReusableBoxFuture<T> {
123    type Output = T;
124
125    /// Polls the future stored inside this box.
126    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
127        Pin::into_inner(self).get_pin().poll(cx)
128    }
129}
130
131// The future stored inside ReusableBoxFuture<T> must be Send.
132unsafe impl<T> Send for ReusableBoxFuture<T> {}
133
134// The only method called on self.boxed is poll, which takes &mut self, so this
135// struct being Sync does not permit any invalid access to the Future, even if
136// the future is not Sync.
137unsafe impl<T> Sync for ReusableBoxFuture<T> {}
138
139// Just like a Pin<Box<dyn Future>> is always Unpin, so is this type.
140impl<T> Unpin for ReusableBoxFuture<T> {}
141
142impl<T> Drop for ReusableBoxFuture<T> {
143    fn drop(&mut self) {
144        unsafe {
145            drop(Box::from_raw(self.boxed.as_ptr()));
146        }
147    }
148}
149
150impl<T> fmt::Debug for ReusableBoxFuture<T> {
151    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
152        f.debug_struct("ReusableBoxFuture").finish()
153    }
154}
155
156#[cfg(test)]
157mod test {
158    use super::ReusableBoxFuture;
159    use futures::future::FutureExt;
160    use std::alloc::Layout;
161    use std::future::Future;
162    use std::pin::Pin;
163    use std::task::{Context, Poll};
164
165    #[test]
166    fn test_different_futures() {
167        let fut = async move { 10 };
168        // Not zero sized!
169        assert_eq!(Layout::for_value(&fut).size(), 1);
170
171        let mut b = ReusableBoxFuture::new(fut);
172
173        assert_eq!(b.get_pin().now_or_never(), Some(10));
174
175        b.try_set(async move { 20 })
176            .unwrap_or_else(|_| panic!("incorrect size"));
177
178        assert_eq!(b.get_pin().now_or_never(), Some(20));
179
180        b.try_set(async move { 30 })
181            .unwrap_or_else(|_| panic!("incorrect size"));
182
183        assert_eq!(b.get_pin().now_or_never(), Some(30));
184    }
185
186    #[test]
187    fn test_different_sizes() {
188        let fut1 = async move { 10 };
189        let val = [0u32; 1000];
190        let fut2 = async move { val[0] };
191        let fut3 = ZeroSizedFuture {};
192
193        assert_eq!(Layout::for_value(&fut1).size(), 1);
194        assert_eq!(Layout::for_value(&fut2).size(), 4004);
195        assert_eq!(Layout::for_value(&fut3).size(), 0);
196
197        let mut b = ReusableBoxFuture::new(fut1);
198        assert_eq!(b.get_pin().now_or_never(), Some(10));
199        b.set(fut2);
200        assert_eq!(b.get_pin().now_or_never(), Some(0));
201        b.set(fut3);
202        assert_eq!(b.get_pin().now_or_never(), Some(5));
203    }
204
205    struct ZeroSizedFuture {}
206    impl Future for ZeroSizedFuture {
207        type Output = u32;
208        fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<u32> {
209            Poll::Ready(5)
210        }
211    }
212
213    #[test]
214    fn test_zero_sized() {
215        let fut = ZeroSizedFuture {};
216        // Zero sized!
217        assert_eq!(Layout::for_value(&fut).size(), 0);
218
219        let mut b = ReusableBoxFuture::new(fut);
220
221        assert_eq!(b.get_pin().now_or_never(), Some(5));
222        assert_eq!(b.get_pin().now_or_never(), Some(5));
223
224        b.try_set(ZeroSizedFuture {})
225            .unwrap_or_else(|_| panic!("incorrect size"));
226
227        assert_eq!(b.get_pin().now_or_never(), Some(5));
228        assert_eq!(b.get_pin().now_or_never(), Some(5));
229    }
230}
tokio/signal/reusable_box.rs

tokio/signal/
reusable_box.rs
