1 use std::alloc::Layout; 2 use std::future::Future; 3 use std::panic::AssertUnwindSafe; 4 use std::pin::Pin; 5 use std::ptr::{self, NonNull}; 6 use std::task::{Context, Poll}; 7 use 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. 13 pub(crate) struct ReusableBoxFuture<T> { 14 boxed: NonNull<dyn Future<Output = T> + Send>, 15 } 16 17 impl<T> ReusableBoxFuture<T> { 18 /// Create a new `ReusableBoxFuture<T>` containing the provided future. new<F>(future: F) -> Self where F: Future<Output = T> + Send + 'static,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. set<F>(&mut self, future: F) where F: Future<Output = T> + Send + 'static,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. try_set<F>(&mut self, future: F) -> Result<(), F> where F: Future<Output = T> + Send + 'static,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`. set_same_layout<F>(&mut self, future: F) where F: Future<Output = T> + Send + 'static,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(|| { 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 ptr::write(self_ptr, future); 92 93 // Update the vtable of self.boxed. The pointer is not null because we 94 // just got it from self.boxed, which is not null. 95 self.boxed = NonNull::new_unchecked(self_ptr); 96 97 // If the old future's destructor panicked, resume unwinding. 98 match result { 99 Ok(()) => {} 100 Err(payload) => { 101 panic::resume_unwind(payload); 102 } 103 } 104 } 105 106 /// Gets a pinned reference to the underlying future. get_pin(&mut self) -> Pin<&mut (dyn Future<Output = T> + Send)>107 pub(crate) fn get_pin(&mut self) -> Pin<&mut (dyn Future<Output = T> + Send)> { 108 // SAFETY: The user of this box cannot move the box, and we do not move it 109 // either. 110 unsafe { Pin::new_unchecked(self.boxed.as_mut()) } 111 } 112 113 /// Polls the future stored inside this box. poll(&mut self, cx: &mut Context<'_>) -> Poll<T>114 pub(crate) fn poll(&mut self, cx: &mut Context<'_>) -> Poll<T> { 115 self.get_pin().poll(cx) 116 } 117 } 118 119 impl<T> Future for ReusableBoxFuture<T> { 120 type Output = T; 121 122 /// Polls the future stored inside this box. poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T>123 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> { 124 Pin::into_inner(self).get_pin().poll(cx) 125 } 126 } 127 128 // The future stored inside ReusableBoxFuture<T> must be Send. 129 unsafe impl<T> Send for ReusableBoxFuture<T> {} 130 131 // The only method called on self.boxed is poll, which takes &mut self, so this 132 // struct being Sync does not permit any invalid access to the Future, even if 133 // the future is not Sync. 134 unsafe impl<T> Sync for ReusableBoxFuture<T> {} 135 136 // Just like a Pin<Box<dyn Future>> is always Unpin, so is this type. 137 impl<T> Unpin for ReusableBoxFuture<T> {} 138 139 impl<T> Drop for ReusableBoxFuture<T> { drop(&mut self)140 fn drop(&mut self) { 141 unsafe { 142 drop(Box::from_raw(self.boxed.as_ptr())); 143 } 144 } 145 } 146 147 impl<T> fmt::Debug for ReusableBoxFuture<T> { fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result148 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 149 f.debug_struct("ReusableBoxFuture").finish() 150 } 151 } 152 153 #[cfg(test)] 154 mod test { 155 use super::ReusableBoxFuture; 156 use futures::future::FutureExt; 157 use std::alloc::Layout; 158 use std::future::Future; 159 use std::pin::Pin; 160 use std::task::{Context, Poll}; 161 162 #[test] test_different_futures()163 fn test_different_futures() { 164 let fut = async move { 10 }; 165 // Not zero sized! 166 assert_eq!(Layout::for_value(&fut).size(), 1); 167 168 let mut b = ReusableBoxFuture::new(fut); 169 170 assert_eq!(b.get_pin().now_or_never(), Some(10)); 171 172 b.try_set(async move { 20 }) 173 .unwrap_or_else(|_| panic!("incorrect size")); 174 175 assert_eq!(b.get_pin().now_or_never(), Some(20)); 176 177 b.try_set(async move { 30 }) 178 .unwrap_or_else(|_| panic!("incorrect size")); 179 180 assert_eq!(b.get_pin().now_or_never(), Some(30)); 181 } 182 183 #[test] test_different_sizes()184 fn test_different_sizes() { 185 let fut1 = async move { 10 }; 186 let val = [0u32; 1000]; 187 let fut2 = async move { val[0] }; 188 let fut3 = ZeroSizedFuture {}; 189 190 assert_eq!(Layout::for_value(&fut1).size(), 1); 191 assert_eq!(Layout::for_value(&fut2).size(), 4004); 192 assert_eq!(Layout::for_value(&fut3).size(), 0); 193 194 let mut b = ReusableBoxFuture::new(fut1); 195 assert_eq!(b.get_pin().now_or_never(), Some(10)); 196 b.set(fut2); 197 assert_eq!(b.get_pin().now_or_never(), Some(0)); 198 b.set(fut3); 199 assert_eq!(b.get_pin().now_or_never(), Some(5)); 200 } 201 202 struct ZeroSizedFuture {} 203 impl Future for ZeroSizedFuture { 204 type Output = u32; poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<u32>205 fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<u32> { 206 Poll::Ready(5) 207 } 208 } 209 210 #[test] test_zero_sized()211 fn test_zero_sized() { 212 let fut = ZeroSizedFuture {}; 213 // Zero sized! 214 assert_eq!(Layout::for_value(&fut).size(), 0); 215 216 let mut b = ReusableBoxFuture::new(fut); 217 218 assert_eq!(b.get_pin().now_or_never(), Some(5)); 219 assert_eq!(b.get_pin().now_or_never(), Some(5)); 220 221 b.try_set(ZeroSizedFuture {}) 222 .unwrap_or_else(|_| panic!("incorrect size")); 223 224 assert_eq!(b.get_pin().now_or_never(), Some(5)); 225 assert_eq!(b.get_pin().now_or_never(), Some(5)); 226 } 227 } 228