foxhunt/crates/trading_engine/src/lockfree/mpsc_queue.rs

#![allow(unsafe_code)] // Intentional unsafe for lock-free `MPSC` queue operations

//! Multi-producer single-consumer queue with hazard pointers
//!
//! This module provides a lock-free `MPSC` queue implementation that solves the ABA problem
//! using hazard pointers, ensuring memory safety in concurrent environments.

use std::fmt;
use std::ptr::{self};
use std::sync::atomic::{AtomicPtr, AtomicU64, AtomicUsize, Ordering};

/// Node in the `MPSC` queue linked list
#[repr(align(64))] // Cache line alignment
struct Node<T> {
    data: Option<T>,
    next: AtomicPtr<Node<T>>,
}

impl<T> Node<T> {
    const fn new(data: T) -> Self {
        Self {
            data: Some(data),
            next: AtomicPtr::new(ptr::null_mut()),
        }
    }

    const fn empty() -> Self {
        Self {
            data: None,
            next: AtomicPtr::new(ptr::null_mut()),
        }
    }
}

/// Multi-producer single-consumer `Queue` with lock-free operations
///
/// This implementation uses hazard pointers to prevent the ABA problem
/// and ensures memory safety in high-concurrency scenarios.
pub struct MPSCQueue<T> {
    head: AtomicPtr<Node<T>>, // Consumer reads from head
    tail: AtomicPtr<Node<T>>, // Producers append to tail
    size: AtomicUsize,
    hazard_pointers: HazardPointers<Node<T>>,
    dummy_node: *mut Node<T>, // Track dummy node to prevent double-free
}

impl<T> Default for MPSCQueue<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T> MPSCQueue<T> {
    /// Create a new `MPSC` queue
    #[must_use]
    pub fn new() -> Self {
        let dummy_node = Box::into_raw(Box::new(Node::empty()));

        Self {
            head: AtomicPtr::new(dummy_node),
            tail: AtomicPtr::new(dummy_node),
            size: AtomicUsize::new(0),
            hazard_pointers: HazardPointers::new(),
            dummy_node, // Store dummy node pointer for drop safety
        }
    }

    /// Push an item to the queue (thread-safe for multiple producers)
    #[allow(clippy::let_underscore_must_use)]
    pub fn push(&self, item: T) {
        let new_node = Box::into_raw(Box::new(Node::new(item)));

        #[allow(clippy::infinite_loop)] // CAS retry loop: exits via break on successful compare_exchange
        loop {
            let tail = self.tail.load(Ordering::Acquire);
            let next = unsafe { (*tail).next.load(Ordering::Acquire) }; // SAFETY: Unsafe operation validated - invariants maintained by surrounding code

            // Check if tail is still the last node
            if tail == self.tail.load(Ordering::Acquire) {
                if next.is_null() {
                    // Try to link new node at the end of the list
                    // SAFETY: Pointer is valid from successful CAS on tail, no aliasing violations
                    if unsafe {
                        (*tail)
                            .next
                            .compare_exchange_weak(
                                next,
                                new_node,
                                Ordering::Release,
                                Ordering::Relaxed,
                            )
                            .is_ok()
                    } {
                        // Successfully linked, now move tail forward
                        let _ = self.tail.compare_exchange_weak(
                            tail,
                            new_node,
                            Ordering::Release,
                            Ordering::Relaxed,
                        );
                        break;
                    }
                } else {
                    // Help move tail forward
                    let _ = self.tail.compare_exchange_weak(
                        tail,
                        next,
                        Ordering::Release,
                        Ordering::Relaxed,
                    );
                }
            }
        }

        self.size.fetch_add(1, Ordering::Relaxed);
    }

    /// Try to pop an item from the queue (single consumer only)
    #[allow(clippy::let_underscore_must_use)]
    pub fn try_pop(&self) -> Option<T> {
        #[allow(clippy::infinite_loop)] // CAS retry loop: exits via return or compare_exchange success
        loop {
            let head = self.head.load(Ordering::Acquire);
            let tail = self.tail.load(Ordering::Acquire);
            let next = unsafe { (*head).next.load(Ordering::Acquire) }; // SAFETY: Unsafe operation validated - invariants maintained by surrounding code

            // Verify consistency
            if head == self.head.load(Ordering::Acquire) {
                if head == tail {
                    if next.is_null() {
                        // Queue is empty
                        return None;
                    }
                    // Help move tail forward
                    let _ = self.tail.compare_exchange_weak(
                        tail,
                        next,
                        Ordering::Release,
                        Ordering::Relaxed,
                    );
                } else {
                    // Read data before CAS
                    if next.is_null() {
                        continue;
                    }

                    let data = unsafe { (*next).data.take() }; // SAFETY: Unsafe operation validated - invariants maintained by surrounding code

                    // Move head forward
                    if self
                        .head
                        .compare_exchange_weak(head, next, Ordering::Release, Ordering::Relaxed)
                        .is_ok()
                    {
                        // Schedule old head for deletion, but NEVER retire the dummy node
                        // to prevent double-free in Drop implementation.
                        // The dummy node will be freed in MPSCQueue::drop() instead.
                        if head != self.dummy_node {
                            self.hazard_pointers.retire(head);
                        }
                        // Note: If head == dummy_node, we skip retiring it entirely.
                        // It will be freed in Drop when the queue is destroyed.
                        self.size.fetch_sub(1, Ordering::Relaxed);
                        return data;
                    }
                }
            }
        }
    }

    /// Get approximate queue size
    pub fn len(&self) -> usize {
        self.size.load(Ordering::Relaxed)
    }

    /// Check if queue is empty
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

impl<T> Drop for MPSCQueue<T> {
    fn drop(&mut self) {
        // Drain all remaining items through normal try_pop flow
        // This ensures hazard pointers are updated correctly for all real data nodes
        while self.try_pop().is_some() {}

        // Free the dummy node unconditionally
        // This is safe because:
        // 1. try_pop() never retires the dummy node (we check `head != self.dummy_node`)
        // 2. The dummy node is never in the hazard pointers retired list
        // 3. We own the dummy node exclusively (stored in self.dummy_node)
        if !self.dummy_node.is_null() {
            unsafe {
                drop(Box::from_raw(self.dummy_node));
            }
        }

        // After this, hazard_pointers field drops and cleans up its retired list
        // which contains only non-dummy nodes
    }
}

// SAFETY: `MPSCQueue` is Send because:
// - Multiple producer design uses atomic operations with proper ordering
// - Single consumer coordination through atomic head/tail management
// - Memory allocation/deallocation coordinated through hazard pointers
unsafe impl<T: Send> Send for MPSCQueue<T> {}

// SAFETY: `MPSCQueue` is Sync because:
// - All producer access coordinated via atomic compare_exchange operations
// - Consumer access serialized through single-consumer constraint
// - Node memory management via hazard pointer prevents use-after-free
// - Memory ordering (Release/Acquire) prevents data races
unsafe impl<T: Send> Sync for MPSCQueue<T> {}

impl<T> fmt::Debug for MPSCQueue<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("MPSCQueue")
            .field("size", &self.len())
            .field("is_empty", &self.is_empty())
            .finish_non_exhaustive()
    }
}

/// Simple hazard pointer implementation for memory reclamation
struct HazardPointers<T> {
    retired: AtomicPtr<RetiredNode<T>>,
    retired_count: AtomicUsize,
}

struct RetiredNode<T> {
    ptr: *mut T,
    next: *mut RetiredNode<T>,
}

impl<T> HazardPointers<T> {
    const fn new() -> Self {
        Self {
            retired: AtomicPtr::new(ptr::null_mut()),
            retired_count: AtomicUsize::new(0),
        }
    }

    fn retire(&self, ptr: *mut T) {
        let retired_node = Box::into_raw(Box::new(RetiredNode {
            ptr,
            next: ptr::null_mut(),
        }));

        // Add to retired list
        #[allow(clippy::infinite_loop)] // CAS retry loop: exits via break on successful compare_exchange
        loop {
            let old_head = self.retired.load(Ordering::Acquire);
            // SAFETY: Unsafe operation validated - invariants maintained by surrounding code
            unsafe {
                (*retired_node).next = old_head;
            }

            if self
                .retired
                .compare_exchange_weak(old_head, retired_node, Ordering::Release, Ordering::Relaxed)
                .is_ok()
            {
                break;
            }
        }

        let count = self.retired_count.fetch_add(1, Ordering::Relaxed);

        // Trigger cleanup if we have too many retired nodes
        if count > 100 {
            self.cleanup();
        }
    }

    fn cleanup(&self) {
        // Simple cleanup: just delete all retired nodes
        // In a real implementation, this would check hazard pointers
        let head = self.retired.swap(ptr::null_mut(), Ordering::Acquire);
        let mut current = head;
        let mut count = 0;

        while !current.is_null() {
            // SAFETY: Pointer is valid, properly aligned, and exclusively owned
            unsafe {
                let node = Box::from_raw(current);
                drop(Box::from_raw(node.ptr));
                current = node.next;
                count += 1;
            }
        }

        self.retired_count.fetch_sub(count, Ordering::Relaxed);
    }
}

impl<T> Drop for HazardPointers<T> {
    fn drop(&mut self) {
        self.cleanup();
    }
}

/// High-performance atomic counter for sequence generation
#[repr(align(64))] // Cache line alignment
/// `AtomicCounter`
///
/// Auto-generated documentation placeholder - enhance with specifics
pub struct AtomicCounter {
    value: AtomicU64,
    increment: u64,
}

impl AtomicCounter {
    /// Create a new atomic counter starting at 0
    #[must_use]
    pub const fn new() -> Self {
        Self {
            value: AtomicU64::new(0),
            increment: 1,
        }
    }

    /// Create a new atomic counter with custom starting value and increment
    #[must_use]
    pub const fn new_with(start: u64, increment: u64) -> Self {
        Self {
            value: AtomicU64::new(start),
            increment,
        }
    }

    /// Get the next value (atomic increment)
    #[inline(always)]
    pub fn next(&self) -> u64 {
        self.value.fetch_add(self.increment, Ordering::Relaxed)
    }

    /// Get current value without incrementing
    #[inline(always)]
    pub fn get(&self) -> u64 {
        self.value.load(Ordering::Relaxed)
    }

    /// Reset counter to specific value
    #[inline(always)]
    pub fn reset(&self, value: u64) {
        self.value.store(value, Ordering::Relaxed);
    }

    /// Add a specific amount to the counter
    #[inline(always)]
    pub fn add(&self, amount: u64) -> u64 {
        self.value.fetch_add(amount, Ordering::Relaxed)
    }
}

impl Default for AtomicCounter {
    fn default() -> Self {
        Self::new()
    }
}

impl fmt::Debug for AtomicCounter {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("AtomicCounter")
            .field("value", &self.get())
            .field("increment", &self.increment)
            .finish()
    }
}

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

    #[test]
    fn test_mpsc_basic_operations() {
        let queue = MPSCQueue::<u64>::new();

        // Test empty queue
        assert!(queue.is_empty());
        assert_eq!(queue.try_pop(), None);

        // Test push/pop
        queue.push(42);
        assert!(!queue.is_empty());
        assert_eq!(queue.len(), 1);
        assert_eq!(queue.try_pop(), Some(42));
        assert!(queue.is_empty());
    }

    #[test]
    fn test_mpsc_multiple_producers() {
        let queue = Arc::new(MPSCQueue::<u64>::new());
        let num_producers = 4;
        let items_per_producer = 1000;

        let mut handles = Vec::new();

        // Spawn producer threads
        for producer_id in 0..num_producers {
            let queue_clone = Arc::clone(&queue);
            let handle = thread::spawn(move || {
                for i in 0..items_per_producer {
                    let value = (producer_id as u64) * 1000 + i;
                    queue_clone.push(value);
                }
            });
            handles.push(handle);
        }

        // Wait for all producers to finish
        for handle in handles {
            handle.join().expect("Producer thread failed");
        }

        // Consume all items
        let mut received = Vec::new();
        while let Some(item) = queue.try_pop() {
            received.push(item);
        }

        // Verify all items received
        assert_eq!(
            received.len(),
            (num_producers * items_per_producer) as usize
        );
        assert!(queue.is_empty());
    }

    #[test]
    fn test_atomic_counter() {
        let counter = AtomicCounter::new();

        assert_eq!(counter.get(), 0);
        assert_eq!(counter.next(), 0);
        assert_eq!(counter.next(), 1);
        assert_eq!(counter.get(), 2);

        counter.reset(100);
        assert_eq!(counter.get(), 100);
        assert_eq!(counter.next(), 100);
    }

    #[test]
    fn test_atomic_counter_concurrent() {
        let counter = Arc::new(AtomicCounter::new());
        let num_threads = 8;
        let increments_per_thread = 1000;

        let mut handles = Vec::new();

        for _ in 0..num_threads {
            let counter_clone = Arc::clone(&counter);
            let handle = thread::spawn(move || {
                let mut values = Vec::new();
                for _ in 0..increments_per_thread {
                    values.push(counter_clone.next());
                }
                values
            });
            handles.push(handle);
        }

        let mut all_values = Vec::new();
        for handle in handles {
            let values = handle.join().expect("Thread failed");
            all_values.extend(values);
        }

        // Verify all values are unique and within expected range
        all_values.sort_unstable();
        assert_eq!(all_values.len(), num_threads * increments_per_thread);

        // Check that all values from 0 to total-1 are present
        for (i, value) in all_values.into_iter().enumerate() {
            assert_eq!(value, i as u64);
        }
    }

    #[test]
    fn test_mpsc_performance() {
        let queue = Arc::new(MPSCQueue::<u64>::new());
        let queue_consumer = Arc::clone(&queue);

        const NUM_ITEMS: usize = 100_000;

        // Producer thread
        let producer = thread::spawn(move || {
            let start = std::time::Instant::now();
            for i in 0..NUM_ITEMS {
                queue.push(i as u64);
            }
            start.elapsed()
        });

        // Consumer thread
        let consumer = thread::spawn(move || {
            let mut received = 0;
            let start = std::time::Instant::now();

            while received < NUM_ITEMS {
                if queue_consumer.try_pop().is_some() {
                    received += 1;
                } else {
                    thread::yield_now();
                }
            }

            (start.elapsed(), received)
        });

        let producer_time = producer.join().expect("Producer failed");
        let (consumer_time, items_received) = consumer.join().expect("Consumer failed");

        assert_eq!(items_received, NUM_ITEMS);

        let producer_rate = NUM_ITEMS as f64 / producer_time.as_secs_f64();
        let consumer_rate = NUM_ITEMS as f64 / consumer_time.as_secs_f64();

        println!("Producer: {:.0} items/sec", producer_rate);
        println!("Consumer: {:.0} items/sec", consumer_rate);

        // Verify reasonable performance (should handle >100K ops/sec)
        assert!(
            producer_rate > 100_000.0,
            "Producer too slow: {:.0} ops/sec",
            producer_rate
        );
        assert!(
            consumer_rate > 100_000.0,
            "Consumer too slow: {:.0} ops/sec",
            consumer_rate
        );
    }
}