flume/

select.rs

//! Types that permit waiting upon multiple blocking operations using the [`Selector`] interface.

use crate::*;
use spin1::Mutex as Spinlock;
use std::{any::Any, marker::PhantomData};

#[cfg(feature = "eventual-fairness")]
use nanorand::Rng;

// A unique token corresponding to an event in a selector
type Token = usize;

struct SelectSignal(
    thread::Thread,
    Token,
    AtomicBool,
    Arc<Spinlock<VecDeque<Token>>>,
);

impl Signal for SelectSignal {
    fn fire(&self) -> bool {
        self.2.store(true, Ordering::SeqCst);
        self.3.lock().push_back(self.1);
        self.0.unpark();
        false
    }

    fn as_any(&self) -> &(dyn Any + 'static) {
        self
    }
    fn as_ptr(&self) -> *const () {
        self as *const _ as *const ()
    }
}

trait Selection<'a, T> {
    fn init(&mut self) -> Option<T>;
    fn poll(&mut self) -> Option<T>;
    fn deinit(&mut self);
}

/// An error that may be emitted when attempting to wait for a value on a receiver.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum SelectError {
    /// A timeout occurred when waiting on a `Selector`.
    Timeout,
}

impl fmt::Display for SelectError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            SelectError::Timeout => "timeout occurred".fmt(f),
        }
    }
}

impl std::error::Error for SelectError {}

/// A type used to wait upon multiple blocking operations at once.
///
/// A [`Selector`] implements [`select`](https://en.wikipedia.org/wiki/Select_(Unix))-like behaviour,
/// allowing a thread to wait upon the result of more than one operation at once.
///
/// # Examples
/// ```
/// let (tx0, rx0) = flume::unbounded();
/// let (tx1, rx1) = flume::unbounded();
///
/// std::thread::spawn(move || {
///     tx0.send(true).unwrap();
///     tx1.send(42).unwrap();
/// });
///
/// flume::Selector::new()
///     .recv(&rx0, |b| println!("Received {:?}", b))
///     .recv(&rx1, |n| println!("Received {:?}", n))
///     .wait();
/// ```
pub struct Selector<'a, T: 'a> {
    selections: Vec<Box<dyn Selection<'a, T> + 'a>>,
    next_poll: usize,
    signalled: Arc<Spinlock<VecDeque<Token>>>,
    #[cfg(feature = "eventual-fairness")]
    rng: nanorand::WyRand,
    phantom: PhantomData<*const ()>,
}

impl<'a, T: 'a> Default for Selector<'a, T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<'a, T: 'a> fmt::Debug for Selector<'a, T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Selector").finish()
    }
}

impl<'a, T> Selector<'a, T> {
    /// Create a new selector.
    pub fn new() -> Self {
        Self {
            selections: Vec::new(),
            next_poll: 0,
            signalled: Arc::default(),
            phantom: PhantomData::default(),
            #[cfg(feature = "eventual-fairness")]
            rng: nanorand::WyRand::new(),
        }
    }

    /// Add a send operation to the selector that sends the provided value.
    ///
    /// Once added, the selector can be used to run the provided handler function on completion of this operation.
    pub fn send<U, F: FnMut(Result<(), SendError<U>>) -> T + 'a>(
        mut self,
        sender: &'a Sender<U>,
        msg: U,
        mapper: F,
    ) -> Self {
        struct SendSelection<'a, T, F, U> {
            sender: &'a Sender<U>,
            msg: Option<U>,
            token: Token,
            signalled: Arc<Spinlock<VecDeque<Token>>>,
            hook: Option<Arc<Hook<U, SelectSignal>>>,
            mapper: F,
            phantom: PhantomData<T>,
        }

        impl<'a, T, F, U> Selection<'a, T> for SendSelection<'a, T, F, U>
        where
            F: FnMut(Result<(), SendError<U>>) -> T,
        {
            fn init(&mut self) -> Option<T> {
                let token = self.token;
                let signalled = self.signalled.clone();
                let r = self.sender.shared.send(
                    self.msg.take().unwrap(),
                    true,
                    |msg| {
                        Hook::slot(
                            Some(msg),
                            SelectSignal(
                                thread::current(),
                                token,
                                AtomicBool::new(false),
                                signalled,
                            ),
                        )
                    },
                    // Always runs
                    |h| {
                        self.hook = Some(h);
                        Ok(())
                    },
                );

                if self.hook.is_none() {
                    Some((self.mapper)(match r {
                        Ok(()) => Ok(()),
                        Err(TrySendTimeoutError::Disconnected(msg)) => Err(SendError(msg)),
                        _ => unreachable!(),
                    }))
                } else {
                    None
                }
            }

            fn poll(&mut self) -> Option<T> {
                let res = if self.sender.shared.is_disconnected() {
                    // Check the hook one last time
                    if let Some(msg) = self.hook.as_ref()?.try_take() {
                        Err(SendError(msg))
                    } else {
                        Ok(())
                    }
                } else if self.hook.as_ref().unwrap().is_empty() {
                    // The message was sent
                    Ok(())
                } else {
                    return None;
                };

                Some((&mut self.mapper)(res))
            }

            fn deinit(&mut self) {
                if let Some(hook) = self.hook.take() {
                    // Remove hook
                    let hook: Arc<Hook<U, dyn Signal>> = hook;
                    wait_lock(&self.sender.shared.chan)
                        .sending
                        .as_mut()
                        .unwrap()
                        .1
                        .retain(|s| s.signal().as_ptr() != hook.signal().as_ptr());
                }
            }
        }

        let token = self.selections.len();
        self.selections.push(Box::new(SendSelection {
            sender,
            msg: Some(msg),
            token,
            signalled: self.signalled.clone(),
            hook: None,
            mapper,
            phantom: Default::default(),
        }));

        self
    }

    /// Add a receive operation to the selector.
    ///
    /// Once added, the selector can be used to run the provided handler function on completion of this operation.
    pub fn recv<U, F: FnMut(Result<U, RecvError>) -> T + 'a>(
        mut self,
        receiver: &'a Receiver<U>,
        mapper: F,
    ) -> Self {
        struct RecvSelection<'a, T, F, U> {
            receiver: &'a Receiver<U>,
            token: Token,
            signalled: Arc<Spinlock<VecDeque<Token>>>,
            hook: Option<Arc<Hook<U, SelectSignal>>>,
            mapper: F,
            received: bool,
            phantom: PhantomData<T>,
        }

        impl<'a, T, F, U> Selection<'a, T> for RecvSelection<'a, T, F, U>
        where
            F: FnMut(Result<U, RecvError>) -> T,
        {
            fn init(&mut self) -> Option<T> {
                let token = self.token;
                let signalled = self.signalled.clone();
                let r = self.receiver.shared.recv(
                    true,
                    || {
                        Hook::trigger(SelectSignal(
                            thread::current(),
                            token,
                            AtomicBool::new(false),
                            signalled,
                        ))
                    },
                    // Always runs
                    |h| {
                        self.hook = Some(h);
                        Err(TryRecvTimeoutError::Timeout)
                    },
                );

                if self.hook.is_none() {
                    Some((self.mapper)(match r {
                        Ok(msg) => Ok(msg),
                        Err(TryRecvTimeoutError::Disconnected) => Err(RecvError::Disconnected),
                        _ => unreachable!(),
                    }))
                } else {
                    None
                }
            }

            fn poll(&mut self) -> Option<T> {
                let res = if let Ok(msg) = self.receiver.try_recv() {
                    self.received = true;
                    Ok(msg)
                } else if self.receiver.shared.is_disconnected() {
                    Err(RecvError::Disconnected)
                } else {
                    return None;
                };

                Some((&mut self.mapper)(res))
            }

            fn deinit(&mut self) {
                if let Some(hook) = self.hook.take() {
                    // Remove hook
                    let hook: Arc<Hook<U, dyn Signal>> = hook;
                    wait_lock(&self.receiver.shared.chan)
                        .waiting
                        .retain(|s| s.signal().as_ptr() != hook.signal().as_ptr());
                    // If we were woken, but never polled, wake up another
                    if !self.received
                        && hook
                            .signal()
                            .as_any()
                            .downcast_ref::<SelectSignal>()
                            .unwrap()
                            .2
                            .load(Ordering::SeqCst)
                    {
                        wait_lock(&self.receiver.shared.chan).try_wake_receiver_if_pending();
                    }
                }
            }
        }

        let token = self.selections.len();
        self.selections.push(Box::new(RecvSelection {
            receiver,
            token,
            signalled: self.signalled.clone(),
            hook: None,
            mapper,
            received: false,
            phantom: Default::default(),
        }));

        self
    }

    fn wait_inner(mut self, deadline: Option<Instant>) -> Option<T> {
        #[cfg(feature = "eventual-fairness")]
        {
            self.next_poll = self.rng.generate_range(0..self.selections.len());
        }

        let res = 'outer: loop {
            // Init signals
            for _ in 0..self.selections.len() {
                if let Some(val) = self.selections[self.next_poll].init() {
                    break 'outer Some(val);
                }
                self.next_poll = (self.next_poll + 1) % self.selections.len();
            }

            // Speculatively poll
            if let Some(msg) = self.poll() {
                break 'outer Some(msg);
            }

            loop {
                if let Some(deadline) = deadline {
                    if let Some(dur) = deadline.checked_duration_since(Instant::now()) {
                        thread::park_timeout(dur);
                    }
                } else {
                    thread::park();
                }

                if deadline.map(|d| Instant::now() >= d).unwrap_or(false) {
                    break 'outer self.poll();
                }

                let token = if let Some(token) = self.signalled.lock().pop_front() {
                    token
                } else {
                    // Spurious wakeup, park again
                    continue;
                };

                // Attempt to receive a message
                if let Some(msg) = self.selections[token].poll() {
                    break 'outer Some(msg);
                }
            }
        };

        // Deinit signals
        for s in &mut self.selections {
            s.deinit();
        }

        res
    }

    fn poll(&mut self) -> Option<T> {
        for _ in 0..self.selections.len() {
            if let Some(val) = self.selections[self.next_poll].poll() {
                return Some(val);
            }
            self.next_poll = (self.next_poll + 1) % self.selections.len();
        }
        None
    }

    /// Wait until one of the events associated with this [`Selector`] has completed. If the `eventual-fairness`
    /// feature flag is enabled, this method is fair and will handle a random event of those that are ready.
    pub fn wait(self) -> T {
        self.wait_inner(None).unwrap()
    }

    /// Wait until one of the events associated with this [`Selector`] has completed or the timeout has expired. If the
    /// `eventual-fairness` feature flag is enabled, this method is fair and will handle a random event of those that
    /// are ready.
    pub fn wait_timeout(self, dur: Duration) -> Result<T, SelectError> {
        self.wait_inner(Some(Instant::now() + dur))
            .ok_or(SelectError::Timeout)
    }

    /// Wait until one of the events associated with this [`Selector`] has completed or the deadline has been reached.
    /// If the `eventual-fairness` feature flag is enabled, this method is fair and will handle a random event of those
    /// that are ready.
    pub fn wait_deadline(self, deadline: Instant) -> Result<T, SelectError> {
        self.wait_inner(Some(deadline)).ok_or(SelectError::Timeout)
    }
}