RawWakerVTable functions, did I write them correctly?

help
1

Hi folks, I am trying to understand how tokio does what it does by writing a small executor of mine. I am at a point where am writing a VTable for constructing a Waker. I have written all the 4 functions which are required to construct a RawWakerVTable. Please let me know if they look okay to you..thanks!

also, std::mem::forget(task) is the correct thing to do cause we are only given a pointer and we do not own an Arc<Task> therefore we shouldn't drop it, correct?

// RawWakerVTable functions
// for each operation on the Waker, the associated function in the vtable will be called.
mod vtable {
    use crate::*;
    pub fn clone(data: *const ()) -> RawWaker {
        RawWaker::new(data, &Executor::VTABLE)
    }

    pub fn wake(data: *const ()) {
        let task: Arc<Task> = unsafe { Arc::from_raw(data as _) };
        // schedule the task on the executor again
        task.send(task.clone());
        // we don't own `data` so we don't drop it?
        std::mem::forget(task);
    }

    // I don't know how this would differ from wake...
    pub fn wake_by_ref(data: *const ()) {
        let task: Arc<Task> = unsafe { Arc::from_raw(data as _) };
        // schedule the task on the executor again
        task.send(task.clone());
        // we don't own `data` so we don't drop it?
        std::mem::forget(task);
    }

    pub fn drop(data: *const ()) {
        unsafe { Arc::<Task>::from_raw(data as _) };
    }
}

   const VTABLE: RawWakerVTable = RawWakerVTable::new(
        vtable::clone,
        vtable::wake,
        vtable::wake_by_ref,
        vtable::drop,
    );

at its core the main purpose of writing all of this is to schedule the task back to the executor so we could call poll on it again, correct?

Ik, I could have also implemented ArcWake but, I wanted know what is actually happening under the hood so I just did it this way :slight_smile:

2

whether the implementation is good also depends on the definition of the executor, and tasks, etc, the vtable alone is not enough to determine the correctness, since the types are erased.

it is different for wake() and wake_by_ref().

for wake(), its argument is an owned pointer, while for wake_by_ref(), its argument is a non-owning pointer. these correspond to Waker::wake(self) and Waker::wake_by_ref(&self), respectively.

although I can't determine the correctness, assuming your waker is constructed with Arc::into_raw(), your implementation of wake() is leaky, because of the mem::forget().

for the wake_by_ref() function, on the other hand, it is probably ok. if you can avoid cloning the Arc<Task>, you don't need to use Arc::from_raw() just to forget() it later, you may just reborrow from raw pointer by dereferencing it.

if your implementation does require cloning the task Arc, it's better to use ManuallyDrop instead of forget(), just in case task.send() might panic.

that's the purpose of the waker of your executor, the vtable is the low level implementation of your waker.

it is recommended to use the Wake trait instead of RawWaker and let the compiler handle the vtables. but if you want to learn the underlying mechanism, you can read topics on virtual functions and dynamic dispatching in C++, you can find much more resources in C++ than the same topics in rust.

3

when you are given temporary access of the raw Arc data i think it's preferable to do

let data:*const ()={
        // 'recover' the Arc used to generate the pointer marking the start of the pseudo borrow
        let task: Arc<Task> = unsafe { Arc::from_raw(data as _) };
        
       // use it ......

       //turn it back into a pointer marking the end of your pseudo borrow
       task.into_raw();
}

shadowing the name of the original pointer is a nice plus to reduce the risk of accidentally double use it

4

Both this and the original forget() are either unsound or fragile. The problem is that if the code in the middle (task.send(task.clone());) panics, the Arc will be dropped, decrementing the reference count when it shouldn’t be and causing use-after-free later.

Using ManuallyDrop as @nerditation mentioned is a better option. In general, any time you find yourself using mem::forget in unsafe code to manipulate ownership, it is probably more correct to use ManuallyDrop.

    pub fn wake_by_ref(data: *const ()) {
        let task: ManuallyDrop<Arc<Task>> =
            ManuallyDrop::new(unsafe { Arc::from_raw(data.cast::<Task>())) };
        task.send(Arc::clone(&task));
    }
5

yeah, you're right. After you've pointed out there signature, that's definitely a leak.

    pub fn wake(data: *const ()) {
        let task: Arc<Task> = unsafe { Arc::from_raw(data as _) };
        // schedule the task on the executor again
        task.send(task.clone());
    }

does the first line (from_raw) of this wake runs the drop function that I've written or does it runs the destructor for Arc?

6

yeah, I see. But how would it lead to a use-after-free as it's just decrementing it and not incrementing it..?

7

If the strong count reaches zero, then Task is deallocated. If you later call wake or wake_by_ref and recreate Arc from the raw pointer to the deallocated Task (or even Arc control block) then when you dereference this Arc you have classic use-after-free.

8

I would implement them like this:

wake This method consumes ownership over one refcount. So it should call Arc::from_raw and then just send the arc.

let task: Arc<Task> = unsafe { Arc::from_raw(data as _) };
task.send(task); // no clone needed
// no mem::forget

wake_by_ref This method does not consume ownership over a refcount, so since you want to send off a cloned Arc, you must increment the refcount. The current code is correct, but I think it's cleaner to write like this:

Arc::increment_strong_count(data as *const Task);
let task: Arc<Task> = unsafe { Arc::from_raw(data as _) };
task.send(task);

or even

Arc::increment_strong_count(data as *const Task);
wake(data);

reusing the wake implementation.

9

While it would be elegant if this worked, isn't it going to "error: use of moved value" (regardless of the signature of send) since there are either two moves or one borrow and one move?

10

Hey, so I've been looking at how Waker::wake is written and it looks something like this:

  pub fn wake(self) {
        // The actual wakeup call is delegated through a virtual function call
        // to the implementation which is defined by the executor.

        // Don't call `drop` -- the waker will be consumed by `wake`.
        let this = ManuallyDrop::new(self);

        // SAFETY: This is safe because `Waker::from_raw` is the only way
        // to initialize `wake` and `data` requiring the user to acknowledge
        // that the contract of `RawWaker` is upheld.
        unsafe { (this.waker.vtable.wake)(this.waker.data) };
    }

when we call Waker::wake it consumes the Waker and the desctructor for Waker never runs. Also, it ends up calling the wake function which we have written for RawWakerVTable. In the RawWakerVTable::wake function we decremented the reference count using Arc::from_raw (as Arc<Task> will be dropped at the end of the scope...) and incremented it by cloning task (task.clone())

In all of this I've observed that the Waker never gets dropped and we cannot use it again...

Does all of this sound about right?

11

Semantically, the Waker is not dropped because it is being consumed instead. If its Drop::drop() were to run, the raw waker would be double-freed.

12

where does the RawWaker gets freed for the first time?

13

When the wake function is called. As the documentation for RawWakerVTable specifies:

The implementation of this function must make sure to release any resources that are associated with this instance of a RawWaker and associated task.

Every waker either gets vtable.wake or vtable.drop called (or is leaked), not both.

14 
    pub fn wake(data: *const ()) {
        let task: Arc<Task> = unsafe { Arc::from_raw(data as _) };
        // schedule the task on the executor again
        task.send(task.clone());
    }

so is our current implementation of RawWakerVTable::wake incorrect?

15

That looks correct.

16

but we are not releasing any resources associated with the RawWaker (apart from the ref count..)
Does the consumed Waker take care of that? if yes then could you please elaborate a bit like how's that actually happening..thanks!

17

A waker just owns a refcount on the Arc stored inside it ... and that's it. There are no other resources to manage.

18

turns out my clone implementation was incorrect as well :slight_smile:
this implementation works tho:

    // we have to retain all the resources, so we don't drop them?
    pub fn clone(data: *const ()) -> RawWaker {
        let task: ManuallyDrop<Arc<Task>> =
            ManuallyDrop::new(unsafe { Arc::from_raw(data.cast::<Task>()) });

        let _ = task.clone();

        RawWaker::new(data, &Executor::VTABLE)
    }

the very first implementation which I have posted above doesn't work...

documentation of RawWakerVTable::clone:
This function will be called when the RawWaker gets cloned, e.g. when
the Waker/LocalWaker in which the RawWaker is stored gets cloned.

The implementation of this function must retain all resources that are
required for this additional instance of a RawWaker and associated
task. Calling wake on the resulting RawWaker should result in a wakeup
of the same task that would have been awoken by the original RawWaker.