moirai-sync 0.7.0

#![cfg_attr(test, allow(clippy::unwrap_used, reason = "test scope"))]

use std::hint;
use std::sync::Arc;
use std::thread;

use moirai_core::pool::LockFreeStack;

use super::concurrent_hash_map::ConcurrentHashMap;
use super::futex_mutex::FutexMutex;
use super::spin_lock::SpinLock;
use super::wait_group::WaitGroup;

const TEST_THREAD_COUNT: usize = 10;
const OPERATIONS_PER_THREAD: usize = 100;
const TEST_ELEMENT_COUNT: usize = 1000;

#[test]
fn test_wait_group() {
    use std::sync::Barrier;

    let wg = Arc::new(WaitGroup::new());
    let mut handles = vec![];
    let start = Arc::new(Barrier::new(4));
    let release = Arc::new(Barrier::new(4));

    wg.add(3);

    for _ in 0..3 {
        let wg = wg.clone();
        let start = Arc::clone(&start);
        let release = Arc::clone(&release);
        handles.push(thread::spawn(move || {
            start.wait();
            release.wait();
            wg.done();
        }));
    }

    start.wait();
    let waiter = {
        let wg = Arc::clone(&wg);
        thread::spawn(move || wg.wait())
    };
    release.wait();
    waiter.join().unwrap();

    for handle in handles {
        handle.join().unwrap();
    }
}

#[test]
fn test_futex_mutex() {
    let mutex = Arc::new(FutexMutex::new(0));
    let mut handles = vec![];

    for _ in 0..TEST_THREAD_COUNT {
        let mutex = mutex.clone();
        handles.push(thread::spawn(move || {
            for _ in 0..OPERATIONS_PER_THREAD {
                let mut guard = mutex.lock();
                *guard += 1;
            }
        }));
    }

    for handle in handles {
        handle.join().unwrap();
    }

    assert_eq!(*mutex.lock(), 1000);
}

#[test]
fn test_futex_mutex_blocking_path_no_lost_wakeup() {
    // Regression for the non-Linux fallback lost-wakeup: `unlock` checked the
    // `waiters` count with a StoreLoad-reorderable load relative to releasing
    // `locked`, so a waiter registering concurrently could be skipped and left
    // parked on the condvar forever. To exercise the *blocking* path (not the
    // adaptive spin), each critical section holds the lock long enough that
    // contending threads exhaust their spin budget and must block. If a wakeup
    // is lost, a thread never reacquires and the test hangs into the nextest
    // timeout. The exact final count also proves no update was lost.
    const THREADS: usize = 16;
    const ITERS: usize = 200;

    let mutex = Arc::new(FutexMutex::new(0u64));
    let mut handles = vec![];
    for _ in 0..THREADS {
        let mutex = mutex.clone();
        handles.push(thread::spawn(move || {
            for _ in 0..ITERS {
                let mut guard = mutex.lock();
                let prev = *guard;
                // Spin briefly while holding the lock to force other threads
                // past the spin budget and onto the blocking/condvar path.
                for _ in 0..256 {
                    hint::spin_loop();
                }
                *guard = prev + 1;
            }
        }));
    }
    for handle in handles {
        handle.join().unwrap();
    }
    assert_eq!(*mutex.lock(), (THREADS * ITERS) as u64);
}

#[test]
fn test_futex_mutex_try_lock() {
    let mutex = FutexMutex::new(42);

    // Uncontended try_lock should succeed
    {
        let guard = mutex
            .try_lock()
            .expect("an uncontended try_lock must succeed");
        assert_eq!(*guard, 42);
    }

    // Contended try_lock should return None
    {
        let _guard = mutex.lock();
        let try_guard = mutex.try_lock();
        assert!(try_guard.is_none());
    }
}

#[test]
fn test_lock_free_stack() {
    let stack = Arc::new(LockFreeStack::new());
    let mut handles = vec![];

    // Push from multiple threads
    for i in 0..10 {
        let stack = stack.clone();
        handles.push(thread::spawn(move || {
            // Capacity is the default 1024; these 10 pushes always succeed.
            stack
                .push(i)
                .expect("push into a non-full stack must succeed");
        }));
    }

    for handle in handles {
        handle.join().unwrap();
    }

    // Pop all items
    let mut items = vec![];
    while let Some(item) = stack.pop() {
        items.push(item);
    }

    items.sort();
    assert_eq!(items, (0..10).collect::<Vec<_>>());
}

#[test]
fn test_concurrent_hashmap() {
    let map = ConcurrentHashMap::new();

    // Insert some values
    assert!(map.insert("key1", 100).unwrap().is_none());
    assert!(map.insert("key2", 200).unwrap().is_none());

    // Test retrieval
    assert_eq!(map.get(&"key1").unwrap(), Some(100));
    assert_eq!(map.get(&"key2").unwrap(), Some(200));
    assert_eq!(map.get(&"key3").unwrap(), None);

    // Test removal
    assert_eq!(map.remove(&"key1").unwrap(), Some(100));
    assert_eq!(map.get(&"key1").unwrap(), None);
}

#[test]
fn test_concurrent_hashmap_segment_distribution() {
    use std::collections::HashSet;

    // Create a map with 16 segments
    let map = ConcurrentHashMap::<i32, i32>::with_segments(16);

    // Track which segments are used
    let mut segments_used = HashSet::new();

    // Insert many keys and track segment distribution
    for i in 0..TEST_ELEMENT_COUNT {
        let key = i as i32;
        map.insert(key, key).unwrap();
        let segment_idx = map.segment_index(&key);
        segments_used.insert(segment_idx);
    }

    // With proper distribution, we should use most segments
    // With 1000 keys across 16 segments, we expect to use all segments
    assert!(
        segments_used.len() >= 14,
        "Poor segment distribution: only {} of 16 segments used",
        segments_used.len()
    );

    // Verify all keys can be retrieved
    for i in 0..TEST_ELEMENT_COUNT {
        let key = i as i32;
        assert_eq!(map.get(&key).unwrap(), Some(key));
    }
}

#[test]
fn test_spinlock_basic_functionality() {
    let lock = SpinLock::new(0);

    // Test basic lock/unlock
    {
        let mut guard = lock.lock();
        *guard = 42;
    }

    // Test that value was updated
    {
        let guard = lock.lock();
        assert_eq!(*guard, 42);
    }
}

#[test]
fn test_spinlock_try_lock() {
    let lock = SpinLock::new(0);

    // Should be able to try_lock on unlocked
    let guard1 = lock
        .try_lock()
        .expect("an unlocked SpinLock must yield a guard");
    assert_eq!(*guard1, 0);

    // Should fail to try_lock when locked
    let guard2 = lock.try_lock();
    assert!(guard2.is_none());

    // Should succeed after first guard is dropped
    drop(guard1);
    let guard3 = lock
        .try_lock()
        .expect("the lock must be reacquirable after the first guard drops");
    assert_eq!(*guard3, 0);
}

#[test]
fn test_spinlock_contention() {
    let lock = Arc::new(SpinLock::new(0));
    let mut handles = vec![];

    // Spawn threads that increment a counter
    for _ in 0..TEST_THREAD_COUNT {
        let lock = lock.clone();
        handles.push(thread::spawn(move || {
            for _ in 0..OPERATIONS_PER_THREAD {
                let mut guard = lock.lock();
                *guard += 1;
                // Hold the lock briefly to create contention
                for _ in 0..10 {
                    hint::spin_loop();
                }
            }
        }));
    }

    // Wait for all threads to complete
    for handle in handles {
        handle.join().unwrap();
    }

    // Verify final count
    let guard = lock.lock();
    assert_eq!(*guard, TEST_THREAD_COUNT * OPERATIONS_PER_THREAD);
}

#[test]
fn test_spinlock_drop_behavior() {
    let lock = SpinLock::new(vec![1, 2, 3]);

    // Test that guard properly derefs
    {
        let guard = lock.lock();
        assert_eq!(guard.len(), 3);
        assert_eq!(guard[0], 1);
    }

    // Test that guard properly derefs mutably
    {
        let mut guard = lock.lock();
        guard.push(4);
        assert_eq!(guard.len(), 4);
    }

    // Verify changes persisted
    {
        let guard = lock.lock();
        assert_eq!(*guard, vec![1, 2, 3, 4]);
    }
}

#[test]
fn test_spinlock_send_sync() {
    // Test that SpinLock implements Send + Sync
    fn assert_send_sync<T: Send + Sync>() {}
    assert_send_sync::<SpinLock<i32>>();

    // Test that we can move SpinLock across threads
    let lock = SpinLock::new(42);
    let handle = thread::spawn(move || {
        let guard = lock.lock();
        *guard
    });

    assert_eq!(handle.join().unwrap(), 42);
}

use super::resource_pool::{ShardedResourcePool, SizeBounded};

struct TestResource {
    id: usize,
    size: u64,
}

impl SizeBounded for TestResource {
    fn size(&self) -> u64 {
        self.size
    }
}

#[test]
fn test_sharded_resource_pool_basic() {
    let pool = ShardedResourcePool::<TestResource>::new(16, 1024);

    // Recycle some elements
    pool.recycle(TestResource { id: 1, size: 100 });
    pool.recycle(TestResource { id: 2, size: 100 });

    // Since we are on the same thread, we should pop the same shard in LIFO order (id: 2 first)
    let res1 = pool.take_at_least(100).expect("should find resource");
    assert_eq!(res1.id, 2);

    let res2 = pool.take_at_least(100).expect("should find resource");
    assert_eq!(res2.id, 1);

    // Binned matching: requesting size 50 should return any resource of size >= 50
    pool.recycle(TestResource { id: 3, size: 128 });
    let res3 = pool.take_at_least(50).expect("should match bin >= 50");
    assert_eq!(res3.size, 128);
}

#[test]
fn test_sharded_resource_pool_steals_from_other_shards() {
    let pool = Arc::new(ShardedResourcePool::<TestResource>::new(16, 1024));

    // Thread 1 recycles an item
    let pool_clone = pool.clone();
    let handle1 = thread::spawn(move || {
        pool_clone.recycle(TestResource { id: 42, size: 256 });
    });
    handle1.join().unwrap();

    // Thread 2 pops the item (forcing a cross-shard steal)
    let pool_clone = pool.clone();
    let handle2 = thread::spawn(move || {
        let res = pool_clone
            .take_at_least(256)
            .expect("should steal resource");
        assert_eq!(res.id, 42);
    });
    handle2.join().unwrap();
}

#[test]
fn test_sharded_resource_pool_fifo_eviction() {
    // Max buffers = 4 per pool, which maps to (4/4).max(1) = 1 max buffer per shard.
    // Max bytes = 1000 per pool, which maps to 250 max bytes per shard.
    let pool = ShardedResourcePool::<TestResource>::new(4, 1000);

    // Recycle elements onto local shard. Since shard limit is 1, each new recycle will evict the previous one.
    pool.recycle(TestResource { id: 1, size: 100 });
    pool.recycle(TestResource { id: 2, size: 100 });

    // The first element (id: 1) should be evicted, leaving only id: 2
    let res = pool.take_at_least(100).expect("should find resource");
    assert_eq!(res.id, 2);
    assert!(pool.take_at_least(100).is_none());
}

#[test]
fn test_concurrent_hashmap_get_or_insert_with() {
    let map = Arc::new(ConcurrentHashMap::<String, i32>::new());
    let mut handles = vec![];

    for i in 0..10 {
        let map = map.clone();
        handles.push(thread::spawn(move || {
            let val = map
                .get_or_insert_with("shared_key".to_string(), || i)
                .unwrap();
            // All threads should resolve to the same value (the first thread that gets write lock)
            assert!((0..10).contains(&val));
            val
        }));
    }

    let mut values = vec![];
    for handle in handles {
        values.push(handle.join().unwrap());
    }

    // Verify all threads got the exact same value
    let first = values[0];
    assert!(values.iter().all(|&v| v == first));
}

#[test]
fn test_sharded_resource_pool_vecdeque_fifo() {
    // Max buffers = 8, which means 2 max buffers per shard
    let pool = ShardedResourcePool::<TestResource>::new(8, 1000);

    // Recycle three resources. With max buffer limit of 2, the first should be evicted.
    pool.recycle(TestResource { id: 1, size: 100 });
    pool.recycle(TestResource { id: 2, size: 100 });
    pool.recycle(TestResource { id: 3, size: 100 });

    // Resource 1 should be evicted (FIFO eviction order for oldest)
    // The remaining should be id: 3 and id: 2, retrieved in LIFO order (3 first, then 2)
    let res1 = pool.take_at_least(100).expect("should retrieve");
    assert_eq!(res1.id, 3);

    let res2 = pool.take_at_least(100).expect("should retrieve");
    assert_eq!(res2.id, 2);

    assert!(pool.take_at_least(100).is_none());
}

#[test]
fn test_sharded_resource_pool_recycle_enforces_byte_budget() {
    // 4 shards: per-shard byte cap = 1000/4 = 250, buffer cap = 100/4 = 25.
    // All recycles on this thread land in one shard, so the byte cap binds: at
    // most two 100-byte items (200 <= 250; a third would be 300 > 250).
    let pool = ShardedResourcePool::<TestResource>::new(100, 1000);
    for id in 0..10 {
        pool.recycle(TestResource { id, size: 100 });
    }

    let mut total_bytes = 0;
    let mut count = 0;
    while let Some(item) = pool.take_at_least(1) {
        total_bytes += item.size();
        count += 1;
    }
    assert!(
        total_bytes <= 250,
        "retained bytes {total_bytes} exceeds the per-shard byte budget of 250"
    );
    assert!(
        count <= 2,
        "retained count {count} exceeds the byte-derived bound of 2"
    );
}

#[test]
fn test_sharded_resource_pool_concurrent_recycle_respects_total_cap() {
    // Aggregate cap = 4 shards * (16 / 4) = 16 buffers; the byte cap is large so
    // the buffer cap binds. Reserve-before-insert recycle must keep each shard
    // within its cap under concurrent recycle+take, and the counter arithmetic
    // must never underflow (which would panic under `overflow-checks`). A persisted
    // overshoot or underflow is the regression this guards.
    const MAX_BUFFERS: usize = 16;
    let pool = Arc::new(ShardedResourcePool::<TestResource>::new(
        MAX_BUFFERS,
        1 << 20,
    ));

    let mut handles = Vec::new();
    for t in 0..8 {
        let pool = Arc::clone(&pool);
        handles.push(thread::spawn(move || {
            for i in 0..2000 {
                pool.recycle(TestResource {
                    id: t * 100_000 + i,
                    size: 64,
                });
                if i % 3 == 0 {
                    let _ = pool.take_at_least(1);
                }
            }
        }));
    }
    for handle in handles {
        handle.join().expect("recycle worker must not panic");
    }

    let mut retained = 0;
    while pool.take_at_least(1).is_some() {
        retained += 1;
    }
    assert!(
        retained <= MAX_BUFFERS,
        "retained {retained} exceeds the aggregate cap of {MAX_BUFFERS}"
    );
}

#[test]
fn test_sharded_resource_pool_clear_serializes_reservation_and_insertion() {
    use std::sync::{Barrier, mpsc::sync_channel};

    let (recycle_entered_tx, recycle_entered_rx) = sync_channel(0);
    let (clear_started_tx, clear_started_rx) = sync_channel(0);
    let (clear_done_tx, clear_done_rx) = sync_channel(0);
    let release = Arc::new(Barrier::new(2));
    let pool = Arc::new(ShardedResourcePool::<TestResource>::new(16, 1024));
    let _hook = pool.install_test_hook(recycle_entered_tx, clear_started_tx, Arc::clone(&release));

    let recycle_pool = Arc::clone(&pool);
    let recycler = thread::spawn(move || {
        recycle_pool.recycle(TestResource { id: 7, size: 100 });
    });
    recycle_entered_rx
        .recv()
        .expect("recycle must reach the reservation/insertion boundary");

    let clear_pool = Arc::clone(&pool);
    let clearer = thread::spawn(move || {
        clear_pool.clear();
        clear_done_tx
            .send(())
            .expect("clear completion receiver remains active");
    });
    clear_started_rx
        .recv()
        .expect("clear must enter before waiting on the target bin");
    assert!(
        clear_done_rx.try_recv().is_err(),
        "clear must not complete while recycle owns the target bin"
    );

    release.wait();
    recycler.join().expect("recycle worker must not panic");
    clear_done_rx
        .recv()
        .expect("clear must complete after recycle publishes its item");
    clearer.join().expect("clear worker must not panic");

    assert!(
        pool.take_at_least(100).is_none(),
        "clear must remove the item published at the interleaving boundary"
    );

    drop(_hook);
    pool.recycle(TestResource { id: 8, size: 100 });
    assert_eq!(pool.take_at_least(100).map(|resource| resource.id), Some(8));
    assert!(pool.take_at_least(100).is_none());
}