Skip to content

Commit

Permalink
util: simplify ReusableBoxFuture::try_set logic
Browse files Browse the repository at this point in the history
  • Loading branch information
yukibtc committed Dec 30, 2024
1 parent b3ff911 commit fe96b41
Showing 1 changed file with 7 additions and 78 deletions.
85 changes: 7 additions & 78 deletions tokio-util/src/sync/reusable_box.rs
Original file line number Diff line number Diff line change
@@ -1,9 +1,6 @@
use std::alloc::Layout;
use std::fmt;
use std::future::{self, Future};
use std::mem::{self, ManuallyDrop};
use std::future::Future;
use std::pin::Pin;
use std::ptr;
use std::task::{Context, Poll};

/// A reusable `Pin<Box<dyn Future<Output = T> + Send + 'a>>`.
Expand Down Expand Up @@ -47,24 +44,14 @@ impl<'a, T> ReusableBoxFuture<'a, T> {
where
F: Future<Output = T> + Send + 'a,
{
// If we try to inline the contents of this function, the type checker complains because
// the bound `T: 'a` is not satisfied in the call to `pending()`. But by putting it in an
// inner function that doesn't have `T` as a generic parameter, we implicitly get the bound
// `F::Output: 'a` transitively through `F: 'a`, allowing us to call `pending()`.
#[inline(always)]
fn real_try_set<'a, F>(
this: &mut ReusableBoxFuture<'a, F::Output>,
future: F,
) -> Result<(), F>
where
F: Future + Send + 'a,
if size_of_val(&*self.boxed) == size_of_val(&future)
&& align_of_val(&*self.boxed) == align_of_val(&future)
{
// future::Pending<T> is a ZST so this never allocates.
let boxed = mem::replace(&mut this.boxed, Box::pin(future::pending()));
reuse_pin_box(boxed, future, |boxed| this.boxed = Pin::from(boxed))
self.boxed = Box::pin(future);
Ok(())
} else {
Err(future)
}

real_try_set(self, future)
}

/// Get a pinned reference to the underlying future.
Expand Down Expand Up @@ -97,61 +84,3 @@ impl<T> fmt::Debug for ReusableBoxFuture<'_, T> {
f.debug_struct("ReusableBoxFuture").finish()
}
}

fn reuse_pin_box<T: ?Sized, U, O, F>(boxed: Pin<Box<T>>, new_value: U, callback: F) -> Result<O, U>
where
F: FnOnce(Box<U>) -> O,
{
let layout = Layout::for_value::<T>(&*boxed);
if layout != Layout::new::<U>() {
return Err(new_value);
}

// SAFETY: We don't ever construct a non-pinned reference to the old `T` from now on, and we
// always drop the `T`.
let raw: *mut T = Box::into_raw(unsafe { Pin::into_inner_unchecked(boxed) });

// When dropping the old value panics, we still want to call `callback` — so move the rest of
// the code into a guard type.
let guard = CallOnDrop::new(|| {
let raw: *mut U = raw.cast::<U>();
unsafe { raw.write(new_value) };

// SAFETY:
// - `T` and `U` have the same layout.
// - `raw` comes from a `Box` that uses the same allocator as this one.
// - `raw` points to a valid instance of `U` (we just wrote it in).
let boxed = unsafe { Box::from_raw(raw) };

callback(boxed)
});

// Drop the old value.
unsafe { ptr::drop_in_place(raw) };

// Run the rest of the code.
Ok(guard.call())
}

struct CallOnDrop<O, F: FnOnce() -> O> {
f: ManuallyDrop<F>,
}

impl<O, F: FnOnce() -> O> CallOnDrop<O, F> {
fn new(f: F) -> Self {
let f = ManuallyDrop::new(f);
Self { f }
}
fn call(self) -> O {
let mut this = ManuallyDrop::new(self);
let f = unsafe { ManuallyDrop::take(&mut this.f) };
f()
}
}

impl<O, F: FnOnce() -> O> Drop for CallOnDrop<O, F> {
fn drop(&mut self) {
let f = unsafe { ManuallyDrop::take(&mut self.f) };
f();
}
}

0 comments on commit fe96b41

Please sign in to comment.