path: root/tokio/src/io/driver/scheduled_io.rs

use super::{Ready, ReadyEvent, Tick};
use crate::loom::sync::atomic::AtomicUsize;
use crate::loom::sync::Mutex;
use crate::util::bit;
use crate::util::slab::Entry;

use std::sync::atomic::Ordering::{AcqRel, Acquire, Release};
use std::task::{Context, Poll, Waker};

use super::Direction;

cfg_io_readiness! {
    use crate::util::linked_list::{self, LinkedList};

    use std::cell::UnsafeCell;
    use std::future::Future;
    use std::marker::PhantomPinned;
    use std::pin::Pin;
    use std::ptr::NonNull;
}

/// Stored in the I/O driver resource slab.
#[derive(Debug)]
pub(crate) struct ScheduledIo {
    /// Packs the resource's readiness with the resource's generation.
    readiness: AtomicUsize,

    waiters: Mutex<Waiters>,
}

cfg_io_readiness! {
    type WaitList = LinkedList<Waiter, <Waiter as linked_list::Link>::Target>;
}

#[derive(Debug, Default)]
struct Waiters {
    #[cfg(feature = "net")]
    /// List of all current waiters
    list: WaitList,

    /// Waker used for AsyncRead
    reader: Option<Waker>,

    /// Waker used for AsyncWrite
    writer: Option<Waker>,

    /// True if this ScheduledIo has been killed due to IO driver shutdown
    is_shutdown: bool,
}

cfg_io_readiness! {
    #[derive(Debug)]
    struct Waiter {
        pointers: linked_list::Pointers<Waiter>,

        /// The waker for this task
        waker: Option<Waker>,

        /// The interest this waiter is waiting on
        interest: mio::Interest,

        is_ready: bool,

        /// Should never be `!Unpin`
        _p: PhantomPinned,
    }

    /// Future returned by `readiness()`
    struct Readiness<'a> {
        scheduled_io: &'a ScheduledIo,

        state: State,

        /// Entry in the waiter `LinkedList`.
        waiter: UnsafeCell<Waiter>,
    }

    enum State {
        Init,
        Waiting,
        Done,
    }
}

// The `ScheduledIo::readiness` (`AtomicUsize`) is packed full of goodness.
//
// | reserved | generation |  driver tick | readinesss |
// |----------+------------+--------------+------------|
// |   1 bit  |   7 bits   +    8 bits    +   16 bits  |

const READINESS: bit::Pack = bit::Pack::least_significant(16);

const TICK: bit::Pack = READINESS.then(8);

const GENERATION: bit::Pack = TICK.then(7);

#[test]
fn test_generations_assert_same() {
    assert_eq!(super::GENERATION, GENERATION);
}

// ===== impl ScheduledIo =====

impl Entry for ScheduledIo {
    fn reset(&self) {
        let state = self.readiness.load(Acquire);

        let generation = GENERATION.unpack(state);
        let next = GENERATION.pack_lossy(generation + 1, 0);

        self.readiness.store(next, Release);
    }
}

impl Default for ScheduledIo {
    fn default() -> ScheduledIo {
        ScheduledIo {
            readiness: AtomicUsize::new(0),
            waiters: Mutex::new(Default::default()),
        }
    }
}

impl ScheduledIo {
    pub(crate) fn generation(&self) -> usize {
        GENERATION.unpack(self.readiness.load(Acquire))
    }

    /// Invoked when the IO driver is shut down; forces this ScheduledIo into a
    /// permanently ready state.
    pub(super) fn shutdown(&self) {
        self.wake0(Ready::ALL, true)
    }

    /// Sets the readiness on this `ScheduledIo` by invoking the given closure on
    /// the current value, returning the previous readiness value.
    ///
    /// # Arguments
    /// - `token`: the token for this `ScheduledIo`.
    /// - `tick`: whether setting the tick or trying to clear readiness for a
    ///    specific tick.
    /// - `f`: a closure returning a new readiness value given the previous
    ///   readiness.
    ///
    /// # Returns
    ///
    /// If the given token's generation no longer matches the `ScheduledIo`'s
    /// generation, then the corresponding IO resource has been removed and
    /// replaced with a new resource. In that case, this method returns `Err`.
    /// Otherwise, this returns the previous readiness.
    pub(super) fn set_readiness(
        &self,
        token: Option<usize>,
        tick: Tick,
        f: impl Fn(Ready) -> Ready,
    ) -> Result<(), ()> {
        let mut current = self.readiness.load(Acquire);

        loop {
            let current_generation = GENERATION.unpack(current);

            if let Some(token) = token {
                // Check that the generation for this access is still the
                // current one.
                if GENERATION.unpack(token) != current_generation {
                    return Err(());
                }
            }

            // Mask out the tick/generation bits so that the modifying
            // function doesn't see them.
            let current_readiness = Ready::from_usize(current);
            let new = f(current_readiness);

            let packed = match tick {
                Tick::Set(t) => TICK.pack(t as usize, new.as_usize()),
                Tick::Clear(t) => {
                    if TICK.unpack(current) as u8 != t {
                        // Trying to clear readiness with an old event!
                        return Err(());
                    }

                    TICK.pack(t as usize, new.as_usize())
                }
            };

            let next = GENERATION.pack(current_generation, packed);

            match self
                .readiness
                .compare_exchange(current, next, AcqRel, Acquire)
            {
                Ok(_) => return Ok(()),
                // we lost the race, retry!
                Err(actual) => current = actual,
            }
        }
    }

    /// Notifies all pending waiters that have registered interest in `ready`.
    ///
    /// There may be many waiters to notify. Waking the pending task **must** be
    /// done from outside of the lock otherwise there is a potential for a
    /// deadlock.
    ///
    /// A stack array of wakers is created and filled with wakers to notify, the
    /// lock is released, and the wakers are notified. Because there may be more
    /// than 32 wakers to notify, if the stack array fills up, the lock is
    /// released, the array is cleared, and the iteration continues.
    pub(super) fn wake(&self, ready: Ready) {
        self.wake0(ready, false);
    }

    fn wake0(&self, ready: Ready, shutdown: bool) {
        const NUM_WAKERS: usize = 32;

        let mut wakers: [Option<Waker>; NUM_WAKERS] = Default::default();
        let mut curr = 0;

        let mut waiters = self.waiters.lock();

        waiters.is_shutdown |= shutdown;

        // check for AsyncRead slot
        if ready.is_readable() {
            if let Some(waker) = waiters.reader.take() {
                wakers[curr] = Some(waker);
                curr += 1;
            }
        }

        // check for AsyncWrite slot
        if ready.is_writable() {
            if let Some(waker) = waiters.writer.take() {
                wakers[curr] = Some(waker);
                curr += 1;
            }
        }

        #[cfg(feature = "net")]
        'outer: loop {
            let mut iter = waiters.list.drain_filter(|w| ready.satisfies(w.interest));

            while curr < NUM_WAKERS {
                match iter.next() {
                    Some(waiter) => {
                        let waiter = unsafe { &mut *waiter.as_ptr() };

                        if let Some(waker) = waiter.waker.take() {
                            waiter.is_ready = true;
                            wakers[curr] = Some(waker);
                            curr += 1;
                        }
                    }
                    None => {
                        break 'outer;
                    }
                }
            }

            drop(waiters);

            for waker in wakers.iter_mut().take(curr) {
                waker