fastars 0.1.0

//! Lock-free SPSC (Single-Producer Single-Consumer) queue.
//!
//! This module provides a bounded, lock-free queue optimized for
//! the single-producer single-consumer pattern. Each worker gets
//! its own queue, eliminating contention.
//!
//! # Design
//!
//! The queue uses a ring buffer with power-of-2 capacity for fast modulo
//! operations (bitwise AND instead of division). Head and tail positions
//! use wrapping arithmetic to allow unlimited producer/consumer progression.
//!
//! # Safety
//!
//! This queue is safe for exactly one producer thread and one consumer thread.
//! Using it with multiple producers or consumers leads to undefined behavior.
//!
//! # Example
//!
//! ```
//! use fastars::pipeline::spsc_queue::spsc_channel;
//!
//! let (producer, consumer) = spsc_channel::<i32>(16);
//!
//! // Producer pushes items
//! producer.push(42);
//!
//! // Consumer pops items
//! assert_eq!(consumer.pop(), Some(42));
//! ```

use std::cell::UnsafeCell;
use std::mem::MaybeUninit;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

/// A bounded, lock-free SPSC queue.
///
/// This queue is safe for exactly one producer thread and one consumer thread.
/// Using it with multiple producers or consumers leads to undefined behavior.
pub struct SpscQueue<T> {
    /// Ring buffer storage
    buffer: Box<[UnsafeCell<MaybeUninit<T>>]>,
    /// Capacity of the queue (power of 2)
    capacity: usize,
    /// Mask for modulo operation (capacity - 1)
    mask: usize,
    /// Head position (consumer reads from here)
    head: AtomicUsize,
    /// Tail position (producer writes here)
    tail: AtomicUsize,
}

// Safety: SpscQueue is Send if T is Send
// Only one thread reads (consumer) and one writes (producer)
unsafe impl<T: Send> Send for SpscQueue<T> {}
unsafe impl<T: Send> Sync for SpscQueue<T> {}

impl<T> SpscQueue<T> {
    /// Create a new SPSC queue with the given capacity.
    ///
    /// Capacity will be rounded up to the next power of 2.
    ///
    /// # Panics
    ///
    /// Panics if capacity is 0.
    pub fn new(capacity: usize) -> Self {
        assert!(capacity > 0, "capacity must be greater than 0");

        // Round up to power of 2
        let capacity = capacity.next_power_of_two();
        let mask = capacity - 1;

        // Allocate uninitialized buffer
        let buffer: Vec<UnsafeCell<MaybeUninit<T>>> = (0..capacity)
            .map(|_| UnsafeCell::new(MaybeUninit::uninit()))
            .collect();

        Self {
            buffer: buffer.into_boxed_slice(),
            capacity,
            mask,
            head: AtomicUsize::new(0),
            tail: AtomicUsize::new(0),
        }
    }

    /// Try to push an item to the queue.
    ///
    /// Returns `Ok(())` if successful, `Err(item)` if the queue is full.
    ///
    /// # Safety
    ///
    /// Must be called from exactly one producer thread.
    #[inline]
    pub fn push(&self, item: T) -> Result<(), T> {
        let tail = self.tail.load(Ordering::Relaxed);
        let head = self.head.load(Ordering::Acquire);

        // Check if queue is full
        if tail.wrapping_sub(head) >= self.capacity {
            return Err(item);
        }

        // Write item
        let idx = tail & self.mask;
        unsafe {
            (*self.buffer[idx].get()).write(item);
        }

        // Publish
        self.tail.store(tail.wrapping_add(1), Ordering::Release);
        Ok(())
    }

    /// Blocking push - spins until space is available.
    ///
    /// # Safety
    ///
    /// Must be called from exactly one producer thread.
    #[inline]
    pub fn push_blocking(&self, mut item: T) {
        loop {
            match self.push(item) {
                Ok(()) => return,
                Err(returned) => {
                    item = returned;
                    std::hint::spin_loop();
                }
            }
        }
    }

    /// Try to pop an item from the queue.
    ///
    /// Returns `Some(item)` if successful, `None` if the queue is empty.
    ///
    /// # Safety
    ///
    /// Must be called from exactly one consumer thread.
    #[inline]
    pub fn pop(&self) -> Option<T> {
        let head = self.head.load(Ordering::Relaxed);
        let tail = self.tail.load(Ordering::Acquire);

        // Check if queue is empty
        if head == tail {
            return None;
        }

        // Read item
        let idx = head & self.mask;
        let item = unsafe { (*self.buffer[idx].get()).assume_init_read() };

        // Publish
        self.head.store(head.wrapping_add(1), Ordering::Release);
        Some(item)
    }

    /// Check if the queue is empty.
    #[inline]
    pub fn is_empty(&self) -> bool {
        let head = self.head.load(Ordering::Acquire);
        let tail = self.tail.load(Ordering::Acquire);
        head == tail
    }

    /// Check if the queue is full.
    #[inline]
    pub fn is_full(&self) -> bool {
        let head = self.head.load(Ordering::Acquire);
        let tail = self.tail.load(Ordering::Acquire);
        tail.wrapping_sub(head) >= self.capacity
    }

    /// Get the current number of items in the queue.
    #[inline]
    pub fn len(&self) -> usize {
        let head = self.head.load(Ordering::Acquire);
        let tail = self.tail.load(Ordering::Acquire);
        tail.wrapping_sub(head)
    }

    /// Get the capacity of the queue.
    #[inline]
    pub fn capacity(&self) -> usize {
        self.capacity
    }
}

impl<T> Drop for SpscQueue<T> {
    fn drop(&mut self) {
        // Drop any remaining items
        while self.pop().is_some() {}
    }
}

/// A handle to the producer side of an SPSC queue.
pub struct SpscProducer<T> {
    queue: Arc<SpscQueue<T>>,
}

impl<T> SpscProducer<T> {
    /// Push an item, blocking if the queue is full.
    #[inline]
    pub fn push(&self, item: T) {
        self.queue.push_blocking(item);
    }

    /// Try to push without blocking.
    #[inline]
    pub fn try_push(&self, item: T) -> Result<(), T> {
        self.queue.push(item)
    }

    /// Check if the queue is full.
    #[inline]
    pub fn is_full(&self) -> bool {
        self.queue.is_full()
    }

    /// Get the current number of items in the queue.
    #[inline]
    pub fn len(&self) -> usize {
        self.queue.len()
    }

    /// Get the capacity of the queue.
    #[inline]
    pub fn capacity(&self) -> usize {
        self.queue.capacity()
    }
}

/// A handle to the consumer side of an SPSC queue.
pub struct SpscConsumer<T> {
    queue: Arc<SpscQueue<T>>,
}

impl<T> SpscConsumer<T> {
    /// Pop an item if available.
    #[inline]
    pub fn pop(&self) -> Option<T> {
        self.queue.pop()
    }

    /// Check if empty.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.queue.is_empty()
    }

    /// Get the current number of items in the queue.
    #[inline]
    pub fn len(&self) -> usize {
        self.queue.len()
    }

    /// Get the capacity of the queue.
    #[inline]
    pub fn capacity(&self) -> usize {
        self.queue.capacity()
    }
}

/// Create a new SPSC channel pair.
///
/// Returns a producer and consumer handle that can be used from separate threads.
/// The producer can push items and the consumer can pop them.
///
/// # Arguments
///
/// * `capacity` - The maximum number of items the queue can hold. Will be
///   rounded up to the next power of 2.
///
/// # Example
///
/// ```
/// use fastars::pipeline::spsc_queue::spsc_channel;
/// use std::thread;
///
/// let (producer, consumer) = spsc_channel::<i32>(16);
///
/// let producer_thread = thread::spawn(move || {
///     for i in 0..10 {
///         producer.push(i);
///     }
/// });
///
/// let consumer_thread = thread::spawn(move || {
///     let mut sum = 0;
///     let mut count = 0;
///     while count < 10 {
///         if let Some(v) = consumer.pop() {
///             sum += v;
///             count += 1;
///         } else {
///             thread::yield_now();
///         }
///     }
///     sum
/// });
///
/// producer_thread.join().unwrap();
/// let sum = consumer_thread.join().unwrap();
/// assert_eq!(sum, 45); // 0+1+2+...+9 = 45
/// ```
pub fn spsc_channel<T>(capacity: usize) -> (SpscProducer<T>, SpscConsumer<T>) {
    let queue = Arc::new(SpscQueue::new(capacity));
    (
        SpscProducer {
            queue: queue.clone(),
        },
        SpscConsumer { queue },
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::thread;

    #[test]
    fn test_spsc_basic() {
        let queue = SpscQueue::<i32>::new(4);

        assert!(queue.is_empty());
        assert!(!queue.is_full());

        queue.push(1).unwrap();
        queue.push(2).unwrap();

        assert_eq!(queue.len(), 2);

        assert_eq!(queue.pop(), Some(1));
        assert_eq!(queue.pop(), Some(2));
        assert_eq!(queue.pop(), None);
    }

    #[test]
    fn test_spsc_full() {
        let queue = SpscQueue::<i32>::new(2);

        queue.push(1).unwrap();
        queue.push(2).unwrap();

        assert!(queue.is_full());
        assert!(queue.push(3).is_err());

        queue.pop();
        queue.push(3).unwrap();
    }

    #[test]
    fn test_spsc_threaded() {
        let (producer, consumer) = spsc_channel::<i32>(16);

        let producer_thread = thread::spawn(move || {
            for i in 0..1000 {
                producer.push(i);
            }
        });

        let consumer_thread = thread::spawn(move || {
            let mut sum = 0i64;
            let mut count = 0;
            while count < 1000 {
                if let Some(v) = consumer.pop() {
                    sum += v as i64;
                    count += 1;
                } else {
                    thread::yield_now();
                }
            }
            sum
        });

        producer_thread.join().unwrap();
        let sum = consumer_thread.join().unwrap();

        // Sum of 0..1000 = 499500
        assert_eq!(sum, 499500);
    }

    #[test]
    fn test_capacity_power_of_two() {
        let queue = SpscQueue::<i32>::new(5);
        assert_eq!(queue.capacity(), 8); // Rounded up to 8

        let queue = SpscQueue::<i32>::new(8);
        assert_eq!(queue.capacity(), 8); // Already power of 2
    }

    #[test]
    fn test_spsc_wrap_around() {
        // Test that wrapping arithmetic works correctly
        let queue = SpscQueue::<i32>::new(4);

        // Fill and drain multiple times to force wrap-around
        for round in 0..10 {
            for i in 0..4 {
                queue.push(round * 4 + i).unwrap();
            }
            assert!(queue.is_full());

            for i in 0..4 {
                assert_eq!(queue.pop(), Some(round * 4 + i));
            }
            assert!(queue.is_empty());
        }
    }

    #[test]
    fn test_spsc_producer_consumer_handles() {
        let (producer, consumer) = spsc_channel::<String>(8);

        producer.push("hello".to_string());
        producer.push("world".to_string());

        assert_eq!(producer.len(), 2);
        assert_eq!(consumer.len(), 2);

        assert_eq!(consumer.pop(), Some("hello".to_string()));
        assert_eq!(consumer.pop(), Some("world".to_string()));
        assert!(consumer.is_empty());
    }

    #[test]
    fn test_spsc_drop_remaining() {
        use std::sync::atomic::{AtomicUsize, Ordering};

        static DROP_COUNT: AtomicUsize = AtomicUsize::new(0);

        #[derive(Debug)]
        struct DropCounter;

        impl Drop for DropCounter {
            fn drop(&mut self) {
                DROP_COUNT.fetch_add(1, Ordering::SeqCst);
            }
        }

        DROP_COUNT.store(0, Ordering::SeqCst);

        {
            let queue = SpscQueue::<DropCounter>::new(4);
            queue.push(DropCounter).unwrap();
            queue.push(DropCounter).unwrap();
            queue.push(DropCounter).unwrap();
            // Queue dropped here with 3 items remaining
        }

        assert_eq!(DROP_COUNT.load(Ordering::SeqCst), 3);
    }

    #[test]
    fn test_spsc_high_throughput() {
        let (producer, consumer) = spsc_channel::<u64>(1024);

        let items = 100_000u64;

        let producer_thread = thread::spawn(move || {
            for i in 0..items {
                producer.push(i);
            }
        });

        let consumer_thread = thread::spawn(move || {
            let mut received = 0u64;
            let mut expected = 0u64;
            while received < items {
                if let Some(v) = consumer.pop() {
                    assert_eq!(v, expected, "Items received out of order");
                    expected += 1;
                    received += 1;
                } else {
                    thread::yield_now();
                }
            }
            received
        });

        producer_thread.join().unwrap();
        let received = consumer_thread.join().unwrap();

        assert_eq!(received, items);
    }

    #[test]
    fn test_try_push() {
        let (producer, _consumer) = spsc_channel::<i32>(2);

        assert!(producer.try_push(1).is_ok());
        assert!(producer.try_push(2).is_ok());
        assert!(producer.try_push(3).is_err()); // Queue full
    }

    #[test]
    #[should_panic(expected = "capacity must be greater than 0")]
    fn test_zero_capacity_panics() {
        let _queue = SpscQueue::<i32>::new(0);
    }
}