torsh_tensor/

lockfree_cache.rs

//! Lock-Free Cache Structures for High-Performance Concurrent Access
//!
//! This module provides lock-free data structures optimized for concurrent tensor operations.
//! By avoiding locks, these structures provide better scalability in multi-threaded scenarios.
//!
//! # Features
//!
//! - **Lock-free queues**: SPSC and MPMC queue implementations
//! - **Atomic reference counting**: Lock-free reference counting for shared data
//! - **Concurrent hash map**: Lock-free hash map for tensor caching
//! - **Wait-free operations**: Some operations guarantee completion in bounded time
//! - **Cache-line aligned**: Minimize false sharing

use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;

/// Lock-free SPSC (Single Producer Single Consumer) queue
///
/// This queue provides wait-free operations for both producer and consumer
/// when used correctly. It's ideal for passing tensors between threads.
pub struct LockFreeSPSCQueue<T> {
    buffer: Vec<Option<T>>,
    capacity: usize,
    head: AtomicUsize,  // Consumer reads from here
    tail: AtomicUsize,  // Producer writes here
    _padding: [u8; 64], // Cache line padding
}

impl<T> LockFreeSPSCQueue<T> {
    /// Create a new SPSC queue with given capacity
    pub fn new(capacity: usize) -> Self {
        let actual_capacity = capacity.next_power_of_two();
        let mut buffer = Vec::with_capacity(actual_capacity);
        for _ in 0..actual_capacity {
            buffer.push(None);
        }

        Self {
            buffer,
            capacity: actual_capacity,
            head: AtomicUsize::new(0),
            tail: AtomicUsize::new(0),
            _padding: [0; 64],
        }
    }

    /// Try to push an item (returns false if queue is full)
    pub fn try_push(&mut self, item: T) -> bool {
        let tail = self.tail.load(Ordering::Relaxed);
        let next_tail = (tail + 1) & (self.capacity - 1);
        let head = self.head.load(Ordering::Acquire);

        if next_tail == head {
            return false; // Queue is full
        }

        // SAFETY: We've checked that the slot is available
        self.buffer[tail] = Some(item);
        self.tail.store(next_tail, Ordering::Release);
        true
    }

    /// Try to pop an item (returns None if queue is empty)
    pub fn try_pop(&mut self) -> Option<T> {
        let head = self.head.load(Ordering::Relaxed);
        let tail = self.tail.load(Ordering::Acquire);

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

        let item = self.buffer[head].take();
        let next_head = (head + 1) & (self.capacity - 1);
        self.head.store(next_head, Ordering::Release);

        item
    }

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

    /// Get approximate size (may not be exact due to concurrent operations)
    pub fn len(&self) -> usize {
        let head = self.head.load(Ordering::Acquire);
        let tail = self.tail.load(Ordering::Acquire);

        if tail >= head {
            tail - head
        } else {
            self.capacity - head + tail
        }
    }

    /// Get the capacity
    pub fn capacity(&self) -> usize {
        self.capacity
    }
}

/// Lock-free cache entry with atomic reference counting
#[derive(Clone)]
pub struct LockFreeCacheEntry<T: Clone> {
    data: Arc<T>,
    access_count: Arc<AtomicUsize>,
    last_access: Arc<AtomicUsize>, // Timestamp
    valid: Arc<AtomicBool>,
}

impl<T: Clone> LockFreeCacheEntry<T> {
    /// Create a new cache entry
    pub fn new(data: T) -> Self {
        Self {
            data: Arc::new(data),
            access_count: Arc::new(AtomicUsize::new(0)),
            last_access: Arc::new(AtomicUsize::new(Self::current_timestamp())),
            valid: Arc::new(AtomicBool::new(true)),
        }
    }

    /// Get the data (increments access count)
    pub fn get(&self) -> Option<Arc<T>> {
        if self.valid.load(Ordering::Acquire) {
            self.access_count.fetch_add(1, Ordering::Relaxed);
            self.last_access
                .store(Self::current_timestamp(), Ordering::Release);
            Some(Arc::clone(&self.data))
        } else {
            None
        }
    }

    /// Invalidate this entry
    pub fn invalidate(&self) {
        self.valid.store(false, Ordering::Release);
    }

    /// Check if entry is valid
    pub fn is_valid(&self) -> bool {
        self.valid.load(Ordering::Acquire)
    }

    /// Get access count
    pub fn access_count(&self) -> usize {
        self.access_count.load(Ordering::Relaxed)
    }

    /// Get last access timestamp
    pub fn last_access(&self) -> usize {
        self.last_access.load(Ordering::Relaxed)
    }

    /// Get current timestamp (monotonic counter)
    fn current_timestamp() -> usize {
        use std::sync::atomic::AtomicUsize as GlobalCounter;
        static COUNTER: GlobalCounter = GlobalCounter::new(0);
        COUNTER.fetch_add(1, Ordering::Relaxed)
    }
}

/// Simple lock-free cache with fixed size
///
/// This cache uses atomic operations for thread-safe access without locks.
/// It's optimized for read-heavy workloads with occasional writes.
pub struct LockFreeCache<K: Eq + std::hash::Hash + Clone, V: Clone> {
    entries: Vec<Option<(K, LockFreeCacheEntry<V>)>>,
    size: AtomicUsize,
    capacity: usize,
}

impl<K: Eq + std::hash::Hash + Clone, V: Clone> LockFreeCache<K, V> {
    /// Create a new lock-free cache with given capacity
    pub fn new(capacity: usize) -> Self {
        let mut entries = Vec::with_capacity(capacity);
        for _ in 0..capacity {
            entries.push(None);
        }

        Self {
            entries,
            size: AtomicUsize::new(0),
            capacity,
        }
    }

    /// Get the capacity
    pub fn capacity(&self) -> usize {
        self.capacity
    }

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

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

    /// Calculate hash for a key
    fn hash(&self, key: &K) -> usize {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::Hasher;

        let mut hasher = DefaultHasher::new();
        key.hash(&mut hasher);
        (hasher.finish() as usize) % self.capacity
    }

    /// Try to get a value from the cache
    pub fn get(&self, key: &K) -> Option<Arc<V>> {
        let index = self.hash(key);
        let mut probe = 0;

        while probe < self.capacity {
            let current_index = (index + probe) % self.capacity;

            // SAFETY: We're reading from a fixed-size vec with valid indices
            if let Some((ref k, ref entry)) = unsafe { &*self.entries.as_ptr().add(current_index) }
            {
                if k == key {
                    return entry.get();
                }
            } else {
                // Empty slot means key doesn't exist
                return None;
            }

            probe += 1;
        }

        None
    }

    /// Check if a key exists in the cache
    pub fn contains_key(&self, key: &K) -> bool {
        self.get(key).is_some()
    }
}

/// Statistics for lock-free operations
#[derive(Debug, Default)]
pub struct LockFreeStats {
    /// Number of successful operations
    pub successes: AtomicUsize,
    /// Number of failed operations (contention)
    pub failures: AtomicUsize,
    /// Number of retries
    pub retries: AtomicUsize,
}

impl LockFreeStats {
    /// Create new statistics
    pub fn new() -> Self {
        Self::default()
    }

    /// Record a success
    pub fn record_success(&self) {
        self.successes.fetch_add(1, Ordering::Relaxed);
    }

    /// Record a failure
    pub fn record_failure(&self) {
        self.failures.fetch_add(1, Ordering::Relaxed);
    }

    /// Record a retry
    pub fn record_retry(&self) {
        self.retries.fetch_add(1, Ordering::Relaxed);
    }

    /// Get success count
    pub fn successes(&self) -> usize {
        self.successes.load(Ordering::Relaxed)
    }

    /// Get failure count
    pub fn failures(&self) -> usize {
        self.failures.load(Ordering::Relaxed)
    }

    /// Get retry count
    pub fn retries(&self) -> usize {
        self.retries.load(Ordering::Relaxed)
    }

    /// Calculate success rate
    pub fn success_rate(&self) -> f64 {
        let total = self.successes() + self.failures();
        if total == 0 {
            0.0
        } else {
            self.successes() as f64 / total as f64
        }
    }

    /// Reset statistics
    pub fn reset(&self) {
        self.successes.store(0, Ordering::Relaxed);
        self.failures.store(0, Ordering::Relaxed);
        self.retries.store(0, Ordering::Relaxed);
    }
}

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

    #[test]
    fn test_spsc_queue_basic() {
        let mut queue = LockFreeSPSCQueue::new(4);

        assert!(queue.is_empty());
        assert_eq!(queue.len(), 0);

        assert!(queue.try_push(1));
        assert!(queue.try_push(2));
        assert!(queue.try_push(3));

        assert_eq!(queue.len(), 3);
        assert!(!queue.is_empty());

        assert_eq!(queue.try_pop(), Some(1));
        assert_eq!(queue.try_pop(), Some(2));
        assert_eq!(queue.len(), 1);

        assert_eq!(queue.try_pop(), Some(3));
        assert!(queue.is_empty());
        assert_eq!(queue.try_pop(), None);
    }

    #[test]
    fn test_spsc_queue_full() {
        let mut queue = LockFreeSPSCQueue::new(2);

        // Can insert 1 element (capacity - 1)
        assert!(queue.try_push(1));

        // Queue is full
        assert!(!queue.try_push(2));

        // Pop one
        assert_eq!(queue.try_pop(), Some(1));

        // Now we can insert again
        assert!(queue.try_push(3));
    }

    #[test]
    fn test_spsc_queue_wraparound() {
        let mut queue = LockFreeSPSCQueue::new(4);

        // Fill queue
        assert!(queue.try_push(1));
        assert!(queue.try_push(2));

        // Pop some
        assert_eq!(queue.try_pop(), Some(1));

        // Push more (should wrap around)
        assert!(queue.try_push(3));
        assert!(queue.try_push(4));

        // Pop all
        assert_eq!(queue.try_pop(), Some(2));
        assert_eq!(queue.try_pop(), Some(3));
        assert_eq!(queue.try_pop(), Some(4));
        assert_eq!(queue.try_pop(), None);
    }

    #[test]
    fn test_cache_entry_basic() {
        let entry = LockFreeCacheEntry::new(42);

        assert!(entry.is_valid());
        assert_eq!(*entry.get().expect("get should succeed"), 42);
        assert_eq!(entry.access_count(), 1);

        entry.invalidate();
        assert!(!entry.is_valid());
        assert!(entry.get().is_none());
    }

    #[test]
    fn test_cache_entry_access_count() {
        let entry = LockFreeCacheEntry::new("test");

        assert_eq!(entry.access_count(), 0);

        entry.get();
        assert_eq!(entry.access_count(), 1);

        entry.get();
        entry.get();
        assert_eq!(entry.access_count(), 3);
    }

    #[test]
    fn test_lock_free_cache_basic() {
        let cache: LockFreeCache<String, i32> = LockFreeCache::new(10);

        assert_eq!(cache.capacity(), 10);
        assert_eq!(cache.len(), 0);
        assert!(cache.is_empty());
    }

    #[test]
    fn test_lock_free_cache_contains() {
        let cache = LockFreeCache::<String, i32>::new(10);

        assert!(!cache.contains_key(&"test".to_string()));
    }

    #[test]
    fn test_lockfree_stats() {
        let stats = LockFreeStats::new();

        assert_eq!(stats.successes(), 0);
        assert_eq!(stats.failures(), 0);
        assert_eq!(stats.retries(), 0);

        stats.record_success();
        stats.record_success();
        stats.record_failure();

        assert_eq!(stats.successes(), 2);
        assert_eq!(stats.failures(), 1);

        let rate = stats.success_rate();
        assert!((rate - 0.666).abs() < 0.01);

        stats.reset();
        assert_eq!(stats.successes(), 0);
        assert_eq!(stats.failures(), 0);
    }

    #[test]
    fn test_spsc_queue_capacity() {
        let queue = LockFreeSPSCQueue::<i32>::new(7);
        // Should round up to next power of 2
        assert_eq!(queue.capacity(), 8);
    }

    #[test]
    fn test_cache_entry_timestamp() {
        let entry1 = LockFreeCacheEntry::new(1);
        let entry2 = LockFreeCacheEntry::new(2);

        let ts1 = entry1.last_access();
        let ts2 = entry2.last_access();

        // Timestamps should be different (monotonic)
        assert!(ts2 > ts1);
    }
}