async_selector/

selector.rs

//! Fast and flexible [`Future`]/[`Stream`] selector.

use std::{
    fmt,
    marker::PhantomData,
    ops::{ControlFlow, Not},
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
};

use futures::Stream;

use crate::{
    list::{
        IntrusiveList,
        cursor::{Cursor, CursorMut},
    },
    mpsc,
    pollable::{PollStrategy, PollWith},
    selector::{
        ext::{WithExt, WithExtAndId, WithId},
        iter::{ExtractIf, IntoIter, Iter, IterMut},
    },
    task::Task,
};

pub use borrowed::{Borrowed, BorrowedMut};
pub use id::Id;
pub use removed::Removed;

mod borrowed;
pub mod ext;
mod id;
pub mod iter;
mod removed;

/// Selector over a dynamic set of pollable tasks, for example [`Future`]s and [`Stream`]s.
///
/// Designed for optimal performance when polling a large number of tasks
/// (see [example](https://github.com/Razz4780/async-selector/blob/main/examples/speed.rs)).
///
/// Allows for:
/// 1. Safely injecting shared state into the tasks (see [`PollWith`]).
/// 2. Accessing and removing the tasks by unique ids.
///
/// Unless you want to exercise the full flexibility of this type,
/// you can use the specialized [`FutureSelector`](crate::FutureSelector)
/// and [`StreamSelector`](crate::StreamSelector).
///
/// # Removal
///
/// The selector creates a heap allocation for each stored task.
/// Removing a task from the selector does not instantly free that memory.
/// The memory can only be freed when:
/// 1. all [`Id`] instances for this task are dropped, AND
/// 2. [`Removed`] instance is dropped or consumed, AND
/// 3. the [`Waker`](std::task::Waker) (and all its clones)
///    passed when polling the task are dropped or consumed, AND
/// 4. the selector observes the task removal.
///
/// The selector always eventually observes the removal when polled.
///
/// # Wakeups
///
/// The selector uses a smart strategy for polling the tasks.
/// A task is **only** polled in the following cases:
/// 1. after it is pushed into the selector
/// 2. after it yields a non-terminal value
/// 3. after the waker passed to [`PollWith::poll_progress`] inside [`Context`] is woken
///
/// To avoid nasty surprises, keep this in mind when:
/// 1. Modifying a task borrowed from the selector
/// 2. Polling the selector with different extension types
///
/// See [example](https://github.com/Razz4780/async-selector/blob/main/examples/extensions.rs).
///
/// # Panic
///
/// If the task's [`PollWith::poll_progress`] implementation panics,
/// the task is removed from the selector and dropped.
/// The selector remains valid.
pub struct Selector<S: PollStrategy> {
    /// Queue of tasks that were woken.
    ready_rx: mpsc::Receiver<Task<S::Pollable>>,
    /// List of all tasks.
    list: IntrusiveList<Task<S::Pollable>>,
    /// [`PollStrategy`] determining how we poll tasks.
    _phantom: PhantomData<fn() -> S>,
}

impl<S: PollStrategy> Selector<S> {
    /// Pushes a new task into the selector.
    ///
    /// This method is O(1).
    pub fn push(&mut self, pollable: S::Pollable) {
        let node = self
            .list
            .insert(Task::empty(self.ready_rx.weak_sender()), pollable);
        self.ready_rx.send(node);
    }

    /// Pushes a new task into the selector and returns its unique id.
    ///
    /// This method is O(1).
    pub fn push_with_id(&mut self, pollable: S::Pollable) -> Id<S::Pollable> {
        let node = self
            .list
            .insert(Task::empty(self.ready_rx.weak_sender()), pollable);
        let id = Id(Arc::downgrade(&node));
        self.ready_rx.send(node);
        id
    }

    /// Returns whether the selector is empty.
    ///
    /// This method is O(1).
    pub fn is_empty(&self) -> bool {
        self.list.is_empty()
    }

    /// Returns the number of tasks in the selector.
    ///
    /// This method is O(1).
    pub fn len(&self) -> usize {
        self.list.len()
    }

    /// Manually wakes all tasks in the selector.
    ///
    /// Depending on the tasks' [`PollWith`] implementation,
    /// this might be required when polling with different extension types.
    /// See the wakeups [section](Selector#wakeups).
    ///
    /// This method is O(n).
    pub fn wake_all(&self) {
        self.iter().for_each(|borrowed| borrowed.wake());
    }

    /// Returns an iterator over the tasks in the selector.
    ///
    /// The tasks are visited in the insertion order.
    pub fn iter(&self) -> Iter<'_, S::Pollable> {
        Iter {
            cursor: Cursor::new(&self.list),
            queue: &self.ready_rx,
        }
    }

    /// Returns an iterator that allows modifying each task in the selector.
    ///
    /// The tasks are visited in the insertion order.
    ///
    /// **Important:** before modifying tasks stored in the selector, see the wakeups [section](Selector#wakeups).
    pub fn iter_mut(&mut self) -> IterMut<'_, S::Pollable> {
        IterMut {
            cursor: CursorMut::new(&mut self.list),
            queue: &self.ready_rx,
        }
    }

    /// Creates an iterator which uses a closure to determine if a task should be removed.
    ///
    /// If the closure returns true, the task is removed from the selector and yielded.
    ///
    /// If the returned [`ExtractIf`] is not exhausted, e.g. because it is dropped without iterating or the iteration short-circuits,
    /// then the remaining tasks will be retained.
    ///
    /// **Important:** before removing tasks from the selector, see the removal [section](Selector#removal).
    #[must_use = "ExtractIf does not remove any elements unless consumed"]
    pub fn extract_if<F>(&mut self, pred: F) -> ExtractIf<'_, S::Pollable, F>
    where
        F: FnMut(Pin<&mut S::Pollable>) -> bool,
    {
        ExtractIf {
            cursor: CursorMut::new(&mut self.list),
            pred,
        }
    }

    /// Returns an immutable reference to the task with the given id.
    ///
    /// Returns `None` if the task is not found in this selector,
    /// for example because it was removed or has already finished.
    ///
    /// This method is O(1).
    pub fn get(&self, id: &Id<S::Pollable>) -> Option<Borrowed<'_, S::Pollable>> {
        let task = id.0.upgrade()?;
        if self.ready_rx.is_parent(task.ready_tx()).not() {
            return None;
        }
        let node = unsafe {
            // SAFETY: we just checked that this id comes from this selector.
            // Therefore, the task cannot be stored in any other list.
            self.list.get(&task)
        }?;
        Some(Borrowed {
            node,
            queue: &self.ready_rx,
        })
    }

    /// Returns a mutable reference to the target with the given id.
    ///
    /// Returns `None` if the task is not found in this selector,
    /// for example because it was removed or has already finished.
    ///
    /// This method is O(1).
    ///
    /// **Important:** before modifying tasks stored in the selector, see the wakeups [section](Selector#wakeups).
    pub fn get_mut(&mut self, id: &Id<S::Pollable>) -> Option<BorrowedMut<'_, S::Pollable>> {
        let task = id.0.upgrade()?;
        if self.ready_rx.is_parent(task.ready_tx()).not() {
            return None;
        }
        let node = unsafe {
            // SAFETY: we just checked that this id comes from this selector.
            // Therefore, the task cannot be stored in any other list.
            self.list.get_mut(&task)
        }?;
        Some(BorrowedMut {
            node,
            queue: &self.ready_rx,
        })
    }

    /// Removes the task with the given id from the selector.
    ///
    /// Returns `None` if the task is not found in the selector,
    /// for example because it was removed or has already finished.
    ///
    /// This method is O(1).
    ///
    /// **Important:** before removing tasks from the selector, see the removal [section](Selector#removal).
    pub fn remove(&mut self, id: &Id<S::Pollable>) -> Option<Removed<S::Pollable>> {
        let task = id.0.upgrade()?;
        if self.ready_rx.is_parent(task.ready_tx()).not() {
            return None;
        }
        let removed = unsafe {
            // SAFETY: we just checked that this id comes from this selector.
            // Therefore, the task cannot be stored in any other list.
            self.list.remove(&task)?
        };
        Some(Removed(removed))
    }

    /// Borrows a [`Stream`] that will pass the given extensions to the tasks when polling.
    ///
    /// **Important:** before polling the tasks with different extension types, see the wakeups [section](Selector#wakeups).
    pub fn with_ext<'s, 'e, 'emut, E: ?Sized, EMut: ?Sized>(
        &'s mut self,
        ext: &'e E,
        ext_mut: &'emut mut EMut,
    ) -> WithExt<'s, 'e, 'emut, S, E, EMut> {
        WithExt {
            selector: self,
            ext,
            ext_mut,
        }
    }

    /// Borrows a [`Stream`] that will return a task [`Id`] with every item.
    pub fn with_id(&mut self) -> WithId<'_, S> {
        WithId { selector: self }
    }

    /// Borrows a [`Stream`] that will pass the given extensions to the tasks when polling,
    /// and return a task [`Id`] with every item.
    ///
    /// **Important:** before polling the tasks with different extension types, see the wakeups [section](Selector#wakeups).
    pub fn with_ext_and_id<'s, 'e, 'emut, E: ?Sized, EMut: ?Sized>(
        &'s mut self,
        ext: &'e E,
        ext_mut: &'emut mut EMut,
    ) -> WithExtAndId<'s, 'e, 'emut, S, E, EMut> {
        WithExtAndId {
            selector: self,
            ext,
            ext_mut,
        }
    }

    /// Returns the next ready item from one of the tasks stored in the selector.
    ///
    /// Provided extensions will be passed down to the tasks as arguments to [`PollWith::poll_progress`].
    /// Provided `extra` lambda can be used to extract additional data from the [`Task`] instance.
    ///
    /// Returns `None` if the selector is empty.
    #[allow(clippy::type_complexity)]
    fn poll_next_inner<'a, E, EMut, F, T>(
        &mut self,
        ext: &'a E,
        ext_mut: &mut EMut,
        extra: F,
        cx: &mut Context<'_>,
    ) -> Poll<Option<(<S as PollWith<'a, E, EMut>>::Progress, T)>>
    where
        S: PollWith<'a, E, EMut>,
        E: ?Sized,
        EMut: ?Sized,
        F: FnOnce(&Arc<Task<S::Pollable>>) -> T,
    {
        let marker = self.ready_rx.register(cx.waker());
        if marker.is_null() {
            return if self.list.is_empty() {
                Poll::Ready(None)
            } else {
                Poll::Pending
            };
        }

        let mut polled_all_queue = false;
        while polled_all_queue.not() {
            let task = match self.ready_rx.recv() {
                Some(task) => {
                    polled_all_queue = std::ptr::eq(task.as_ref(), marker);
                    task
                }
                None if self.list.is_empty() => return Poll::Ready(None),
                None => return Poll::Pending,
            };

            let mut guard = {
                let guard = unsafe {
                    // SAFETY: we received this task from our ready queue,
                    // so it must be ours.
                    self.list.access(&task)
                };
                match guard {
                    Some(guard) => guard,
                    None => continue,
                }
            };
            let waker = task.borrow_waker();
            let mut cx = Context::from_waker(&waker);

            // If this method panics, the guard will automatically remove the task from the list and drop it.
            let result = S::poll_progress(guard.get(), ext, ext_mut, &mut cx);

            match result {
                Poll::Ready(ControlFlow::Continue(item)) => {
                    guard.forget();
                    let extra = extra(&task);
                    self.ready_rx.send(task);
                    return Poll::Ready(Some((item, extra)));
                }
                Poll::Ready(ControlFlow::Break(Some(item))) => {
                    let extra = extra(&task);
                    return Poll::Ready(Some((item, extra)));
                }
                Poll::Ready(ControlFlow::Break(None)) => {}
                Poll::Pending => guard.forget(),
            }
        }

        if self.list.is_empty() {
            Poll::Ready(None)
        } else {
            Poll::Pending
        }
    }
}

impl<S: PollWith<'static, (), ()>> Stream for Selector<S> {
    type Item = S::Progress;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = unsafe {
            // SAFETY: no field is ever moved in memory.
            self.get_unchecked_mut()
        };
        this.poll_next_inner(&(), &mut (), |_| (), cx)
            .map(|opt| opt.map(|(item, ())| item))
    }
}

impl<S: PollStrategy> Default for Selector<S> {
    fn default() -> Self {
        Self {
            ready_rx: mpsc::Receiver::new(Task::empty),
            list: Default::default(),
            _phantom: Default::default(),
        }
    }
}

impl<S: PollStrategy> fmt::Debug for Selector<S> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Selector")
            .field("len", &self.len())
            .finish()
    }
}

impl<S: PollStrategy> Extend<S::Pollable> for Selector<S> {
    fn extend<T: IntoIterator<Item = S::Pollable>>(&mut self, iter: T) {
        for pollable in iter {
            self.push(pollable);
        }
    }
}

impl<S: PollStrategy> FromIterator<S::Pollable> for Selector<S> {
    fn from_iter<T: IntoIterator<Item = S::Pollable>>(iter: T) -> Self {
        let mut this = Self::default();
        this.extend(iter);
        this
    }
}

impl<S: PollStrategy> IntoIterator for Selector<S> {
    type IntoIter = IntoIter<S::Pollable>;
    type Item = Removed<S::Pollable>;

    fn into_iter(self) -> Self::IntoIter {
        IntoIter(self.list)
    }
}

#[cfg(test)]
mod test {
    use std::{
        ops::Not,
        panic::{AssertUnwindSafe, catch_unwind},
        pin::Pin,
        sync::Arc,
        task::{Context, Poll, Waker},
    };

    use futures::{FutureExt, StreamExt, channel::oneshot, task::AtomicWaker};
    use rstest::rstest;

    use crate::{pollable::PollAsFuture, selector::Selector};

    #[tokio::test]
    async fn basic() {
        let (tx, rx) = oneshot::channel::<()>();
        let mut selector = Selector::<PollAsFuture<_>>::default();
        selector.push(rx);
        assert!(selector.next().now_or_never().is_none());
        assert_eq!(selector.len(), 1);
        tx.send(()).unwrap();
        assert!(selector.next().await.is_some());
        assert_eq!(selector.len(), 0);
    }

    /// Verifies that [`Selector`] respects the inner item's yield when polled,
    /// and does not poll the same item twice in a single [`Selector::poll_next_with_ext`] call.
    #[rstest]
    #[tokio::test]
    async fn task_yield_is_respected(#[values(1, 4, 8)] futures: usize) {
        struct Fut {
            polled: bool,
        }

        impl Future for Fut {
            type Output = ();

            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
                let this = self.get_mut();
                if this.polled.not() {
                    this.polled = true;
                    cx.waker().wake_by_ref();
                    Poll::Pending
                } else {
                    Poll::Ready(())
                }
            }
        }

        let mut selector = Selector::<PollAsFuture<_>>::default();
        for _ in 0..futures {
            selector.push(Fut { polled: false });
        }

        assert!(
            selector
                .poll_next_unpin(&mut Context::from_waker(Waker::noop()))
                .is_pending()
        );
        for fut in selector.iter() {
            assert!(fut.polled);
        }

        for _ in 0..futures {
            assert_eq!(
                selector.poll_next_unpin(&mut Context::from_waker(Waker::noop())),
                Poll::Ready(Some(())),
            );
        }
    }

    #[test]
    fn stale_wakeups_on_removed_tasks_still_report_empty_selector() {
        struct StoreWaker(Arc<AtomicWaker>);

        impl Future for StoreWaker {
            type Output = usize;

            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
                self.0.register(cx.waker());
                Poll::Pending
            }
        }

        let slot = Arc::new(AtomicWaker::new());
        let mut selector = Selector::<PollAsFuture<_>>::default();
        let id = selector.push_with_id(StoreWaker(slot.clone()).boxed());
        let mut cx = Context::from_waker(Waker::noop());

        assert!(selector.poll_next_unpin(&mut cx).is_pending());
        let waker = slot.take().unwrap();
        let _ = selector.remove(&id).unwrap();
        assert!(selector.is_empty());

        for _ in 0..3 {
            waker.wake_by_ref();
            assert_eq!(selector.poll_next_unpin(&mut cx), Poll::Ready(None));
        }

        selector.push(std::future::ready(7).boxed());
        assert_eq!(selector.poll_next_unpin(&mut cx), Poll::Ready(Some(7)));
        assert_eq!(selector.poll_next_unpin(&mut cx), Poll::Ready(None));
    }

    #[test]
    fn panicking_task_is_removed_and_selector_remains_valid() {
        struct PanicOnPoll {
            _shared: Arc<()>,
        }

        impl Future for PanicOnPoll {
            type Output = usize;

            fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Self::Output> {
                panic!("boom");
            }
        }

        let drops = Arc::new(());
        let drops_weak = Arc::downgrade(&drops);
        let mut selector = Selector::<PollAsFuture<_>>::default();
        selector.push(PanicOnPoll { _shared: drops }.boxed());

        let mut cx = Context::from_waker(Waker::noop());
        let result = catch_unwind(AssertUnwindSafe(|| selector.poll_next_unpin(&mut cx)));

        assert!(result.is_err());
        assert!(selector.is_empty());
        assert_eq!(selector.len(), 0);
        assert!(drops_weak.upgrade().is_none());

        selector.push(std::future::ready(11).boxed());
        assert_eq!(selector.poll_next_unpin(&mut cx), Poll::Ready(Some(11)));
        assert_eq!(selector.poll_next_unpin(&mut cx), Poll::Ready(None));
    }
}