nuclei 0.5.0

//! What a shard's poller watches, by token: its IO sources, and the
//! children whose exits it waits for.
//!
//! A token is a slot index and the slot's generation, and is what the
//! kernel hands back with each event (epoll data, kevent udata): never a
//! pointer. A slot's generation moves on when its registration goes, so an
//! event still in flight for a registration that is gone, or for a slot
//! reused since, finds nothing. Tokens fit a pointer-sized word, since that
//! is all kqueue carries. The word's top bit says which table a token is
//! of: a source's has it clear, a child's set.

use std::io;
use std::sync::Arc;

use crate::io::Source;

/// Bits of a token that index the slot; the rest carry the generation.
const INDEX_BITS: u32 = if usize::BITS >= 64 { 32 } else { 20 };
const INDEX_MASK: u64 = (1 << INDEX_BITS) - 1;
/// The top bit of a pointer-sized word: set in every token of the children.
pub(crate) const CHILD: u64 = 1 << (usize::BITS - 1);
/// Generations wrap within what a token holds between the index and the
/// top bit.
const GEN_MASK: u64 = ((usize::MAX as u64) >> 1) >> INDEX_BITS;

struct Slot<T> {
    generation: u64,
    entry: Option<T>,
}

/// Registrations by token, every token carrying `TAG` (0 or [`CHILD`]).
pub(crate) struct Tokens<T, const TAG: u64> {
    slots: Vec<Slot<T>>,
    free: Vec<u32>,
}

/// The shard's IO sources.
pub(crate) type Sources = Tokens<Arc<Source>, 0>;

impl<T, const TAG: u64> Default for Tokens<T, TAG> {
    fn default() -> Self {
        Self {
            slots: Vec::new(),
            free: Vec::new(),
        }
    }
}

impl<T, const TAG: u64> Tokens<T, TAG> {
    fn token(index: u32, generation: u64) -> u64 {
        TAG | (generation << INDEX_BITS) | u64::from(index)
    }

    /// The index of the slot `token` names, if it is of this table and of
    /// the slot's generation.
    fn index_of(&self, token: u64) -> Option<usize> {
        if token & CHILD != TAG {
            return None;
        }
        let index = (token & INDEX_MASK) as usize;
        let slot = self.slots.get(index)?;
        (slot.generation == (token & !CHILD) >> INDEX_BITS).then_some(index)
    }

    /// Stores `entry` and returns its token.
    pub(crate) fn insert(&mut self, entry: T) -> io::Result<u64> {
        if let Some(index) = self.free.pop() {
            let slot = &mut self.slots[index as usize];
            slot.entry = Some(entry);
            return Ok(Self::token(index, slot.generation));
        }
        let index = self.slots.len();
        // The last index is never handed out, so no token is all ones (the
        // pollers' own wake and "no registration" use that).
        if index as u64 >= INDEX_MASK {
            return Err(io::Error::new(
                io::ErrorKind::OutOfMemory,
                "the shard watches as many sources as a token can name",
            ));
        }
        self.slots.push(Slot {
            generation: 0,
            entry: Some(entry),
        });
        Ok(Self::token(index as u32, 0))
    }

    /// The entry registered under `token`, if that registration is still
    /// in place.
    #[cfg(any(readiness, test))]
    pub(crate) fn get(&self, token: u64) -> Option<&T> {
        self.slots[self.index_of(token)?].entry.as_ref()
    }

    /// The entry registered under `token`, to change.
    #[cfg(port)]
    pub(crate) fn get_mut(&mut self, token: u64) -> Option<&mut T> {
        let index = self.index_of(token)?;
        self.slots[index].entry.as_mut()
    }

    /// Ends the registration under `token`; a stale token, or one of the
    /// other table, does nothing.
    pub(crate) fn remove(&mut self, token: u64) -> Option<T> {
        let index = self.index_of(token)?;
        let slot = &mut self.slots[index];
        let entry = slot.entry.take()?;
        slot.generation = (slot.generation + 1) & GEN_MASK;
        self.free.push(index as u32);
        Some(entry)
    }

    /// Takes every registered entry out.
    pub(crate) fn drain(&mut self) -> impl Iterator<Item = T> + '_ {
        self.drain_tokens().map(|(_, entry)| entry)
    }

    /// Takes every registered entry out, each with its token.
    pub(crate) fn drain_tokens(&mut self) -> impl Iterator<Item = (u64, T)> + '_ {
        self.free.clear();
        self.slots
            .drain(..)
            .enumerate()
            .filter_map(|(index, slot)| {
                let token = Self::token(index as u32, slot.generation);
                slot.entry.map(|entry| (token, entry))
            })
    }

    /// How many entries are registered.
    #[cfg(test)]
    pub(crate) fn len(&self) -> usize {
        self.slots.len() - self.free.len()
    }
}

#[cfg(all(test, not(loom)))]
mod tests {
    use super::{CHILD, INDEX_BITS, Sources, Tokens};
    use crate::io::Source;
    use crate::io::source::Kind;

    fn source() -> std::sync::Arc<Source> {
        Source::new(crate::io::raw::NONE, Kind::Bytes)
    }

    #[test]
    fn a_removed_registration_is_not_found_by_its_old_token() {
        let mut s = Sources::default();
        let a = s.insert(source()).unwrap();
        assert!(s.get(a).is_some());
        assert!(s.remove(a).is_some());
        assert!(s.get(a).is_none());
        // The slot is reused under a new generation.
        let b = s.insert(source()).unwrap();
        assert_eq!(a & ((1 << INDEX_BITS) - 1), b & ((1 << INDEX_BITS) - 1));
        assert_ne!(a, b);
        assert!(s.get(a).is_none(), "the old token names nothing");
        assert!(s.remove(a).is_none(), "nor removes the new registration");
        assert!(s.get(b).is_some());
        assert_eq!(s.len(), 1);
    }

    #[test]
    fn tokens_fit_a_pointer_and_are_never_all_ones() {
        let mut s = Sources::default();
        let mut c = Tokens::<u8, CHILD>::default();
        for _ in 0..1000 {
            let t = s.insert(source()).unwrap();
            let u = c.insert(0).unwrap();
            for token in [t, u] {
                assert!(token <= usize::MAX as u64);
                assert_ne!(token, u64::MAX);
            }
            s.remove(t);
            c.remove(u);
        }
    }

    /// A token of one table finds nothing in the other, even where index
    /// and generation agree.
    #[test]
    fn a_token_of_one_table_finds_nothing_in_the_other() {
        let mut s = Sources::default();
        let mut c = Tokens::<u8, CHILD>::default();
        let t = s.insert(source()).unwrap();
        let u = c.insert(7).unwrap();
        assert_eq!(t & !CHILD, u & !CHILD, "the same index and generation");
        assert!(s.get(u).is_none());
        assert!(c.get(t).is_none());
        assert!(s.remove(u).is_none() && c.remove(t).is_none());
        assert_eq!((s.len(), c.len()), (1, 1));
    }
}