Unsafe no-alloc block_on

block_on.rs

use super::*;

/// Synchronously block until a future completes.
///
/// # Safety
///
/// No waker provided to the future should ever be polled after the future returns `Poll::Ready`. The waker references
/// data on the blocking thread's stack and, as such, access to it after this point is UB because the data is no longer
/// a valid pointer.
#[allow(unused_unsafe)] // It's a stupid warning anyway. No such thing as unsafe code that's too explicit.
pub unsafe fn block_on<F: Future>(mut fut: F) -> F::Output {
    // This code is finicky! Do not touch it unless you're 100% sure that you know what you're doing and what the
    // implications of your change might be.

    use std::{thread, task::{RawWaker, RawWakerVTable}};

    // Straight from `criterion`, use `std::hint::black_box` when it is stable.
    fn black_box<T>(dummy: T) -> T {
        unsafe {
            let ret = std::ptr::read_volatile(&dummy);
            std::mem::forget(dummy);
            ret
        }
    }

    // Ensures the types don't go out of sync
    type WakeData = (AtomicUsize, thread::Thread);

    const UNUSED: usize = 0; // Waker has not been cloned (i.e: it's not being used for wakeups `wake_data` is unused)
    const USING: usize = 1; // The waker is in used but has not been woken yet (i.e: `wake_data` should NOT be dropped)
    const WOKEN: usize = 2; // The waker has been woken and dropping `wake_data` is permitted

    static VTABLE: RawWakerVTable = unsafe {
        let wake = |wake_data_ptr| {
            let wake_data = &*(wake_data_ptr as *const WakeData);
            wake_data.1.unpark();
            // Once the thread has been unparked, signal to the thread that it is permitted to drop `wake_data`.
            wake_data.0.store(WOKEN, Ordering::SeqCst);
        };
        RawWakerVTable::new(
            |wake_data_ptr| {
                let wake_data = &*(wake_data_ptr as *const WakeData);
                wake_data.0.store(USING, Ordering::SeqCst);
                RawWaker::new(wake_data_ptr, &VTABLE)
            },
            wake,
            wake,
            // Do nothing, the data is on the stack of `block_on` and will be dropped in time.
            |_wake_data_ptr| {},
        )
    };

    let wake_data: WakeData = (AtomicUsize::new(UNUSED), thread::current());
    let wake_data_ptr = &wake_data as *const _;

    // Safe because we don't return until the future is ready.
    let waker = unsafe { Waker::from_raw(RawWaker::new(wake_data_ptr as *const _, &VTABLE)) };

    loop {
        // Safe because `fut` isn't going to move until this function returns, at which point we've
        // stop polling anyway.
        let fut = unsafe { Pin::new_unchecked(&mut fut) };
        match fut.poll(&mut Context::from_waker(&waker)) {
            Poll::Ready(item) => {
                // Because spurious wakeups may occur at any time, we must ensure that we aren't dropping `wake_data`
                // before the waker has finished touching it. To do this, we spin until the 'woken' flag has been set.
                // The 'woken' flag gets set *after* the thread is unparked to ensure that unparking cannot occur after
                // `wake_data` has been dropped. In practice, this flag is very rarely observed as being false, but
                // it's a necessary requirement to ensure that UB is not triggered in the 0.001% of cases.
                while wake_data.0.load(Ordering::SeqCst) == USING {}
                // Ensure the compiler doesn't optimise `wake_data` away
                assert!(wake_data_ptr == &wake_data as *const _); // Should get compiled away if they are equal
                black_box(&wake_data);
                break item
            },
            Poll::Pending => thread::park(),
        }
    }
}