The asynchronous drop glue generation design

This text aims to explain the design of my asynchronous drop prototype, which I have been working on for some time now.

Public interface of AsyncDrop

I've tried to make interface of asynchronous drops as similar to the synchronous drops as possible. Take a look at the definition of the most important public trait of my prototype (AsyncDrop trait):

/// Custom code within the asynchronous destructor.
#[lang = "async_drop"]
pub trait AsyncDrop {
    /// A future returned by the [`AsyncDrop::async_drop`] to be part
    /// of the async destructor.
    type Dropper<'a>: Future<Output = ()>
    where
        Self: 'a;

    /// Constructs the asynchronous destructor for this type.
    fn async_drop(self: Pin<&mut Self>) -> Self::Dropper<'_>;
}

Given that we don't have async as a keyword generic I've had to define the entire new trait. It's kinda similar to AsyncFnMut as that trait also mirrors FnMut. Both of these async traits use near to the desugared interface of async functions, returning from sync method a future object of trait's associated type. I've also wrapped mutable reference to self into Pin just to be sure, maybe it'll become useful or detrimental.

Let's imagine its implementation for a new, cancellable during drop task handle in tokio:

impl<T> AsyncDrop for CancellableJoinHandle<T> {
    type Dropper<'a>: impl Future<Output = ()>;

    fn async_drop(self: Pin<&mut Self>) -> Self::Dropper<'_> {
        async move {
            self.join_handle.abort();
            let _ = Pin::new(&mut self.join_handle).await;
        }
    }
}

Here we are wrapping tokio::task::JoinHandle and using JoinHandle::abort to cancel the task if possible, and then awaiting its end. The impl_trait_in_assoc_type feature is used there to not implement futures manually, perhaps this can be simplified further with return-position impl Trait and async methods in traits.

Choosing against poll_drop

You may wonder about possible alternative design of async drop, usually named poll_drop:

#[lang = "async_drop"]
pub trait AsyncDrop {
    fn poll_drop(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()>;
}

We have decided against it since it would require to embed the state of asynchronous destruction into the type itself. For example Vec<T> would need to store an additional index to know which element is currently in the process of asynchronous destruction (unless we poll_drop every element on each parent call, but I imagine that could become expensive quick, and it is not exactly symmetrical to how the regular Drop functions). Also each element of the vector would require additional space for these embedded asynchronous destructors, even tho it would be utilized one at a time.

However there is indeed one benefit of poll_drop which I hypothesized to be a supplemental interface down below.

Asynchronous drop glue

To run async drop glue on a type we can use public async_drop or async_drop_in_place functions, just as with the regular variant of drop. These are the async implementations:

pub async fn async_drop<T>(to_drop: T) {
    let to_drop = MaybeUninit::new(to_drop);
    // SAFETY: we store the object within the `to_drop` variable
    unsafe { async_drop_in_place(to_drop.as_mut_ptr()).await };
}

#[lang = "async_drop_in_place"]
pub unsafe fn async_drop_in_place<T: ?Sized>(
    to_drop: *mut T,
) -> <T as AsyncDestruct>::AsyncDestructor {
    // ...
}

I assume you understand how async_drop function works. However the hard part lies with async_drop_in_place. It is not an async function but merely returns an object of AsyncDestruct::AsyncDestructor type, presumably a future. You can also notice we don't have syntax T: AsyncDestruct. Let's take a closer look of AsyncDestruct trait and its associated type:

#[lang = "async_destruct"]
trait AsyncDestruct {
    type AsyncDestructor: Future<Output = ()>;
}

This trait is internal to the compiler. The AsyncDestructor is actually a future for async drop glue, the code deinitializing the Self object. It is implemented for every type, thus it does not require trait bounds to use on any type. Compiler implements it the same way as the also internal DiscriminantKind trait. Now I should mention that async_drop_in_place's body is also generated by the compiler, but this time it's the same way drop_in_place is generated (via shim).

But what type should we assign to AsyncDestructor? async_drop_in_place simply creates that asynchronous destructor future and does not execute it. I haven't yet found a way to generate coroutines solely from the compiler, but I was given the advice to compose core library types to create such futures. I've defined various future combinators to chain, defer futures or to choose either of two futures and by combining them I've implemented asynchronous destructors for ADTs and other types. Although some code couldn't have been offloaded to the core (I think). For example I've had to precompute a pointer to each field ahead of time inside of the async_drop method.

#[lang = "async_drop_chain"]
async fn chain<F, G>(first: F, second: G)
where
    F: IntoFuture<Output = ()>,
    G: IntoFuture<Output = ()>,
{
    first.await;
    second.await;
}

#[lang = "async_drop_either"]
async unsafe fn either<O: IntoFuture<Output = ()>, M: IntoFuture<Output = ()>, T>(
    other: O,
    matched: M,
    this: *mut T,
    discriminant: <T as DiscriminantKind>::Discriminant,
) {
    if unsafe { discriminant_value(&*this) } == discriminant {
        drop(other);
        matched.await
    } else {
        drop(matched);
        other.await
    }
}

#[lang = "async_drop_defer"]
async unsafe fn defer<T: ?Sized>(to_drop: *mut T) {
    unsafe { async_drop_in_place(to_drop) }.await
}

Since async drop glue could hypothetically in future be executed automatically within the cleanup branches used for unwind, one property I believe AsyncDestructor future should have is that instead of panicking it must simply return Poll::Ready(()) on every poll after future completes. I've called this property future idempotency since it makes sense and have a special fuse combinator wrap around any regular future to have such guarantee.

As of right now (2024-03-29) async drop glue for coroutines (async blocks) and dynamic types (dyn Trait) are not implemented. Coroutines have their special code for generating even regular drop glue, extracting a coroutine_drop branch from coroutine's MIR. Other person works on it. For dynamic types support I have a hypothetical design which I'll describe below. Automatic async drops at the end of the scope aren't implemented too.

Combinator table

You might ask if we even need Noop combinator and can't we not instantiate async destructor for trivially destructible types? But no, this is not possible, since user may call async_drop_in_place on any type, which has to return some future type.

See current implementations of these combinators inside of the library/core/src/future/async_drop.rs.

Visibility problem

If you compare public interface for interacting with value discriminants within the core library with interface described here, you could notice usage of trait's associated type instead of a generic type. Actually directly using this associated type may be problematic as it can possibly leak its special trait and method implementations. Also I believe it would be better to keep AsyncDestruct trait private. At last it perhaps it would be more convenient to use a generic type instead as with Discriminant.

To solve this problem the only way right now would be to define a wrapper struct AsyncDropInPlace<T> around it and forward its Future implementation to the actual async destructor of type T. We would also have a new wrapper function async_drop_in_place to return that wrapper type and would rename compiler generated function which held this name previously into async_drop_in_place_raw.

However, this AsyncDropInPlace could still leak some details of stored inner value, such as any auto trait implementation and a drop check. These can be either left as is (current behavior) or be suppressed with PhantomData<*mut ()> field and with a noop Drop implementation on it. Not sure which one should be chosen.

Generation of async_drop_in_place_raw

The body of async_drop_in_place_raw function is generated by the compiler within the compiler/rustc_mir_transform/src/shim.rs. AsyncDestructorCtorShimBuilder is the core structure of for generating code in form of MIR. Let's take a look at what kind of code is being generated for enum:

chain(
    surface_async_drop_in_place(to_drop),
    either(
        chain(
            async_drop_in_place_raw((*(to_drop as *mut T::Variant1)).field0),
            async_drop_in_place_raw((*(to_drop as *mut T::Variant1)).field1),
        ),
        chain(
            async_drop_in_place_raw((*(to_drop as *mut T::Variant0)).field0),
            async_drop_in_place_raw((*(to_drop as *mut T::Variant0)).field1),
        ),
        to_drop,
        variant0_discriminant,
    ),
)

As you can see it can see it is simply an expression. We can express execution of a single expression with a stack machine, which is actually exactly how AsyncDestructorCtorShimBuilder functions. It stores a stack of operands which are either a moved local, a copied local or a const value (like a discriminant). We allocate and deallocate storage for moved locals on push and pop to the builder's stack. We can assign a value to a local, putting it at the top of the stack or combine operands (but first we pop them) with a function call to put a combinator value at the top of the stack too. Order of arguments for a function call can be summarized as top operand of the stack being the last argument. Then we return the one last stack operand.

This stack machine also allows us to easily create a cleanup branch to drop operands during unwind without redundant drops by reusing drops for stored locals on the stack, forming a kind of tree inside of the MIR control-flow graph.

What's next?

ADT async destructor types

As I've said those future combinators are just a patchwork for current inability to generate ADT futures on the fly. Defining such components inside of the core is fine in some cases, like for async destructor of slice. But for ADTs, tuples, closures the proper solution would be to define the new type kind named something like AdtAsyncDestructor. Given one of those types we could generate a consistent state for the async destructor and then generate its Future::poll function. This way we won't need to calculate and store all the pointers to each field ahead of time.

Ideas for the future

Should async_drop_in_place work with references?

Since async_drop_in_place returns an async destructor future what should reference the dropped object, perhaps it would be more beneficial to have async_drop_in_place use reference &mut ManuallyDrop<T> instead. It would be less unsafe and we won't have to deal with pointers infecting async destructor types with !Send and !Sync.

Async drop glue for dyn Trait

The problem with dynamic types is basically about loosing the type information. We cannot know <dyn Trait as AsyncDestruct>::AsyncDestructor type's size and alignment, thus we cannot know how much stack or coroutine's local space we should allocate for the storage. One approach would be to have type AsyncDestructor = Box<dyn Future> for dynamic types, which could be not ideal. But actually before we coerce static types into dynamic, perhaps we could have a wrapper type which contains space both for T and for <AsyncDestruct as T>::AsyncDestructor?

#[lang = "PollDestruct"]
trait PollDestruct {
    fn poll_drop(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()>;
}

struct PollDrop<T> {
    async_dtor: Option<<T as AsyncDestruct>::AsyncDestructor>,
    value: MaybeUninit<T>,
    _pinned: PhantomPinned,
}

impl<T> PollDestruct for PollDrop<T> {
    fn poll_drop(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
        unsafe {
            let this = self.get_unchecked_mut();
            let dtor = this
                .async_dtor
                .get_or_insert_with(|| async_drop_in_place(this.value.as_mut_ptr()));
            Pin::new_unchecked(dtor).poll(cx)
        }
    }
}

// Have a `PollDrop<Fut> -> Box<dyn Future<Output = ()> + PollDestruct>`

And like that we embedded enough space and type information to unsize these types and work with them, while still being able to be asynchronously destroyed.

Exclusively async drop

It's almost pointless to implement AsyncDrop on your type while it is perfectly valid to synchronously drop your type. There can be a way to restrict sync drops of a type by implementing !Destruct for a type. Compiler should emit a compiler error wherever it tries to synchronously drop a ?Destruct value. It would be fine to asynchronously drop them, which would be done (semi)automatically inside of async code.

While this approach as far as I know preserves backwards compatibility, it would require users to manually add support for T: ?Destruct types inside of their code, which is the reason new ?Trait bounds are considered to be unergonomic by many rustc lead developers. Perhaps it would be fine to have T: Destruct by default for synchronous functions and T: ?Destruct by default for asynchronous ones in the next edition?

But my mentor suggests to try out a different approach: emitting such errors after monomorphization of a generic function, perhaps as a temporary measure before a proper type-level solution is enabled. It does sound like how C++ templates work which come with some issues on their own. But rust already allows post-monomorphization errors like linker errors.

Automatic async drop and implicit cancellation points

The core feature of the hypothetical async drop mechanism is considered to be automatic async cleanup, which requires to add implicit await points inside of the async code wherever it destroys an object with async drop implementation. Currently every await point also creates a cancellation point where future can be cancelled with drop if it is suspended there.

Implicit cancellation point within the async code would probably make it much more difficult to maintain cancellation safety of your async code because of not seeing where exactly your async code can suspend. The simplest solution for this would be to have implicit await point to not generate a cancellation point. This is possible if such async block implements !Destruct (see above) and can only be asynchronously dropped. Then if user starts async drop of that future while it is suspended on implicit await point, the future will continue as usual until it either returns or suspends on explicit await point. User will have to explicitly call and await async_drop to allow cancellation during suspension.

Drop of async destructors

How should drop of an async destructor should function? I see the simplest solution would probably be that async drop of async destructor will simply continue execution of async destructor.

Conclusion

There are still a lot of questions to be answered, but it's important to not put our hands down.

Also I would like to mention this text is based on similar works of many other people, references to which you can find in this MCP: Low level components for async drop.