low_latency_data_structures/spsc/

queue.rs

use crossbeam_utils::CachePadded;

use crate::mem::{Allocation, Allocator};
use crate::shim::sync::atomic::AtomicUsize;
use crate::shim::sync::{Arc, atomic};
use crate::spsc::builder::Options;
use crate::spsc::consumer::Consumer;
use crate::spsc::producer::Producer;

/// Creates a new SPSC queue with `CAPACITY` slots using the specified allocator.
///
/// `CAPACITY` must be a power of two (compile-time enforced). The slot
/// buffer is allocated up front and locked into RAM via `mlock` so the
/// hot path never page-faults.
///
/// # Examples
///
/// ```
/// use low_latency_data_structures::spsc::{self, new};
/// use low_latency_data_structures::mem::global::GlobalAllocator;
///
/// let (mut producer, mut consumer) = new::<u64, 8, GlobalAllocator>(
///     spsc::Options::global_mlocked(),
/// );
/// assert_eq!(producer.push(1), None);
/// assert_eq!(consumer.pop(), Some(1));
/// ```
///
/// Capacities that are not powers of two fail to compile:
///
/// ```compile_fail
/// use low_latency_data_structures::spsc::{self, new};
/// use low_latency_data_structures::mem::global::GlobalAllocator;
/// # use seq_macro::seq;
/// seq!(N in 2..20 {
///     {
///         const CAP: usize = 2usize.wrapping_pow(N);
///         let _fail = new::<u64, { CAP - 1 }, GlobalAllocator>(
///             spsc::Options::global_mlocked(),
///         );
///         let _fail = new::<u64, { CAP + 1 }, GlobalAllocator>(
///             spsc::Options::global_mlocked(),
///         );
///     }
/// });
/// ```
pub fn new<T, const CAPACITY: usize, A>(
    options: Options<A>,
) -> (Producer<T, CAPACITY, A>, Consumer<T, CAPACITY, A>)
where
    A: Allocator,
{
    const {
        assert!(
            CAPACITY.is_power_of_two(),
            "Given capacity is not a power of two!"
        );
    };
    let slots_allocation = A::allocate(CAPACITY, options.alloc);
    let q = Arc::new(Queue {
        head: AtomicUsize::new(0).into(),
        tail: AtomicUsize::new(0).into(),
        slots_allocation,
    });
    let producer = Producer::new(q.clone());
    let consumer = Consumer::new(q);
    (producer, consumer)
}

pub(super) struct Queue<T, const CAPACITY: usize, A>
where
    A: Allocator,
{
    pub(super) slots_allocation: A::Allocation<T>,
    pub(super) head: CachePadded<AtomicUsize>,
    pub(super) tail: CachePadded<AtomicUsize>,
}

impl<T, const CAPACITY: usize, A> Drop for Queue<T, CAPACITY, A>
where
    A: Allocator,
{
    fn drop(&mut self) {
        let head = self.head.load(atomic::Ordering::Relaxed);
        let tail = self.tail.load(atomic::Ordering::Relaxed);
        let count = tail.wrapping_sub(head);
        for k in 0..count {
            let i = head.wrapping_add(k);
            let slot_ptr = self.slots_allocation.ptr().wrapping_add(i & (CAPACITY - 1));
            // SAFETY: it's not null because `i & (sef.capacity - 1)` limits it to [0;
            // allocated_cap). And it can be safely converted to a reference because T is self
            // contained bitwise.
            unsafe {
                slot_ptr.as_ref_unchecked().with_mut(|ptr| {
                    ptr.as_mut_unchecked().assume_init_drop();
                })
            }
        }
    }
}

#[cfg(test)]
#[cfg(feature = "tests_basic")]
mod tests_basic {
    use std::cell::Cell;
    use std::rc::Rc;
    use std::thread;

    use seq_macro::seq;

    use super::new;
    #[cfg(not(feature = "tests_hugepage"))]
    use crate::mem::global::{GlobalAllocator, GlobalAllocatorOptions};
    #[cfg(feature = "tests_hugepage")]
    use crate::mem::hugepages::{HugepageAllocator, HugepageAllocatorOptions, HugepageSize};
    use crate::spsc::builder::Options;

    #[cfg(not(feature = "tests_hugepage"))]
    fn alloc_options() -> GlobalAllocatorOptions {
        GlobalAllocatorOptions::builder().mlock(true).build()
    }

    #[cfg(not(feature = "tests_hugepage"))]
    type Alloc = GlobalAllocator;
    #[cfg(feature = "tests_hugepage")]
    type Alloc = HugepageAllocator;

    #[cfg(feature = "tests_hugepage")]
    fn alloc_options() -> HugepageAllocatorOptions {
        HugepageAllocatorOptions::builder()
            .mlock(true)
            .hugepage_size(HugepageSize::H2MB)
            .build()
    }

    fn spsc_options() -> Options<Alloc> {
        Options::builder().alloc(alloc_options()).build()
    }

    #[test]
    fn move_producer_consumer_to_threads() {
        let (mut producer, mut consumer) = new::<_, 2, Alloc>(spsc_options());
        thread::spawn(move || {
            let _ = producer.push(123);
        })
        .join()
        .unwrap();
        thread::spawn(move || {
            assert_eq!(consumer.pop(), Some(123));
        })
        .join()
        .unwrap();
    }

    #[test]
    fn handoff_one_value() {
        let (mut producer, mut consumer) = new::<_, 2, Alloc>(spsc_options());
        assert_eq!(producer.push(123), None);
        assert_eq!(consumer.pop(), Some(123));
    }

    #[test]
    fn allows_queues_with_powers_of_two_capacity() {
        #[cfg(feature = "tests_hugepage")]
        seq!(N in 0..6 {
            {
                const POWER: usize = 2usize.pow(N);
                new::<u32, POWER, Alloc>(spsc_options());
            }
        });
        #[cfg(not(feature = "tests_hugepage"))]
        seq!(N in 0..20 {
            {
                const POWER: usize = 2usize.pow(N);
                new::<u32, POWER, Alloc>(spsc_options());
            }
        });
    }

    #[test]
    fn drops_unread_items() {
        let counter = Rc::new(Cell::new(0));
        #[derive(Debug, PartialEq)]
        struct Droppable {
            counter: Rc<Cell<usize>>,
        }
        impl Drop for Droppable {
            fn drop(&mut self) {
                let cnt = self.counter.get();
                self.counter.set(cnt + 1);
            }
        }
        const CAPACITY: usize = 64;
        let (mut producer, mut consumer) = new::<_, CAPACITY, Alloc>(spsc_options());
        for _ in 0..CAPACITY {
            let counter = counter.clone();
            assert_eq!(producer.push(Droppable { counter }), None);
        }
        let read = CAPACITY / 2;
        for _ in 0..read {
            assert!(matches!(consumer.pop(), Some(_)));
        }
        assert_eq!(read, counter.get());
        drop(producer);
        drop(consumer);
        assert_eq!(CAPACITY, counter.get());
    }

    #[test]
    fn empty_returns_none() {
        let (_, mut consumer) = new::<i32, 4, Alloc>(spsc_options());
        assert_eq!(consumer.pop(), None);
        assert_eq!(consumer.pop(), None);
    }

    #[test]
    fn full_returns_item_back() {
        let (mut producer, _) = new::<i32, 2, Alloc>(spsc_options());
        assert_eq!(producer.push(1), None);
        assert_eq!(producer.push(2), None);
        // Full: item comes back untouched
        assert_eq!(producer.push(3), Some(3));
        assert_eq!(producer.push(4), Some(4));
    }

    #[test]
    fn fifo_ordering() -> anyhow::Result<()> {
        let (mut producer, mut consumer) = new::<_, 8, Alloc>(spsc_options());
        for i in 0..8 {
            assert_eq!(producer.push(i), None);
        }
        for i in 0..8 {
            assert_eq!(consumer.pop(), Some(i));
        }
        Ok(())
    }

    #[test]
    fn wraparound_n_laps() {
        const CAPACITY: usize = 4;
        let laps = 100;
        let (mut producer, mut consumer) = new::<_, CAPACITY, Alloc>(spsc_options());
        for lap in 0..laps {
            for i in 0..CAPACITY {
                let val = lap * CAPACITY + i;
                assert_eq!(producer.push(val), None, "push failed at lap {lap}, i {i}");
            }
            // Queue is full
            assert_eq!(producer.push(9999), Some(9999));
            for i in 0..CAPACITY {
                let val = lap * CAPACITY + i;
                assert_eq!(consumer.pop(), Some(val), "wrong value at lap {lap}, i {i}");
            }
            // Queue is empty
            assert_eq!(consumer.pop(), None);
        }
    }

    #[test]
    fn interleaved_push_pop() {
        let (mut producer, mut consumer) = new::<_, 2, Alloc>(spsc_options());
        // Push 1, pop 1, repeat: tests wraparound with tiny queue
        for i in 0..1000 {
            assert_eq!(producer.push(i), None);
            assert_eq!(consumer.pop(), Some(i));
        }
    }

    #[test]
    fn capacity_one() {
        let (mut producer, mut consumer) = new::<_, 1, Alloc>(spsc_options());
        assert_eq!(consumer.pop(), None);
        assert_eq!(producer.push(42), None);
        assert_eq!(producer.push(43), Some(43)); // full
        assert_eq!(consumer.pop(), Some(42));
        assert_eq!(consumer.pop(), None); // empty again
    }

    #[test]
    fn move_only_type() {
        // Verify non-Copy, non-Clone types work
        let (mut producer, mut consumer) = new::<_, 4, Alloc>(spsc_options());
        let s = String::from("hello");
        assert_eq!(producer.push(s), None);
        let got = consumer.pop().unwrap();
        assert_eq!(got, "hello");
    }
}

#[cfg(test)]
#[cfg(feature = "tests_loom")]
mod tests_loom {
    use super::*;
    use crate::mem::loom::{LoomVecAllocator, LoomVecAllocatorOptions};

    static_assertions::assert_cfg!(
        not(feature = "tests_hugepage"),
        "tests_loom incompatible with tests_hugepage as loom uses custom buffer allocator",
    );

    fn spsc_options() -> Options<LoomVecAllocator> {
        Options::builder().alloc(LoomVecAllocatorOptions).build()
    }

    #[test]
    fn concurrent_push_pop() {
        loom::model(|| {
            let (mut producer, mut consumer) = new::<i32, 4, _>(spsc_options());

            let t1 = loom::thread::spawn(move || {
                let _ = producer.push(1);
                let _ = producer.push(2);
                let _ = producer.push(3);
            });

            let t2 = loom::thread::spawn(move || {
                let mut collected = Vec::new();
                while collected.len() < 3 {
                    if let Some(v) = consumer.pop() {
                        collected.push(v);
                    } else {
                        loom::thread::yield_now();
                    }
                }
                collected
            });

            t1.join().unwrap();
            let values = t2.join().unwrap();
            // FIFO: must be in order
            assert_eq!(values, vec![1, 2, 3]);
        });
    }

    /// Producer must retry if the queue is full.
    #[test]
    fn concurrent_with_full_queue() {
        loom::model(|| {
            let (mut producer, mut consumer) = new::<i32, 1, _>(spsc_options());

            let t1 = loom::thread::spawn(move || {
                for i in 0..3 {
                    while producer.push(i).is_some() {
                        loom::thread::yield_now();
                    }
                }
            });

            let t2 = loom::thread::spawn(move || {
                let mut collected = Vec::new();
                while collected.len() < 3 {
                    if let Some(v) = consumer.pop() {
                        collected.push(v);
                    } else {
                        loom::thread::yield_now();
                    }
                }
                collected
            });

            t1.join().unwrap();
            let values = t2.join().unwrap();
            assert_eq!(values, vec![0, 1, 2]);
        });
    }

    #[test]
    fn concurrent_no_values_lost() {
        loom::model(|| {
            let (mut producer, mut consumer) = new::<i32, 2, _>(spsc_options());

            let t1 = loom::thread::spawn(move || {
                for i in 0..3 {
                    while producer.push(i).is_some() {
                        loom::thread::yield_now();
                    }
                }
            });

            let t2 = loom::thread::spawn(move || {
                let mut sum = 0;
                let mut count = 0;
                while count < 3 {
                    if let Some(v) = consumer.pop() {
                        sum += v;
                        count += 1;
                    } else {
                        loom::thread::yield_now();
                    }
                }
                (count, sum)
            });

            t1.join().unwrap();
            let (count, sum) = t2.join().unwrap();
            assert_eq!(count, 3);
            assert_eq!(sum, 0 + 1 + 2);
        });
    }
}

#[cfg(test)]
#[cfg(feature = "tests_dhat")]
mod tests_dhat {
    use super::*;
    use crate::mem::global::GlobalAllocator;

    #[test]
    fn hot_path_zero_allocations() {
        let _profiler = dhat::Profiler::builder().testing().build();
        let (mut producer, mut consumer) =
            new::<u64, 1024, GlobalAllocator>(Options::global_mlocked());

        // Warm up to absorb any one-time platform allocations: lazy symbol
        // resolution in the dynamic linker, libstd TLS init, debug-build
        // slow-path stubs that LLVM emits. None of these scale with iteration
        // count; the per-iteration regression we actually care about would
        // dwarf the slack budget below.
        for i in 0..1024u64 {
            let _ = producer.push(i);
        }
        for _ in 0..1024 {
            let _ = consumer.pop();
        }
        let _ = dhat::HeapStats::get();

        let before = dhat::HeapStats::get();

        for i in 0..1024u64 {
            let _ = producer.push(i);
        }
        for _ in 0..1024 {
            let _ = consumer.pop();
        }

        let after = dhat::HeapStats::get();
        let allocs = after.total_blocks - before.total_blocks;
        dhat::assert!(
            allocs == 0,
            "hot path allocated {} blocks; a real regression would be O(iters)",
            allocs,
        );
    }
}