Potential memory leak

[abh1nav10]

We recently added the following changes in #901.

    if ::std::thread::panicking() {
        ::std::process::abort()
    }

It was added in Task::drop_future, the drop implementation for Task and Task::drop. The last one seems to be introducing a subtle condition where we might leak memory which I tried to reproduce in this way.

#[test]
#[should_panic]
fn check_drop() {
    use std::sync::mpsc;
    let (tx, rx) = mpsc::channel();
    let first = std::thread::spawn(move || {
        compio_runtime::Runtime::new().unwrap().block_on(async {
            // This future is never polled for the second time because the timer wheel will notify
            // us about the expiration of 1 second first and we will proceed to panic next!
            let handle = compio_runtime::spawn(async {
                sleep(Duration::from_secs(3)).await;
                10
            });
            tx.send(handle).unwrap();
            sleep(Duration::from_secs(1)).await;
            panic!("top future panics!");
        })
    });
    let second = std::thread::spawn(move || {
        compio_runtime::Runtime::new().unwrap().block_on(async {
            let handle = rx.recv().unwrap();
            compio_runtime::spawn(async {
                // output never arrives because the first runtime panics before the spawned future
                // gets to return the output and in that case it clears the executor which leads to
                // drop being called on the task inside of it. However, here since we had registered
                // our waker and the `Task::drop` does not decrement the refcount as it will
                // encounter the panicking check and return early, when this runtime gets dropped,
                // this task will leak memory because its refcount has not reached zero.
                let _output = handle.await;
            })
            .detach();
            sleep(Duration::from_secs(3)).await;
        })
    });
    second.join().unwrap();
    // joining this will panic!
    first.join().unwrap();
} 

This may leak the underlying memory because when the first runtime is clearing the executor, it is calling Task::drop for the task that we spawned there which immediately returns because the check for whether we are panicking or not returns true. Since we return early, the refcount does not get decremented for the waker that we had registered and that leads to the memory for the task on the second runtime from never being deallocated.

[abh1nav10]

I think the fix would be to change Task::drop to

pub unsafe fn drop(&self) {
        instrument!(compio_log::Level::TRACE, "Task::drop", id = ?self.header().id);

        let header = self.header();
        debug_assert!(
            header.tracker.valid(),
            "drop should only be called by Executor"
        );
        let state = header.state.set_dropped();

        header.shared.store(ptr::null_mut(), Release);

        if !state.is_completed() {
            trace!("Dropping future");
            // The task has not completed yet, drop future.
            // `drop_future` returns early if we enter it while panicking
            // and if not so, it ensures that we abort if dropping the
            // future or output panics.
            unsafe { (header.vtable.drop_future)(self.0, false) };
        }

        // If `JoinHandle` is setting the waker remotely, it'll check the state
        // afterwards and drop the waker. Otherwise, we drop the waker here if
        // it exists.
        if state.has_waker() && !state.is_setting_waker() {
            trace!("Dropping waker");
            // If the waker here points to `Arc<Notify>`, then dropping it is just
            // an `Arc` refcount decrement and if the caller spawned another task just
            // to poll the `JoinHandle`(should rarely happen), dropping the waker here
            // calls `Task::drop` which has got a panic guard. So, we dont use a panic
            // guard here. This also ensures that the refcount of this waker is decremented
            // appropriately which would not happen if we return early from this function if we
            // enter into it while panicking.
            header
                .waker
                .with_mut(|ptr| unsafe { drop_in_place(ptr.cast::<Waker>()) });
        }

        trace!("Completed");

[Berrysoft]

It's not correct to remove the panic guard according to your comments. We should not assume what does the waker do.

Leaking is expected. But, if you think that it is essential to drop the wakers, it's OK to try removing the short return.

[abh1nav10]

We should not assume what does the waker do.

Not proposing to remove the guard, but just to understand the execution completely, isn't it true that the waker that is dropped inside of Task::drop can either be the waker that gets propagated from the top level future if we await the JoinHandle and in that case, it was created from Waker::from(t) where t is Arc<Notify>. In that case the RawWakerVtable constructed by the standard library just manages the Arc ref-count on the functions used in the VTABLE and calls for the methods on the Wake trait which we have implemented for Notify. On both IO-URING driver, we use an OwnedFd and on the POLL driver, Poller inside of Notify and I don't think that those types can cause a panic when dropped.
If the joinhandle is not directly awaited and polled by a task spawn, in that case, when the Waker is dropped and ref-count is 0, we are executing our own code like so

unsafe fn drop_waker(ptr: *const ()) {
        drop(unsafe { Task::from_raw(ptr) });
    }

which will call the drop implementation of Task which has already got the panic guards in place.

What I am having a hard time understanding is that where exactly is the breaking point in these scenarios that can cause a panic and for which we place a guard before dropping the Waker in Task::drop.

I am definitely missing something here. Completely from a curious standpoint :) @Berrysoft

[Berrysoft]

The waker is passed from an external Context, e.g., by calling poll of JoinHandle. Users might call this method with any kind of waker, and of course it is safe to make it panic.

[abh1nav10]

Missed that they can call poll directly as well and need not just rely on the machinery always. Sorry for that!