Waker Practice: Attempt to Unwrap a Future

Wu Yu Wei published on
3 min, 583 words

Many enum types like Option and Result provide mehtod like unwrap to retieve the field immediately. It panics if it can not get the value we want. This gives me an idea to do a silly exercise. What if we can also unwrap a future? "But why?" you said. Because why not. Consider when we use .await syntax on futures from async blocks or functions, it will poll the future and see if it is ready. The compiler transforms the whole block into some sort of state machine, and each .await represents a potential state transition. And now I would like to make this exit to be unwrap which attempts to resolve a future immediately.

Wake up Mr. Freeman

The entire method is quite easy actually, we just need to match the poll and then return the value. But the missing puzzle is the waker. In real executor system, it will provide all required contexts for us. But we are going to poll it manually ourselves, we need to define our waker function. Fortunately, the standard library already provide all the components we want. We just have to assemble them, so let's wake up and smell the ashes!

pub fn waker<F: Fn() + Send + Sync + 'static>(f: F) -> Waker {
    let raw = Arc::into_raw(Arc::new(f)) as *const ();
    let vtable = &RawWakerHelper::<F>::VTABLE;
    unsafe { Waker::from_raw(RawWaker::new(raw, vtable)) }
}

We create a waker function which outputs, well, a Waker. This type is a wrapper of inner instance RawWaker which provides customized wakeup behavior. It consists of a pointer to data and vtable to customize the behavior we mentioned before. Since there's a method in vtable for cloning and we want the waker and future we are dealing with are still the same one, we wrap up with Arc and turn it back to raw pointer. I will show you what's going to be next.

struct RawWakerHelper<F>(F);

impl<F: Fn() + Send + Sync + 'static> RawWakerHelper<F> {
    const VTABLE: RawWakerVTable = RawWakerVTable::new(
        Self::clone_waker,
        Self::wake,
        Self::wake_by_ref,
        Self::drop_waker,
    );

    unsafe fn clone_waker(ptr: *const ()) -> RawWaker {
        let arc = ManuallyDrop::new(Arc::from_raw(ptr as *const F));
        mem::forget(arc.clone());
        RawWaker::new(ptr, &Self::VTABLE)
    }

    unsafe fn wake(ptr: *const ()) {
        let arc = Arc::from_raw(ptr as *const F);
        (arc)();
    }

    unsafe fn wake_by_ref(ptr: *const ()) {
        let arc = ManuallyDrop::new(Arc::from_raw(ptr as *const F));
        (arc)();
    }

    unsafe fn drop_waker(ptr: *const ()) {
        drop(Arc::from_raw(ptr as *const F));
    }
}

As you can see the rest parameters that RawWakerVTable requires are here. Since it demands raw pointers which already touch the unsafe field. We are able to play around out Arc pointer in several way and still be sound. Take a look a little bit more if needed, all methods and modules used are documented in standard library.

Time to uwrap

Once our waker is done, the last thing left to do is to write out one and only function unwrap. And it's way simpler than you thought:

pub fn unwrap<F: Future>(f: F) -> F::Output {
    let waker = waker(|| ());
    let mut f = Box::pin(f);
    match f.as_mut().poll(&mut Context::from_waker(&waker)) {
        Poll::Ready(output) => output,
        Poll::Pending       => panic!("future not ready"),
    }
}

That's it! Although this is probably very useless unless you are prototyping some executor or reactor systems. Still a good practice imho. Here's the link to the playground. Hope you can get more familiar with modules beneath the future system.