Is this reusable box sound?

The following code implements a reusable boxed future. I am pretty sure that is is ok, but it would be nice with a second pair of eyes.

use std::alloc::Layout;
use std::future::Future;
use std::panic::AssertUnwindSafe;
use std::pin::Pin;
use std::ptr::{self, NonNull};
use std::task::{Context, Poll};
use std::{fmt, panic};

/// A reusable `Pin<Box<dyn Future<Output = T> + Send>>`.
///
/// This type lets you replace the future stored in the box without
/// reallocating when the size and alignment permits this.
pub struct ReusableBoxFuture<T> {
    boxed: NonNull<dyn Future<Output = T> + Send>,
    layout: Layout,
}

impl<T> ReusableBoxFuture<T> {
    /// Create a new `ReusableBoxFuture<T>` containing the provided future.
    pub fn new<F>(future: F) -> Self
    where
        F: Future<Output = T> + Send + 'static,
    {
        let layout = Layout::for_value(&future);
        let boxed: Box<dyn Future<Output = T> + Send> = Box::new(future);

        let boxed = Box::into_raw(boxed);

        // SAFETY: Box::into_raw does not return null pointers.
        let boxed = unsafe { NonNull::new_unchecked(boxed) };

        Self { boxed, layout }
    }

    /// Replace the future currently stored in this box.
    ///
    /// This reallocates if and only if the layout of the provided future is
    /// different from the layout of the currently stored future.
    pub fn set<F>(&mut self, future: F)
    where
        F: Future<Output = T> + Send + 'static,
    {
        let layout = Layout::for_value(&future);

        if layout == self.layout {
            // SAFETY: We just checked that the layout of F is correct.
            unsafe {
                self.set_same_layout(future);
            }
        } else {
            *self = Self::new(future);
        }
    }

    /// Replace the future currently stored in this box.
    ///
    /// This function never reallocates, but returns an error if the provided
    /// future has a different size or alignment from the currently stored
    /// future.
    pub fn try_set<F>(&mut self, future: F) -> Result<(), F>
    where
        F: Future<Output = T> + Send + 'static,
    {
        let layout = Layout::for_value(&future);

        if layout == self.layout {
            // SAFETY: We just checked that the layout of F is correct.
            unsafe {
                self.set_same_layout(future);
            }

            Ok(())
        } else {
            Err(future)
        }
    }

    /// Set the current future.
    ///
    /// # Safety
    ///
    /// This function requires that the layout of the provided future is the
    /// same as `self.layout`.
    unsafe fn set_same_layout<F>(&mut self, future: F)
    where
        F: Future<Output = T> + Send + 'static,
    {
        // Drop the existing future, catching any panics.
        let result = panic::catch_unwind(AssertUnwindSafe(|| {
            ptr::drop_in_place(self.boxed.as_ptr());
        }));

        // Overwrite the future behind the pointer. This is safe because the
        // allocation was allocated with the same size and alignment as the type F.
        let self_ptr: *mut F = self.boxed.as_ptr() as *mut F;
        ptr::write(self_ptr, future);

        // Update the vtable of self.boxed. The pointer is not null because we
        // just got it from self.boxed, which is not null.
        self.boxed = NonNull::new_unchecked(self_ptr);

        // If the old future's destructor panicked, resume unwinding.
        match result {
            Ok(()) => {}
            Err(payload) => {
                panic::resume_unwind(payload);
            }
        }
    }

    /// Get a pinned reference to the underlying future.
    pub fn get_pin(&mut self) -> Pin<&mut (dyn Future<Output = T> + Send)> {
        unsafe { Pin::new_unchecked(self.boxed.as_mut()) }
    }

    /// Poll the future stored inside this box.
    pub fn poll(&mut self, cx: &mut Context<'_>) -> Poll<T> {
        self.get_pin().poll(cx)
    }
}

impl<T> Future for ReusableBoxFuture<T> {
    type Output = T;

    /// Poll the future stored inside this box.
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
        Pin::into_inner(self).get_pin().poll(cx)
    }
}

// The future stored inside ReusableBoxFuture<T> must be Send.
unsafe impl<T> Send for ReusableBoxFuture<T> {}

// The only method called on self.boxed is poll, which takes &mut self, so this
// struct being Sync does not permit any invalid access to the Future, even if
// the future is not Sync.
unsafe impl<T> Sync for ReusableBoxFuture<T> {}

// Just like a Pin<Box<dyn Future>> is always Unpin, so is this type.
impl<T> Unpin for ReusableBoxFuture<T> {}

impl<T> Drop for ReusableBoxFuture<T> {
    fn drop(&mut self) {
        unsafe {
            drop(Box::from_raw(self.boxed.as_ptr()));
        }
    }
}

impl<T> fmt::Debug for ReusableBoxFuture<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ReusableBoxFuture").finish()
    }
}

relevant PR

(note: the PR has since been changed from the above snippet)

2 Likes

original playground link

This is safe, I just have a few nits:

  • Layout::for_value can be replaced with use Layout::new::<F>
  • On nightly, you don't need to store a Layout at all, you can use Layout::for_value_raw
  • I would use a scopeguard instead of catch_unwind/resume_unwind to allow for no_std+alloc usage

It seems you can extract the internal layout of the &dyn Trait by Layout::for_value(), instead to hold it in a separate field.

1 Like

I actually tried to use a scope-guard first, but found it a lot more annoying to write because the destructor of the guard has to perform a move of the future.

Does this seem right?

pub fn try_set<F>(&mut self, future: F) -> Result<(), F>
where
    F: Future<Output = T> + Send + 'static,
{
    let dyn_future: &(dyn Future<Output = T> + Send) = unsafe {
        self.boxed.as_ref()
    };
    let self_layout = Layout::for_value(dyn_future);

    if Layout::new::<F>() == self_layout {
        // SAFETY: We just checked that the layout of F is correct.
        unsafe {
            self.set_same_layout(future);
        }

        Ok(())
    } else {
        Err(future)
    }
}

Wait, I just looked into this closer, are you trying to overwrite the data portion of the fat-pointer in set_same_layout? If so, that's unspecified behavior (because layouts are unspecified). edit: nope, I'm misreading this

All I'm doing is replacing one NonNull<dyn Future> for another.

At the start of the set/try_set the self.boxed must be a ptr converted from the live box so it should be safe. But it would be slightly safer to wrap the &dyn Future into another scope to make sure it will not be used after the ptr overwrite.

let self_layout = {
    let dyn_future = ...;
    Layout::for_value(dyn_future);
};

I could do it like that.

Here's a version that should work on no_std + alloc, you can even use Box! Rust Playground

1 Like

We are trying to avoid adding new dependencies to the Tokio crates, so I would prefer to avoid the scopeguard dependency. Additionally, the Tokio crates already don't support no_std.

1 Like

Ok, I just wanted to mention it, in case it was possible.

1 Like