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.
 */

//!
//! The guts of the highly concurrent part of the library
//!
//! This module is responsible for the poll, wake, drop cycle that futures go through.
//!

use crate::guard::{OptSlotIdx, PollThread, SlotIdx, StateNodeIdx};
use crate::{EngineActivity, SharedStorage, SharedStorageHeader, Slot, StateNode};
use std::fmt::Debug;
use std::hint::unreachable_unchecked;
use std::marker::PhantomData;
use std::mem::ManuallyDrop;
use std::ops::Deref;
use std::pin::Pin;
use std::ptr;
use std::ptr::NonNull;
use std::sync::atomic::Ordering;
use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
use triomphe::Arc;

/*

The concurrent queue of pollable slots.

We use the intrusive MPSC queue algorithm described on 1024cores.net, and adapt it to be more
greedy. In particular, if an enqueue operation is running concurrently with dequeue, the
consumer will spin in order to obtain the enqueued value. This detail is purely optimization-driven
and could be removed without affecting correctness.

Additionally, the author of this crate was originally using a different algorithm, where "head"
and "tail" have opposite meanings. So, the tail is for enqueue, and head for dequeue. If this
disturbs you, please run a find+replace on the source code, and verify that it's the same algorithm.

https://web.archive.org/web/20190403234426/http://www.1024cores.net/home/lock-free-algorithms/queues/intrusive-mpsc-node-based-queue

 */

/// Removes the first item from the queue
///
/// The poll thread is the single consumer of this queue
fn poll_queue_pop_head(
    _: PollThread,
    head_ref: &mut StateNodeIdx,
    storage: &SharedStorage,
) -> OptSlotIdx {
    /*
    Dequeue. We follow the MPSC algorithm mentioned above.

    When the queue is in an inconsistent state, there's no use in looping/retrying operations.
    That's because the poller being woken up will bring us back to this function.
     */
    unsafe {
        // SAFETY
        // Indices are assumed to be valid insofar as the queue is in an expected state
        // Head is never UNSET, so we can keep using get_poll_queue_link

        let mut head = *head_ref;
        let mut head_next = head.get_poll_queue_link(storage).load(Ordering::Acquire);

        // Skip over stub, if found
        if head == StateNodeIdx::STUB {
            if head_next == StateNodeIdx::UNSET {
                // Empty
                return OptSlotIdx::UNSET;
            }
            *head_ref = head_next;
            head = head_next;
            head_next = head.get_poll_queue_link(storage).load(Ordering::Acquire);
        }

        // Main path: found an element
        if head_next != StateNodeIdx::UNSET {
            *head_ref = head_next;
            debug_assert_ne!(StateNodeIdx::STUB, head, "cannot return stub");
            return head.into_slot_idx().into();
        }

        let tail = storage.header.poll_queue_tail.load(Ordering::Acquire);
        if head != tail {
            // The next item is null, but tail has moved on?
            // This suggests an update is taking place in the enqueue operation
            return OptSlotIdx::UNSET;
        }
        // There is exactly 1 element in the queue, and it isn't the stub.
        // We're the only thread that can re-add the stub. Do this, and try to get the 1 element
        let stub_state_node = StateNodeIdx::STUB.get_state_node(storage);
        poll_queue_push_node_at_index(storage, stub_state_node, StateNodeIdx::STUB);

        head_next = head.get_poll_queue_link(storage).load(Ordering::Acquire);
        if head_next != StateNodeIdx::UNSET {
            // Can't be the stub node.
            *head_ref = head_next;
            debug_assert_ne!(StateNodeIdx::STUB, head, "cannot return stub");
            return head.into_slot_idx().into();
        }
        // The producers must be very aggressive.
        OptSlotIdx::UNSET
    }
}

/// Puts an item back at the front of the queue
///
/// SAFETY
///
/// Caller promises:
/// 1. The state node is not already on the queue
/// 2. The state node index matches the state node provided
/// 3. If not the stub node, the status is either Init or Woken
pub(super) unsafe fn poll_queue_reinsert_head(
    _: PollThread,
    activity: &mut EngineActivity,
    state_node: &StateNode,
    state_node_idx: StateNodeIdx,
) {
    let head = &mut activity.poll_queue_head;

    let prev_head = *head;
    state_node
        .poll_queue_link
        .store(prev_head, Ordering::Relaxed);
    *head = state_node_idx;
}

/// Adds another node to the queue
///
/// SAFETY
///
/// Caller promises:
/// 1. The state node is not already on the queue
/// 2. The state node index matches the state node provided
/// 3. If not the stub node, the status is either Init or Woken
pub(super) unsafe fn poll_queue_push_node_at_index(
    storage: &SharedStorage,
    state_node: &StateNode,
    state_node_idx: StateNodeIdx,
) {
    // Ordering: Relaxed - since slot is not yet on the queue.
    state_node
        .poll_queue_link
        .store(StateNodeIdx::UNSET, Ordering::Relaxed);

    let previous = storage
        .header
        .poll_queue_tail
        .swap(state_node_idx, Ordering::AcqRel);
    let previous_link = unsafe {
        // SAFETY
        // StateNodeIdx::UNSET is never part of the queue
        previous.get_poll_queue_link(storage)
    };
    previous_link.store(state_node_idx, Ordering::Release);
}

/// Removes the head of the pollable queue.
///
/// The status of the removed state node will be either Init or Woken
pub(super) fn pop_pollable<'s>(
    token: PollThread,
    activity: &mut EngineActivity,
    shared_storage: &'s Arc<SharedStorage>,
) -> PollableSlotQuery<'s> {
    let pop_head = poll_queue_pop_head(token, &mut activity.poll_queue_head, shared_storage);
    if pop_head.is_set() {
        let slot_idx = pop_head.into_slot_idx();
        let pollable_slot = PollableSlot::new(shared_storage, slot_idx);
        let status = &pollable_slot.state_node().status;
        /*
        Ordering:
        Relaxed - we just popped this slot off the queue. As such, the queue provides a
        happens-before relationship that lets us observe past state transitions.
         */
        match status.load(Ordering::Relaxed) {
            SlotStatus::Init | SlotStatus::Woken => {
                // Slot is ready to poll
                (pop_head, Some(pollable_slot))
            }
            SlotStatus::UninitButEnqueued => {
                /*
                Race condition; happens with some wake-up calls

                A future can be polled and return completion, and then get dropped. However, it
                might already be enqueued to go another round. This happens because the status
                change can race between the polling thread and the waking thread. The polling
                thread can only react to the future's completion after the future is polled.
                 */
                status.store(SlotStatus::Uninit, Ordering::Relaxed); // reset if needed
                (pop_head, None) // indicates to caller to try again
            }
            SlotStatus::Uninit | SlotStatus::Waiting => {
                // Optimizing really does nothing, but we insert this for debugging and clarity
                unsafe {
                    // SAFETY
                    // Slots with these statuses are never on the queue
                    debug_assert!(false, "Slot with Uninit/Waiting status popped off queue");
                    unreachable_unchecked()
                }
            }
        }
    } else {
        (pop_head, None)
    }
}

fn wake_up_node(storage: &SharedStorage, state_node: &StateNode) {
    /*
    In order for this method call to have any relevance according to the standard library contract,
    it must happen *after* the future has started to be polled. (If it started before the poll
    occurred, then the wake-up could be logically collapsed into that poll).

    Therefore, there is a happens-before relation between the poll and this method. That relationship
    must be established by the future implementation. We can piggyback upon this in order to weaken
    our own orderings.

    Ordering is thus:
    1. Acquire load on success - we need the previously stored waker to be visible, so this creates
    a fence with the waker registration.
    2. Relaxed set on success - the poll queue provides an ordering that makes this write visible.
    3. Relaxed on failure - it is the future's job to create a happens-before relationship.
     */
    let success = state_node.status.compare_and_swap(
        SlotStatus::Waiting,
        SlotStatus::Woken,
        Ordering::Acquire,
        Ordering::Relaxed,
    );
    if success {
        /*
        The wake-up notification happens after the slot is on the queue. This prevents missed
        wake-ups, as it ensures the poller thread can dequeue this slot and poll it.
         */
        unsafe {
            // SAFETY
            // The node isn't on the queue yet, is not the stub, and has status Woken
            let state_node_idx = StateNodeIdx::from(state_node.slot_idx);
            poll_queue_push_node_at_index(storage, state_node, state_node_idx);
        }
        storage.header.main_waker.notify();
    }
}

/// Pretend enum when checking for a pollable slot
///
/// Code:
/// ([OptSlotIdx::UNSET], _) => no slots were on the queue
/// ([SlotIdx], None) => a found slot that wasn't pollable (caller should retry operation)
/// ([SlotIdx], Some) => a found slot that can be polled
pub(super) type PollableSlotQuery<'s> = (OptSlotIdx, Option<PollableSlot<'s>>);

/// A token indicating the future can be polled. Doubles as the state node
pub(super) struct PollableSlot<'s> {
    /*
    For MIRI stacked borrows, this must stay as a pointer. Its offset is manually computed.
    That way, the pointer inherits the access of *const SharedStorage (access to whole allocation)
    https://github.com/rust-lang/unsafe-code-guidelines/issues/134
     */
    state_node_ptr: NonNull<StateNode>,
    lifetime: PhantomData<&'s ()>,
}

impl<'s> PollableSlot<'s> {
    pub(super) fn state_node(&self) -> &'s StateNode {
        unsafe {
            // SAFETY - see struct explanation
            self.state_node_ptr.as_ref()
        }
    }

    fn new(shared_storage: &'s Arc<SharedStorage>, slot_idx: SlotIdx) -> Self {
        unsafe {
            // SAFETY
            // We're performing the inverse operation of recover_shared_storage (below)
            let storage_ptr = Arc::as_ptr(shared_storage);
            let ptr = storage_ptr.byte_add(SharedStorageHeader::BYTE_OFFSET);
            let ptr = ptr as *const StateNode;
            let ptr = ptr.add(1 + slot_idx.value() as usize);
            Self {
                state_node_ptr: NonNull::new_unchecked(ptr as *mut StateNode),
                lifetime: PhantomData,
            }
        }
    }
}

///
/// An internal function to get back the shared storage from a state node pointer
///
/// SAFETY
///
/// 1. The state node pointer must be part of an allocation of SharedStorage
/// 2. It must have provenance over the whole allocation
/// 3. It cannot be the head node
unsafe fn recover_shared_storage(state_node_ptr: *const StateNode) -> *const SharedStorage {
    unsafe {
        // SAFETY
        // Everything is stored in the same allocation
        // 1. Access the state node
        let slot_idx = (*state_node_ptr).slot_idx;
        // 2. Count backwards using the slot index. Add 1 to skip header node
        let ptr = state_node_ptr.sub(1 + slot_idx.value() as usize);
        let ptr = ptr.byte_sub(SharedStorageHeader::BYTE_OFFSET);
        let nodes_len = {
            // Creating a &[T; 0] as part of a &[T; N] will always be safe
            let header = &*(ptr as *const SharedStorageHeader);
            header.nodes_len as usize
        };
        ptr::slice_from_raw_parts(ptr, nodes_len) as *const SharedStorage
    }
}

impl PollableSlot<'_> {
    pub(super) fn call_poll<F: Future>(
        self,
        token: PollThread,
        slot: Pin<&mut Slot<F>>,
    ) -> Poll<(F::Output, Droppable)> {
        // We construct our own waker - a simple, elegant pointer to the state node
        // For efficiency, it is not dropped. Only if the waker is cloned does it become owned
        let our_waker = unsafe {
            // SAFETY
            // 1. The state node comes from the shared storage's allocation, and has provenance
            // 2. The state node is not the stub
            // 3. This waker is not dropped (kept inside ManuallyDrop)
            waker_from_state_node(self.state_node_ptr.as_ptr())
        };
        // A helper responsible for dropping the future if it panics
        struct FutOwner<'poll, F: Future> {
            token: PollThread,
            slot: Pin<&'poll mut Slot<F>>,
            status: &'poll AtomicStatus,
        }
        impl<'poll, F: Future> Drop for FutOwner<'poll, F> {
            fn drop(&mut self) {
                unsafe {
                    // SAFETY
                    // We guarantee the future is initialized
                    call_drop(self.token, self.slot.as_mut(), self.status);
                }
            }
        }
        let mut fut_owner = FutOwner {
            token,
            slot,
            status: &self.state_node().status,
        };
        /*
        Ordering:
        1. Release - in the poll loop, we wrote the latest waker. That write must be visible to
        wake-up calls, so they can take the waker and call it.
        2. Relaxed - this write will itself be visible to wake-up calls, in order to make a status
        change. However, the future is the one responsible for the happens-before relationship.
         */
        fut_owner
            .status
            .store(SlotStatus::Waiting, Ordering::Release);
        let poll_res = {
            let future = fut_owner.slot.as_mut().project().future;
            unsafe {
                // SAFETY
                // We never move out of the future
                // The Pollable token ensures the future is initialized and pollable
                let future = future.map_unchecked_mut(|fut| fut.assume_init_mut());
                // If panic occurs here, FutOwner will drop the future
                future.poll(&mut Context::from_waker(&our_waker))
            }
        };
        std::mem::forget(fut_owner);
        poll_res.map(|val| (val, Droppable(token)))
    }
}

/// Drops the future in the given slot
///
/// SAFETY
///
/// Caller guarantees future is initialized, and the status corresponds to this slot
pub(super) unsafe fn call_drop<F: Future>(
    _: PollThread,
    slot: Pin<&mut Slot<F>>,
    status: &AtomicStatus,
) {
    /*
    Ordering: Relaxed - despite possible contention in the waker, we are on the polling thread.
    Thus, it's possible to "steal" the status atomically.
     */
    let previous_status = status.swap(SlotStatus::Uninit, Ordering::Relaxed);
    if previous_status == SlotStatus::Woken {
        /*
        Already woken, meaning enqueued as pollable.

        Ordering: Relaxed - by now, the state is ours and can't be changed by wake-ups. Note that
        there is a dependency ordering between the previous swap load and this store.
         */
        status.store(SlotStatus::UninitButEnqueued, Ordering::Relaxed);
    }
    unsafe {
        // SAFETY
        // The future stays pinned
        let slot_future = slot.project().future.get_unchecked_mut();
        // StackPoller lives and dies on the polling thread
        // Futures are never added/removed except on that thread
        ptr::drop_in_place(slot_future.as_mut_ptr()); // drop panics OK - guarded
    }
}

/// A token indicating the future can be dropped, because it completed
///
/// This is an internal API, because it's not technically safe. The caller must use the same
/// slot in the call to [Self::drop_future] as the slot from which this type was obtained.
#[must_use]
pub(super) struct Droppable(PollThread);

impl Droppable {
    pub(super) fn drop_future<F: Future>(self, slot: Pin<&mut Slot<F>>, state_node: &StateNode) {
        unsafe {
            // SAFETY
            // This Droppable token guarantees the future is initialized
            // The caller could technically pass a different slot, but this is an internal API
            call_drop(self.0, slot, &state_node.status);
        }
    }
}

impl<F> Slot<F> {
    /// Should be called when free space is available
    ///
    /// Must only be called when the previous future has been dropped, or the slot is empty
    /// (denoted by slot status being [SlotStatus::Uninit]).
    ///
    /// SAFETY
    ///
    /// The state node, slot index, and this slot must all correspond.
    ///
    /// The future in this slot must have been dropped or not initialized.
    pub(super) unsafe fn init_future(
        mut self: Pin<&mut Self>,
        _: PollThread,
        slot_idx: SlotIdx,
        state_node: &StateNode,
        shared_storage: &SharedStorage,
        new_future: F,
    ) {
        /*
        Ordering - Relaxed. We're just crossing the uninit/init line. The past status should be
        Uninit or UninitButEnqueued, so it's okay if a waking thread reads that instead.
         */
        let previous_status = state_node.status.swap(SlotStatus::Init, Ordering::Relaxed);
        unsafe {
            // SAFETY
            // The previous future is dropped, so okay to overwrite memory
            let fut_ptr = self.as_mut().get_unchecked_mut().future.as_mut_ptr();
            ptr::write(fut_ptr, new_future);

            match previous_status {
                // Newly initialized futures are ready to be polled
                SlotStatus::Uninit => {
                    // Safety
                    // The node isn't on the queue yet, is not the stub, and has Init status
                    poll_queue_push_node_at_index(
                        shared_storage,
                        state_node,
                        StateNodeIdx::from(slot_idx),
                    );
                }
                SlotStatus::UninitButEnqueued => {
                    // Ok
                }
                SlotStatus::Waiting | SlotStatus::Woken | SlotStatus::Init => {
                    // Scandal! Horrors! Only an uninitialized future can be overwritten
                    debug_assert!(false, "Pinned future overwritten before drop");
                    unreachable_unchecked()
                }
            }
        }
    }
}

///
/// Creates the waker implementation. It must not be dropped, but lives for as long as the slot
/// state does.
///
/// SAFETY
///
/// The caller promises not to drop the returned value. As for the state node pointer:
/// 1. It must be part of an allocation of SharedStorage
/// 2. It must have provenance over the whole allocation
/// 3. It cannot be the stub node
///
unsafe fn waker_from_state_node(state_node: *const StateNode) -> ManuallyDrop<Waker> {
    /*
    Wakers are implementing by re-using the reference count of the shared storage. This relies on
    the shared storage and state node living in the same allocation, and rewinding a pointer to
    the slot when we need to recover the storage.

    Both the constructor and clone() increment the reference count, while wake and drop use
    decrements. Thus, it's very similar to a reference-counted waker, except that all the wakers
    use the same reference count and are stored statically -- represented as the state node.
     */

    fn clone_impl(data: *const ()) -> RawWaker {
        unsafe {
            // SAFETY
            // Allocation, provenance, and non-head node requirements are upheld by caller
            let shared_storage = recover_shared_storage(data as *const StateNode);
            // Cloning increments the reference count
            let arc = ManuallyDrop::new(Arc::from_raw(shared_storage));
            let _cloned = ManuallyDrop::new(Arc::clone(ManuallyDrop::deref(&arc)));
        }
        // Use the same vtable; everything is otherwise the same
        RawWaker::new(data, VTABLE)
    }
    fn wake_impl(data: *const ()) {
        let state_node = data as *const StateNode;
        unsafe {
            // SAFETY
            // Allocation, provenance, and non-head node requirements are upheld by caller
            let shared_storage = recover_shared_storage(state_node);
            {
                // Create references here
                let state_node = &*state_node;
                let shared_storage = &*shared_storage;
                wake_up_node(shared_storage, state_node);
            }
            // Dropping decrements the reference count
            // Important - the references are no longer in scope
            let _arc = Arc::from_raw(shared_storage);
        }
    }
    fn wake_by_ref_impl(data: *const ()) {
        let state_node = data as *const StateNode;
        unsafe {
            // SAFETY
            // Allocation, provenance, and non-head node requirements are upheld by caller
            let shared_storage = recover_shared_storage(state_node);
            let state_node = &*state_node;
            let shared_storage = &*shared_storage;
            wake_up_node(shared_storage, state_node);
            // No deallocation is performed - just references all around
        }
    }
    fn drop_impl(data: *const ()) {
        let state_node = data as *const StateNode;
        unsafe {
            // SAFETY
            // Allocation, provenance, and non-head node requirements are upheld by caller
            let shared_storage = recover_shared_storage(state_node);
            // Dropping decrements the reference count
            let _arc = Arc::from_raw(shared_storage);
        }
    }

    debug_assert!((|| unsafe {
        let slot_idx = (*state_node).slot_idx;
        let shared_storage = &*recover_shared_storage(state_node);
        let recovered_state_node = slot_idx.get_state_node(shared_storage);
        ptr::eq(state_node, recovered_state_node)
    })());

    const VTABLE: &RawWakerVTable =
        &RawWakerVTable::new(clone_impl, wake_impl, wake_by_ref_impl, drop_impl);
    unsafe {
        // We store the pointer to the state node, and use it to reconstruct the storage later
        let data = state_node as *const ();
        // ManuallyDrop allows us to skip reference count updates for the poll itself
        ManuallyDrop::new(Waker::from_raw(RawWaker::new(data, VTABLE)))
    }
}

pub(super) use slot_status::AtomicStatus;

#[derive(PartialEq, Eq, Debug)]
#[repr(u8)]
pub(super) enum SlotStatus {
    /// No future exists in the slot
    Uninit,
    /// No future exists, due to a drop.
    /// However, a wake was triggered between the call to poll and drop (observed race condition).
    UninitButEnqueued,
    /// The future was initialized.
    /// This must be distinguished from [Self::Woken] in order to disambiguate race conditions
    Init,
    /// The future is susceptible to a wake-up
    Waiting,
    /// The future was woken, as in moved from Waiting to Woken through a waker.
    Woken,
}

mod slot_status {
    use super::SlotStatus;
    use std::fmt::{Debug, Formatter};
    use std::hint::unreachable_unchecked;
    use std::sync::atomic::{AtomicU8, Ordering};

    pub(crate) struct AtomicStatus(AtomicU8);

    impl Debug for AtomicStatus {
        fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
            Debug::fmt(&self.load(Ordering::Relaxed), f)
        }
    }

    impl Default for AtomicStatus {
        fn default() -> Self {
            Self(AtomicU8::new(SlotStatus::Uninit as u8))
        }
    }

    impl AtomicStatus {
        fn extract(val: u8) -> SlotStatus {
            if val == SlotStatus::Uninit as u8 {
                SlotStatus::Uninit
            } else if val == SlotStatus::UninitButEnqueued as u8 {
                SlotStatus::UninitButEnqueued
            } else if val == SlotStatus::Init as u8 {
                SlotStatus::Init
            } else if val == SlotStatus::Waiting as u8 {
                SlotStatus::Waiting
            } else if val == SlotStatus::Woken as u8 {
                SlotStatus::Woken
            } else {
                unsafe {
                    // SAFETY
                    // We never store any other kind of value here
                    unreachable_unchecked()
                }
            }
        }

        pub(crate) fn load(&self, ordering: Ordering) -> SlotStatus {
            Self::extract(self.0.load(ordering))
        }

        pub(super) fn store(&self, val: SlotStatus, ordering: Ordering) {
            self.0.store(val as u8, ordering);
        }

        pub(super) fn compare_and_swap(
            &self,
            current: SlotStatus,
            new: SlotStatus,
            success: Ordering,
            failure: Ordering,
        ) -> bool {
            self.0
                .compare_exchange(current as u8, new as u8, success, failure)
                .is_ok()
        }

        pub(super) fn swap(&self, new: SlotStatus, ordering: Ordering) -> SlotStatus {
            Self::extract(self.0.swap(new as u8, ordering))
        }
    }
}