slotpoller 0.2.1

/*
 * Copyright © 2026 Anand Beh
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//!
//! This library implements a collection of pollable futures via a backing array. It has the
//! following properties:
//! 1. Futures can be stored on the stack.
//! 2. The number of futures is bounded.
//! 3. Futures are only polled when they generate wake-ups.
//!
//! The API reflects these properties. There are three main structs, [StackSlots], [HeapSlots],
//! and [SlotPoller].
//!
//! ### Stack slots
//!
//! Using [StackSlots] requires you to pin it:
//!
//! ```
//! use std::pin::pin;
//! use slotpoller::{SlotPoller, StackSlots};
//!
//! let stack_slots = pin!(StackSlots::<30, _>::default());
//! let mut slot_poller = SlotPoller::new(stack_slots);
//! # slot_poller.try_push(async {});
//! ```
//!
//! ### Heap slots
//!
//! [HeapSlots] lets you set the capacity at runtime. A capacity must be provided:
//!
//! ```
//! use std::num::NonZeroU16;
//! use slotpoller::{HeapSlots, SlotPoller};
//!
//! let capacity = NonZeroU16::new(50).unwrap();
//! let mut slot_poller = SlotPoller::new(HeapSlots::with_capacity(capacity));
//! # slot_poller.try_push(async {});
//! ```
//!
//! ### Basic Usage
//!
//! After construction, most methods can be found on the [SlotPoller] struct.
//!
//! ```
//! use std::io;
//! use std::pin::pin;
//! use slotpoller::{StackSlots, SlotPoller};
//!
//! async fn exec() -> Result<String, io::Error> {
//!   Ok(String::new())
//! }
//!
//! async fn exec_multiple() -> Result<(), io::Error> {
//!   let stack_slots = pin!(StackSlots::<20, _>::default());
//!   let mut slot_poller = SlotPoller::new(stack_slots);
//!
//!   for _ in 0..50 {
//!     let (res, vacancy) = slot_poller.next_vacancy().await;
//!     if let Some(Err(e)) = res {
//!       return Err(e);
//!     }
//!     vacancy.insert(exec());
//!   }
//!
//!   slot_poller.try_drain(|res| { res?; Ok(()) }).await
//! }
//! ```
//!
//! ### Allocation
//!
//! The allocation footprint of this crate is very minimal and stays constant over the lifetime of
//! the poller. Here are the allocating functions in this crate:
//!
//! * [HeapSlots::with_capacity] (2 allocations)
//! * [StackSlots::new] (1 allocation)
//!
//! By design, the Rust async model requires either heap allocation or static mutable memory in
//! order to implement wakers. Users who want an alloc-free solution (using static mutable memory)
//! should look at the `embassy-executor` crate.
//!
//! ### Panics
//!
//! Functions in this library should never generate a panic, with one exception: polling a future after it
//! has completed (see [Future::poll]).
//!

mod guard;
mod memory;
mod polling;
#[cfg(test)]
mod tests;

use crate::guard::{AtomicStateNodeIdx, OptSlotIdx, PollThread, SlotIdx, StateNodeIdx};
use crate::memory::MemoryBacking;
use crate::polling::{AtomicStatus, poll_queue_reinsert_head, pop_pollable};
use cache_padded::CachePadded;
use diatomic_waker::DiatomicWaker;
use std::alloc::Layout;
use std::fmt::{Debug, Formatter};
use std::marker::PhantomData;
use std::mem::MaybeUninit;
use std::ops::Deref;
use std::pin::Pin;
use std::ptr::NonNull;
use std::task::{Context, Poll};
use triomphe::{Arc, HeaderSlice};

///
/// An internal trait, used to identify different memory providers.
///
/// This is what allows the library to support both stack-based polling of a fixed size array,
/// and polling of a dynamically allocated array, at the same time.
///
#[allow(private_bounds)]
pub trait SlotMemory<F: Future>: memory::Memory<F> {}

///
/// Stack-friendly memory that can be used by the poller.
///
/// The generic parameter `N` determines the size of the array. Its value must **not** exceed
/// [u16::MAX], and the implementation checks this requirement. In the future, if the feature
/// `generic_const_exprs` is added to Rust, the type may change to `u16`.
///
/// Usage is simple. [SlotPoller] takes a pinned reference to this struct:
///
/// ```
/// use std::io;
/// use std::pin::pin;
/// use slotpoller::{StackSlots, SlotPoller};
///
/// async fn exec() -> Result<(), io::Error> {
///   Ok(())
/// }
///
/// async fn exec_multiple() {
///   let stack_slots = pin!(StackSlots::<20, _>::default());
///   let mut stack_poller = SlotPoller::new(stack_slots);
///
///   let (_, vacancy) = stack_poller.next_vacancy().await;
///   vacancy.insert(exec());
/// }
/// ```
///
pub struct StackSlots<const N: usize, F: Future> {
    activity: EngineActivity,
    slots: [Slot<F>; N],
    shared_storage: Arc<SharedStorage>,
}

///
/// An intermediate struct that erases the constant parameter on [StackSlots].
///
/// This generic erasure prevents excessive binary size due to runaway monomorphization.
///
// Not exposed: Does this have value to API users? Probably not
struct StackSlotsRef<'pin, F: Future> {
    activity: &'pin mut EngineActivity,
    slots: Pin<&'pin mut [Slot<F>]>,
    shared_storage: &'pin Arc<SharedStorage>,
}

impl<'pin, const N: usize, F: Future> From<Pin<&'pin mut StackSlots<N, F>>>
    for StackSlotsRef<'pin, F>
{
    fn from(memory: Pin<&'pin mut StackSlots<N, F>>) -> Self {
        unsafe {
            // SAFETY
            // Slots stays pinned - we project it, while other fields are Unpin
            let memory = memory.get_unchecked_mut();
            Self {
                activity: &mut memory.activity,
                slots: Pin::new_unchecked(&mut memory.slots),
                shared_storage: &memory.shared_storage,
            }
        }
    }
}

impl<'pin, const N: usize, F: Future> SlotMemory<F> for Pin<&'pin mut StackSlots<N, F>> {}

///
/// Allocates the space for the futures on the heap.
///
/// The advantage of this is that you can choose the buffer size at runtime.
/// This approach is very similar to other crates that implement collections of futures.
///
pub struct HeapSlots<F: Future> {
    activity: EngineActivity,
    // Don't use Box. We must free pointer manually, AFTER futures have been dropped
    slots_ptr: NonNull<[Slot<F>]>,
    shared_storage: Arc<SharedStorage>,
    _marker: PhantomData<[Slot<F>]>,
}

impl<F: Future> SlotMemory<F> for HeapSlots<F> {}

impl<F: Future> Unpin for HeapSlots<F> {}

#[derive(Debug)]
struct EngineActivity {
    /// A LIFO stack of empty slots.
    empty_head: OptSlotIdx,
    /// The number of slots with futures inside of them
    slots_active: u16,
    /// The head of the pollable queue is updated only by the polling thread
    poll_queue_head: StateNodeIdx,
    /// Fairness counter for the current polling loop
    poll_loop_idx: u16,
}

struct EngineFields<'m, F> {
    activity: &'m mut EngineActivity,
    slots: Pin<&'m mut [Slot<F>]>,
    shared_storage: &'m Arc<SharedStorage>,
}

/// The fixed length part of the shared storage
///
/// It is updated by everyone (poller and wakers), so it is wrapped in CachePadded for efficiency
#[derive(Debug)]
struct SharedStorageHeaderInner {
    /// The tail of the pollable linked queue
    poll_queue_tail: AtomicStateNodeIdx,
    /// Used to recover the full DST pointer
    nodes_len: u32,
    /// The waker of the poller's latest operations
    main_waker: DiatomicWaker,
}

#[derive(Debug)]
#[repr(transparent)]
struct SharedStorageHeader(CachePadded<SharedStorageHeaderInner>);

impl SharedStorageHeader {
    // HeaderSlice is repr(C), and creating a [T; 0] from a [T; N] is safe
    const BYTE_OFFSET: usize = Layout::new::<HeaderSlice<Self, [StateNode; 0]>>().size();
}

impl Deref for SharedStorageHeader {
    type Target = SharedStorageHeaderInner;

    #[inline]
    fn deref(&self) -> &Self::Target {
        &*self.0
    }
}

type SharedStorage = HeaderSlice<SharedStorageHeader, [StateNode]>;

///
/// The main poller
///
/// It's called a [SlotPoller] because it can handle either stack or heap memory,
/// and because it polls multiple futures.
///
pub struct SlotPoller<F: Future, M: SlotMemory<F>> {
    token: PollThread,
    memory: M::Backing,
}

impl<F: Future, M: SlotMemory<F>> Debug for SlotPoller<F, M> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SlotPoller")
            .field("token", &self.token)
            .field("memory", &self.memory)
            .finish()
    }
}

impl<F: Future, M: SlotMemory<F>> SlotPoller<F, M> {
    ///
    /// Creates the poller.
    ///
    /// After calling this function, the memory is permanently tied to this poller.
    /// This poller can't be moved to other threads, either.
    ///
    pub fn new(memory: M) -> Self {
        Self {
            token: memory.poll_thread(),
            memory: memory.into_backing(),
        }
    }
}

/*

The lifecycle of a slot is managed through the struggle over its activity and contents.

The poll thread stores a &mut StackPoller and manages the on-stack slots and linked list of empty
slots. The waker can only access the shared heap storage, which is manipulated through atomic
operations. The pollable slot queue is shared between them using a control-exchanging algorithm.

A slot is empty if it has no future (status = Status::Uninit). Empty slots are tracked by a linked
list, in stack memory, that is updated by the polling thread. We call this the "empty slot stack"
since it exhibits LIFO behavior.

We can describe the lifecycle and its state transitions with this diagram:

     Status    | Empty/Pollable queue  |      Poll thread       |     Waker thread      |
 ~~~~~~~~~~~~~ | ~~~~~~~~~~~~~~~~~~~~~ | ~~~~~~~~~~~~~~~~~~~~~~ | ~~~~~~~~~~~~~~~~~~~~~ |
     Uninit    |  Queued as empty
       |               |
       |              \|/
       |              *** <--------------- Find a new slot
       |         Dequeued as empty              ||
      \|/                                       ||
      *** <------------------------------- Future inserted
      New                               (defaults to unpolled)
       |
       |
       |
       |
      \|/
      *** <------------------------------- Poll the future
    Waiting                                Create the waker ------\
       |               .                                           --\
       |               .                                              --\
       |               .                                                 |
       |               |                                                \|/
      \|/              |
      *** <------------|----------------------------------------- 1. CAS to Woken
     Woken             |
       |              \|/ <-------------------------------------- 2. If CAS succeeded
       |              ***                                            A. enqueue as pollable
       |         Enqueued as pollable                                B. call original waker
       |               |
       |              \|/
       |              *** <--------------- Withdrawn from queue
       |         Dequeued as pollable            ||
      \|/              |                         ||
      *** <------------|------------------ Poll the future again
     Waiting           |
    /       \          .                   If incomplete, recreate
   /         \         .                   waker and repeat. Else,
   |         |         .                   future is completed.
   |         |
 Pending   Ready
 (repeat)   ***  <------------------------- Future dropped
   |         |                                   ||
   .         |                                   ||
   .       Uninit      |                         ||
   .                  \|/                        ||
                      *** <---------------- Mark slot empty
                 Enqueued as empty
                 (diagram repeats)

 */

struct Slot<F> {
    /// Only initialized if state is Waiting or Woken
    /// If initialized, then the future has not panicked and has not yet completed
    future: MaybeUninit<F>,
    empty_link: OptSlotIdx,
    /// The index of the loop we were in, when polled. Used to ensure fairness
    last_poll_loop_idx: u16,
}

// Manual pin_project impl

impl<F> Slot<F> {
    fn project(self: Pin<&mut Self>) -> SlotProject<'_, F> {
        unsafe {
            // SAFETY
            // Future stays pinned
            let this = self.get_unchecked_mut();
            SlotProject {
                future: Pin::new_unchecked(&mut this.future),
                empty_link: &mut this.empty_link,
                last_poll_loop_idx: &mut this.last_poll_loop_idx,
            }
        }
    }
}

struct SlotProject<'s, F> {
    future: Pin<&'s mut MaybeUninit<F>>,
    empty_link: &'s mut OptSlotIdx,
    last_poll_loop_idx: &'s mut u16,
}

impl<F> Debug for Slot<F> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        struct Opaque;
        impl Debug for Opaque {
            fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
                f.write_str("opaque")
            }
        }
        f.debug_struct("Slot")
            .field("future", &Opaque)
            .field("empty_link", &self.empty_link)
            .field("last_poll_loop_idx", &self.last_poll_loop_idx)
            .finish()
    }
}

#[derive(Debug)]
struct StateNode {
    status: AtomicStatus,
    /// The next link in the queue. [StateNodeIdx::UNSET] for none.
    poll_queue_link: AtomicStateNodeIdx,
    /// The corresponding slot index.
    /// This value is recovered by the waker in order to reconstruct SharedStorage
    slot_idx: SlotIdx,
}

impl EngineActivity {
    fn empty_stack_pop<'s, F>(
        &mut self,
        slots: Pin<&'s mut [Slot<F>]>,
    ) -> Option<(SlotIdx, Pin<&'s mut Slot<F>>)> {
        let head = self.empty_head;
        if head.is_set() {
            let head = head.into_slot_idx();
            let head_slot = head.get_slot(slots);
            let new_head = head_slot.empty_link;
            self.empty_head = new_head;
            self.slots_active += 1;
            Some((head, head_slot))
        } else {
            None
        }
    }

    fn empty_stack_push<F>(&mut self, slot_idx: SlotIdx, slot: Pin<&mut Slot<F>>) {
        let prev_head = self.empty_head;
        *slot.project().empty_link = prev_head;
        self.empty_head = OptSlotIdx::from(slot_idx);
        self.slots_active -= 1;
    }

    #[inline]
    fn empty_stack_has_any(&self) -> bool {
        self.empty_head.is_set()
    }
}

/**

Dequeues pollable slots, polls them, and handles the result of that poll.

Includes cooperative yielding when a future is polled more than once, in the same loop.

 */
macro_rules! poll_loop {
    ($cx:ident, $token:ident, $activity:ident, $slots:ident, $shared_storage:ident, $slot:ident, $slot_idx:ident, $state_node:ident, $poll_res:ident, $if_none:expr, $on_poll:expr) => {{
        let poll_loop_idx = $activity.poll_loop_idx;
        $activity.poll_loop_idx = poll_loop_idx.wrapping_add(1);

        // Store the waker. If woken, we'll return back to this line
        unsafe {
            // SAFETY
            // register/unregister/wait_until can't be called concurrently
            // We're on the poll thread, and this is the only place we use this function
            $shared_storage.header.main_waker.register($cx.waker());
        }
        loop {
            let pop_pollable = pop_pollable($token, $activity, $shared_storage);
            let ($slot_idx, pollable) = match pop_pollable {
                // None found
                (OptSlotIdx::UNSET, _) => break $if_none,
                // Try again
                (_slot_idx, None) => continue,
                // Pollable slot
                (slot_idx, Some(pollable)) => (slot_idx.into_slot_idx(), pollable),
            };
            let mut $slot = $slot_idx.get_slot($slots.as_mut());
            let $state_node = pollable.state_node();

            if $slot.last_poll_loop_idx == poll_loop_idx {
                // Polled the same slot in this loop previously
                // To be fair (and maintain cooperative scheduling), yield once here
                // https://github.com/rust-lang/futures-rs/issues/2053

                // First, put the slot back on the queue
                unsafe {
                    // SAFETY
                    // 1. This node is not already enqueued
                    // 2. The state node and index are corresponding
                    // 3. It keeps the same status as it had previously (Init or Woken)
                    poll_queue_reinsert_head(
                        $token,
                        $activity,
                        $state_node,
                        StateNodeIdx::from($slot_idx),
                    );
                }
                // Now yield
                $cx.waker().wake_by_ref();
                break Poll::Pending;
            }
            // Ready to poll
            *$slot.as_mut().project().last_poll_loop_idx = poll_loop_idx;
            let $poll_res = pollable.call_poll($token, $slot.as_mut());
            $on_poll
        }
    }};
}

//
// Public API
//

impl<F, M> SlotPoller<F, M>
where
    F: Future,
    M: SlotMemory<F>,
{
    ///
    /// Tries to add another future to the stack. If there is no space, returns it back.
    ///
    pub fn try_push(&mut self, future: F) -> Result<(), TryPushErr<F>> {
        let EngineFields {
            activity,
            slots,
            shared_storage,
        } = self.memory.fields();
        if let Some((slot_idx, slot)) = activity.empty_stack_pop(slots) {
            let shared_storage = shared_storage as &SharedStorage;
            let state_node = slot_idx.get_state_node(shared_storage);
            unsafe {
                // SAFETY
                // slot is empty, and slot/slot_idx/state_node correspond
                slot.init_future(self.token, slot_idx, state_node, shared_storage, future);
            }
            Ok(())
        } else {
            Err(TryPushErr(future))
        }
    }

    ///
    /// Waits for the next vacant slot in the stack poller.
    ///
    /// The future created by this method returns a tuple. The first value is the output of any
    /// future that had to be completed, in order to find a free slot. The second value is the
    /// vacancy which is used to insert the future.
    ///
    #[inline]
    pub fn next_vacancy(&mut self) -> NextVacancyFuture<'_, F, M> {
        NextVacancyFuture {
            token: self.token,
            poller: Some(self),
        }
    }

    ///
    /// Waits for the next future to complete, if there are any futures.
    ///
    /// If there are no active futures, returns `None`. Otherwise, returns a future which polls
    /// tasks until one of them completes.
    ///
    #[inline]
    pub fn next_completion(&mut self) -> Option<NextCompletionFuture<'_, F, M>> {
        let fields = self.memory.fields();
        if fields.activity.slots_active == 0 {
            None
        } else {
            Some(NextCompletionFuture {
                token: self.token,
                poller: self,
            })
        }
    }

    ///
    /// Drains this poller, handling results using the function.
    ///
    /// This will poll all futures in the poller. When one future completes, it is handled using
    /// the given function.
    ///
    #[inline]
    pub fn drain<D>(&mut self, drain_function: D) -> DrainFuture<'_, F, M, D>
    where
        D: FnMut(F::Output),
    {
        DrainFuture {
            token: self.token,
            poller: self,
            drain_function: AssertUnpin(drain_function),
        }
    }

    ///
    /// Drains this poller, handling results using a fallible function.
    ///
    /// This is similar to [Self::drain]. However, if the function returns an error, that error
    /// is returned and the draining stops.
    ///
    #[inline]
    pub fn try_drain<D, E>(&mut self, drain_function: D) -> TryDrainFuture<'_, F, M, D, E>
    where
        D: FnMut(F::Output) -> Result<(), E>,
    {
        TryDrainFuture {
            token: self.token,
            poller: self,
            drain_function: AssertUnpin(drain_function),
        }
    }
}

//
// Pushing
//

///
/// The error produced from [SlotPoller::try_push], containing the future which wasn't pushed.
///
/// Because standard library methods like expect, unwrap, etc. require a [Debug] implementation,
/// this type provides it irrespective of the future type.
///
pub struct TryPushErr<F>(pub F);

impl<F> Debug for TryPushErr<F> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.write_str("Failed to push future: out of space")
    }
}

///
/// An open slot inside the stack poller.
///
/// Dropping this struct is harmless, and will let the slot inside the poller stay empty.
///
pub struct Vacancy<'p, F: Future, M: SlotMemory<F>> {
    token: PollThread,
    poller: &'p mut SlotPoller<F, M>,
}

impl<'p, F, M> Vacancy<'p, F, M>
where
    F: Future,
    M: SlotMemory<F>,
{
    ///
    /// Inserts a future into this vacant slot, consuming the vacancy
    ///
    pub fn insert(self, future: F) {
        let EngineFields {
            activity,
            slots,
            shared_storage,
        } = self.poller.memory.fields();
        let (slot_idx, slot) = activity.empty_stack_pop(slots).expect("Missing empty slot");
        let shared_storage = shared_storage as &SharedStorage;
        let state_node = slot_idx.get_state_node(shared_storage);
        unsafe {
            // SAFETY
            // slot is empty, and slot/slot_idx/state_node correspond
            slot.init_future(self.token, slot_idx, state_node, shared_storage, future);
        }
    }
}

///
/// The future returned from [SlotPoller::next_vacancy]
///
#[must_use = "futures do nothing unless polled"]
pub struct NextVacancyFuture<'p, F: Future, M: SlotMemory<F>> {
    token: PollThread,
    /// None if this future finished
    poller: Option<&'p mut SlotPoller<F, M>>,
}

impl<'p, F, M> Future for NextVacancyFuture<'p, F, M>
where
    F: Future,
    M: SlotMemory<F>,
{
    type Output = (Option<F::Output>, Vacancy<'p, F, M>);

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let token = self.token;
        let EngineFields {
            activity,
            mut slots,
            shared_storage,
        } = self
            .poller
            .as_mut()
            .expect("Future polled after completion")
            .memory
            .fields();
        if activity.empty_stack_has_any() {
            return Poll::Ready((
                None,
                Vacancy {
                    token,
                    poller: self.poller.take().unwrap(),
                },
            ));
        }
        poll_loop!(
            cx,
            token,
            activity,
            slots,
            shared_storage,
            slot,
            slot_idx,
            state_node,
            poll_res,
            Poll::Pending, // nothing to poll => keep waiting
            {
                if let Poll::Ready((res, droppable)) = poll_res {
                    // Future completed - mark as empty / drop it
                    activity.empty_stack_push(slot_idx, slot.as_mut());
                    droppable.drop_future(slot, state_node);

                    return Poll::Ready((
                        Some(res),
                        Vacancy {
                            token,
                            poller: self.poller.take().unwrap(),
                        },
                    ));
                }
            }
        )
    }
}

//
// Streaming
//

///
/// The future returned from [SlotPoller::next_completion].
///
/// This future will poll all the futures inside the slot poller until one of them returns a ready
/// result. Its existence as a struct statically guarantees that at least one future exists in the
/// poller.
///
#[must_use = "futures do nothing unless polled"]
pub struct NextCompletionFuture<'p, F: Future, M: SlotMemory<F>> {
    token: PollThread,
    poller: &'p mut SlotPoller<F, M>,
}

impl<'p, F, M> Future for NextCompletionFuture<'p, F, M>
where
    F: Future,
    M: SlotMemory<F>,
{
    type Output = F::Output;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let token = self.token;
        let EngineFields {
            activity,
            mut slots,
            shared_storage,
        } = self.poller.memory.fields();
        poll_loop!(
            cx,
            token,
            activity,
            slots,
            shared_storage,
            slot,
            slot_idx,
            state_node,
            poll_res,
            Poll::Pending, // nothing to poll => keep waiting
            {
                if let Poll::Ready((res, droppable)) = poll_res {
                    // Future completed - mark as empty / drop it
                    activity.empty_stack_push(slot_idx, slot.as_mut());
                    droppable.drop_future(slot, state_node);
                    return Poll::Ready(res);
                }
            }
        )
    }
}

//
// Draining
//

/// Helper struct that disavows pinning for the inner value
struct AssertUnpin<V>(V);

impl<V> Unpin for AssertUnpin<V> {}

///
/// The future returned from [SlotPoller::drain].
///
/// This future will poll all the futures inside the slot poller. When one of them completes, the
/// result is handled by the drain function.
///
#[must_use = "futures do nothing unless polled"]
pub struct DrainFuture<'p, F: Future, M: SlotMemory<F>, D: FnMut(F::Output)> {
    token: PollThread,
    poller: &'p mut SlotPoller<F, M>,
    drain_function: AssertUnpin<D>,
}

///
/// The future returned from [SlotPoller::try_drain].
///
/// This future will poll all the futures inside the slot poller. When one of them completes, the
/// result is handled by the drain function, which can return an error if it wants to abort the drain.
///
#[must_use = "futures do nothing unless polled"]
pub struct TryDrainFuture<'p, F: Future, M: SlotMemory<F>, D: FnMut(F::Output) -> Result<(), E>, E>
{
    token: PollThread,
    poller: &'p mut SlotPoller<F, M>,
    drain_function: AssertUnpin<D>,
}

impl<'p, F, M, D> Future for DrainFuture<'p, F, M, D>
where
    F: Future,
    M: SlotMemory<F>,
    D: FnMut(F::Output),
{
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let token = self.token;
        let this = self.get_mut();
        let EngineFields {
            activity,
            mut slots,
            shared_storage,
        } = this.poller.memory.fields();
        let drain_function = &mut this.drain_function.0;
        poll_loop!(
            cx,
            token,
            activity,
            slots,
            shared_storage,
            slot,
            slot_idx,
            state_node,
            poll_res,
            {
                // Nothing to poll, so check if complete
                if activity.slots_active == 0 {
                    Poll::Ready(())
                } else {
                    Poll::Pending
                }
            },
            {
                if let Poll::Ready((res, droppable)) = poll_res {
                    activity.empty_stack_push(slot_idx, slot.as_mut());
                    droppable.drop_future(slot, state_node);
                    drain_function(res);
                }
            }
        )
    }
}

impl<'p, F, M, D, E> Future for TryDrainFuture<'p, F, M, D, E>
where
    F: Future,
    M: SlotMemory<F>,
    D: FnMut(F::Output) -> Result<(), E>,
{
    type Output = Result<(), E>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let token = self.token;
        let this = self.get_mut();
        let EngineFields {
            activity,
            mut slots,
            shared_storage,
        } = this.poller.memory.fields();
        let drain_function = &mut this.drain_function.0;
        poll_loop!(
            cx,
            token,
            activity,
            slots,
            shared_storage,
            slot,
            slot_idx,
            state_node,
            poll_res,
            {
                // Nothing to poll, so check if complete
                if activity.slots_active == 0 {
                    Poll::Ready(Ok(()))
                } else {
                    Poll::Pending
                }
            },
            {
                if let Poll::Ready((res, droppable)) = poll_res {
                    activity.empty_stack_push(slot_idx, slot.as_mut());
                    droppable.drop_future(slot, state_node);
                    if let Err(e) = drain_function(res) {
                        return Poll::Ready(Err(e));
                    }
                }
            }
        )
    }
}