path: root/tokio/src/util/slab.rs

#![cfg_attr(not(feature = "rt"), allow(dead_code))]

use crate::loom::cell::UnsafeCell;
use crate::loom::sync::atomic::{AtomicBool, AtomicUsize};
use crate::loom::sync::{Arc, Mutex};
use crate::util::bit;
use std::fmt;
use std::mem;
use std::ops;
use std::ptr;
use std::sync::atomic::Ordering::Relaxed;

/// Amortized allocation for homogeneous data types.
///
/// The slab pre-allocates chunks of memory to store values. It uses a similar
/// growing strategy as `Vec`. When new capacity is needed, the slab grows by
/// 2x.
///
/// # Pages
///
/// Unlike `Vec`, growing does not require moving existing elements. Instead of
/// being a continuous chunk of memory for all elements, `Slab` is an array of
/// arrays. The top-level array is an array of pages. Each page is 2x bigger
/// than the previous one. When the slab grows, a new page is allocated.
///
/// Pages are lazily initialized.
///
/// # Allocating
///
/// When allocating an object, first previously used slots are reused. If no
/// previously used slot is available, a new slot is initialized in an existing
/// page. If all pages are full, then a new page is allocated.
///
/// When an allocated object is released, it is pushed into it's page's free
/// list. Allocating scans all pages for a free slot.
///
/// # Indexing
///
/// The slab is able to index values using an address. Even when the indexed
/// object has been released, it is still safe to index. This is a key ability
/// for using the slab with the I/O driver. Addresses are registered with the
/// OS's selector and I/O resources can be released without synchronizing with
/// the OS.
///
/// # Compaction
///
/// `Slab::compact` will release pages that have been allocated but are no
/// longer used. This is done by scanning the pages and finding pages with no
/// allocated objects. These pages are then freed.
///
/// # Synchronization
///
/// The `Slab` structure is able to provide (mostly) unsynchronized reads to
/// values stored in the slab. Insertions and removals are synchronized. Reading
/// objects via `Ref` is fully unsynchronized. Indexing objects uses amortized
/// synchronization.
///
pub(crate) struct Slab<T> {
    /// Array of pages. Each page is synchronized.
    pages: [Arc<Page<T>>; NUM_PAGES],

    /// Caches the array pointer & number of initialized slots.
    cached: [CachedPage<T>; NUM_PAGES],
}

/// Allocate values in the associated slab.
pub(crate) struct Allocator<T> {
    /// Pages in the slab. The first page has a capacity of 16 elements. Each
    /// following page has double the capacity of the previous page.
    ///
    /// Each returned `Ref` holds a reference count to this `Arc`.
    pages: [Arc<Page<T>>; NUM_PAGES],
}

/// References a slot in the slab. Indexing a slot using an `Address` is memory
/// safe even if the slot has been released or the page has been deallocated.
/// However, it is not guaranteed that the slot has not been reused and is now
/// represents a different value.
///
/// The I/O driver uses a counter to track the slot's generation. Once accessing
/// the slot, the generations are compared. If they match, the value matches the
/// address.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub(crate) struct Address(usize);

/// An entry in the slab.
pub(crate) trait Entry: Default {
    /// Reset the entry's value and track the generation.
    fn reset(&self);
}

/// A reference to a value stored in the slab
pub(crate) struct Ref<T> {
    value: *const Value<T>,
}

/// Maximum number of pages a slab can contain.
const NUM_PAGES: usize = 19;

/// Minimum number of slots a page can contain.
const PAGE_INITIAL_SIZE: usize = 32;
const PAGE_INDEX_SHIFT: u32 = PAGE_INITIAL_SIZE.trailing_zeros() + 1;

/// A page in the slab
struct Page<T> {
    /// Slots
    slots: Mutex<Slots<T>>,

    // Number of slots currently being used. This is not guaranteed to be up to
    // date and should only be used as a hint.
    used: AtomicUsize,

    // Set to `true` when the page has been allocated.
    allocated: AtomicBool,

    // The number of slots the page can hold.
    len: usize,

    // Length of all previous pages combined
    prev_len: usize,
}

struct CachedPage<T> {
    /// Pointer to the page's slots.
    slots: *const Slot<T>,

    /// Number of initialized slots.
    init: usize,
}

/// Page state
struct Slots<T> {
    /// Slots
    slots: Vec<Slot<T>>,

    head: usize,

    /// Number of slots currently in use.
    used: usize,
}

unsafe impl<T: Sync> Sync for Page<T> {}
unsafe impl<T: Sync> Send for Page<T> {}
unsafe impl<T: Sync> Sync for CachedPage<T> {}
unsafe impl<T: Sync> Send for CachedPage<T> {}
unsafe impl<T: Sync> Sync for Ref<T> {}
unsafe impl<T: Sync> Send for Ref<T> {}

/// A slot in the slab. Contains slot-specific metadata.
///
/// `#[repr(C)]` guarantees that the struct starts w/ `value`. We use pointer
/// math to map a value pointer to an index in the page.
#[repr(C)]
struct Slot<T> {
    /// Pointed to by `Ref`.
    value: UnsafeCell<Value<T>>,

    /// Next entry in the free list.
    next: u32,
}

/// Value paired with a reference to the page
struct Value<T> {
    /// Value stored in the value
    value: T,

    /// Pointer to the page containing the slot.
    ///
    /// A raw pointer is used as this creates a ref cycle.
    page: *const Page<T>,
}

impl<T> Slab<T> {
    /// Create a new, empty, slab
    pub(crate) fn new() -> Slab<T> {
        // Initializing arrays is a bit annoying. Instead of manually writing
        // out an array and every single entry, `Default::default()` is used to
        // initialize the array, then the array is iterated and each value is
        // initialized.
        let mut slab = Slab {
            pages: Default::default(),
            cached: Default::default(),
        };

        let mut len = PAGE_INITIAL_SIZE;
        let mut prev_len: usize = 0;

        for page in &mut slab.pages {
            let page = Arc::get_mut(page).unwrap();
            page.len = len;
            page.prev_len = prev_len;
            len *= 2;
            prev_len += page.len;

            // Ensure we don't exceed the max address space.
            debug_assert!(
                page.len - 1 + page.prev_len < (1 << 24),
                "max = {:b}",
                page.len - 1 + page.prev_len
            );
        }

        slab
    }

    /// Returns a new `Allocator`.
    ///
    /// The `Allocator` supports concurrent allocation of objects.
    pub(crate) fn allocator(&self) -> Allocator<T> {
        Allocator {
            pages: self.pages.clone(),
        }
    }

    /// Returns a reference to the value stored at the given address.
    ///
    /// `&mut self` is used as the call may update internal cached state.
    pub(crate) fn get(&mut self, addr: Address) -> Option<&T> {
        let page_idx = addr.page();
        let slot_idx = self.pages[page_idx].slot(addr);

        // If the address references a slot that was last seen as uninitialized,
        // the `CachedPage` is updated. This requires acquiring the page lock
        // and updating the slot pointer and initialized offset.
        if self.cached[page_idx].init <= slot_idx {
            self.cached[