arctic-wt 0.1.5

#![cfg(feature = "rand")]

use core::borrow::Borrow as _;
use core::time::Duration;
use std::sync::Condvar;
use std::sync::Mutex;

use arctic::Key;
use rand::SeedableRng;
use rand::rngs::Xoshiro256PlusPlus;
use rand::seq::SliceRandom;

mod u64 {
    use arctic::Key;

    use super::Workload;
    use super::test_map;

    #[test]
    fn many() {
        test_map(&U64, 16, 10_000_000, false);
    }

    #[test]
    fn two() {
        test_map(&U64, 2, 10_000_000, true);
    }

    #[test]
    fn one() {
        test_map(&U64, 1, 10_000_000, true);
    }

    struct U64;

    impl Workload for U64 {
        type Key<'k> = u64;

        type Value = u64;

        fn key(&self, index: usize) -> Self::Key<'_> {
            index as u64
        }

        fn value(&self, index: usize) -> Self::Value {
            index as u64
        }

        fn validate(
            &self,
            index: usize,
            key: &<Self::Key<'_> as Key>::Borrowed,
            value: &<Self::Value as arctic::concurrent::Value>::Borrowed,
        ) {
            assert_eq!(index as u64, *key);
            assert_eq!(index as u64, *value);
        }
    }
}

mod arced {
    use arctic::Key;

    use super::Workload;
    use super::test_map;

    struct Arced;

    #[test]
    fn many() {
        test_map(&Arced, 16, 10_000_000, false);
    }

    #[test]
    fn two() {
        test_map(&Arced, 2, 10_000_000, false);
    }

    #[test]
    fn one() {
        test_map(&Arced, 1, 10_000_000, false);
    }

    #[derive(Debug, PartialEq, Eq)]
    struct Entry {
        key: u32,
        value: u64,
    }

    impl Entry {
        fn new(index: usize) -> Self {
            Self {
                key: index as u32,
                value: index as u64 + 1,
            }
        }
    }

    impl Workload for Arced {
        type Key<'k> = u32;

        type Value = arctic::concurrent::value::Arc<Entry>;

        fn key(&self, index: usize) -> Self::Key<'_> {
            index as u32
        }

        fn value(&self, index: usize) -> Self::Value {
            arctic::sync::Arc::new(Entry::new(index)).into()
        }

        fn validate(
            &self,
            index: usize,
            key: &<Self::Key<'_> as Key>::Borrowed,
            value: &arctic::concurrent::value::ArcRef<Entry>,
        ) {
            assert_eq!(*key, index as u32);
            assert_eq!(**value, Entry::new(index));
        }
    }
}

mod boxed {
    use arctic::Key;

    use super::Workload;
    use super::test_map;

    struct Boxed;

    #[test]
    fn many() {
        test_map(&Boxed, 16, 10_000_000, false);
    }

    #[test]
    fn two() {
        test_map(&Boxed, 2, 10_000_000, false);
    }

    #[test]
    fn one() {
        test_map(&Boxed, 1, 10_000_000, false);
    }

    #[derive(Debug, PartialEq, Eq)]
    struct Entry {
        key: u32,
        value: u64,
    }

    impl Entry {
        fn new(index: usize) -> Self {
            Self {
                key: index as u32,
                value: index as u64 + 1,
            }
        }
    }

    impl Workload for Boxed {
        type Key<'k> = u32;

        type Value = Box<Entry>;

        fn key(&self, index: usize) -> Self::Key<'_> {
            index as u32
        }

        fn value(&self, index: usize) -> Self::Value {
            Box::new(Entry::new(index))
        }

        fn validate(&self, index: usize, key: &<Self::Key<'_> as Key>::Borrowed, value: &Entry) {
            assert_eq!(*key, index as u32);
            assert_eq!(*value, Entry::new(index));
        }
    }
}

mod boxed_str_terminated {
    use arctic::Key;
    use arctic::key::BoxedStr;
    use arctic::key::Terminated;
    use rand::SeedableRng as _;
    use rand::distr::Distribution as _;
    use rand::distr::SampleString;
    use rand::distr::Uniform;

    use super::Workload;
    use super::test_map;

    struct Bytes;

    #[test]
    fn many() {
        test_map(&Bytes, 16, 10_000_000, false);
    }

    #[test]
    fn two() {
        test_map(&Bytes, 2, 1_000_000, false);
    }

    #[test]
    fn one() {
        test_map(&Bytes, 1, 1_000_000, false);
    }

    impl Workload for Bytes {
        type Key<'k> = BoxedStr<Terminated<0>>;
        type Value = u64;

        fn key(&self, index: usize) -> Self::Key<'_> {
            let mut rng = rand::rngs::Xoshiro256PlusPlus::seed_from_u64(index as u64);
            let len = rand::distr::Uniform::new_inclusive(16usize, 32usize)
                .unwrap()
                .sample(&mut rng);

            let mut buffer = Uniform::new(1 as char, char::MAX)
                .unwrap()
                .sample_string(&mut rng, len);
            buffer.push('\0');

            BoxedStr::<Terminated<0>>::new(buffer).unwrap()
        }

        fn value(&self, index: usize) -> Self::Value {
            index as u64
        }

        fn validate(
            &self,
            index: usize,
            key: &<Self::Key<'_> as Key>::Borrowed,
            value: &<Self::Value as arctic::concurrent::Value>::Borrowed,
        ) {
            assert_eq!(key, self.key(index).as_slice());
            assert_eq!(*value, index as u64);
        }
    }
}

mod slice_non_null {
    use arctic::Key;
    use arctic::key;
    use arctic::key::BoxedSlice;
    use arctic::key::NonNull;
    use rand::SeedableRng as _;
    use rand::distr::Distribution as _;
    use rand::distr::SampleString as _;

    use super::Workload;
    use super::test_map;

    struct Slice(Vec<BoxedSlice<NonNull>>);

    #[test]
    fn many() {
        test_map(&Slice::new(10_000_000), 16, 10_000_000, false);
    }

    #[test]
    fn two() {
        test_map(&Slice::new(1_000_000), 2, 1_000_000, false);
    }

    #[test]
    fn one() {
        test_map(&Slice::new(1_000_000), 1, 1_000_000, false);
    }

    impl Slice {
        fn new(key_count: usize) -> Self {
            let mut outer = Vec::new();
            let mut rng = rand::rngs::Xoshiro256PlusPlus::seed_from_u64(key_count as u64);
            let dist_char =
                rand::distr::uniform::Uniform::<char>::new_inclusive(1 as char, char::MAX).unwrap();
            let dist_len = rand::distr::Uniform::new_inclusive(1usize, 32usize).unwrap();
            for _ in 0..key_count {
                let len = dist_len.sample(&mut rng);
                let inner = dist_char.sample_string(&mut rng, len);
                outer.push(key::BoxedSlice::new(inner.into_bytes()).unwrap());
            }
            Self(outer)
        }
    }

    impl Workload for Slice {
        type Key<'k> = &'k key::Slice<NonNull>;
        type Value = u64;

        fn key(&self, index: usize) -> Self::Key<'_> {
            self.0[index].as_slice()
        }

        fn value(&self, index: usize) -> Self::Value {
            index as u64
        }

        fn validate(
            &self,
            index: usize,
            key: &<Self::Key<'_> as Key>::Borrowed,
            value: &<Self::Value as arctic::concurrent::Value>::Borrowed,
        ) {
            assert!(core::ptr::eq(key, self.key(index)));
            assert_eq!(*value, index as u64);
        }
    }
}

mod array {
    use arctic::Key;
    use rand::RngExt as _;
    use rand::SeedableRng as _;
    use rand::rngs::Xoshiro256PlusPlus;

    use super::Workload;
    use super::test_map;

    struct Array<const N: usize>;

    #[test]
    fn many() {
        test_map(&Array::<12>, 16, 10_000_000, false);
    }

    #[test]
    fn two() {
        test_map(&Array::<19>, 2, 1_000_000, false);
    }

    #[test]
    fn one() {
        test_map(&Array::<21>, 1, 1_000_000, false);
    }

    impl<const N: usize> Workload for Array<N> {
        type Key<'k> = [u8; N];
        type Value = u64;

        fn key(&self, index: usize) -> Self::Key<'_> {
            Xoshiro256PlusPlus::seed_from_u64(index as u64).random()
        }

        fn value(&self, index: usize) -> Self::Value {
            index as u64
        }

        fn validate(
            &self,
            index: usize,
            key: &<Self::Key<'_> as Key>::Borrowed,
            value: &<Self::Value as arctic::concurrent::Value>::Borrowed,
        ) {
            assert_eq!(*key, self.key(index));
            assert_eq!(*value, index as u64);
        }
    }
}

trait Workload: Sized + Sync {
    type Key<'k>: arctic::Key + Sync
    where
        Self: 'k;

    type Value: arctic::concurrent::Value + Send + Sync;

    fn key(&self, index: usize) -> Self::Key<'_>;

    fn value(&self, index: usize) -> Self::Value;

    fn validate<'k>(
        &'k self,
        index: usize,
        key: &<Self::Key<'k> as Key>::Borrowed,
        value: &<Self::Value as arctic::concurrent::Value>::Borrowed,
    );
}

// Can't cfg trait bound in where clause
#[cfg(feature = "smr-hazard")]
trait Hazard: arctic::concurrent::smr::hazard::Key {}
#[cfg(feature = "smr-hazard")]
impl<K: arctic::concurrent::smr::hazard::Key> Hazard for K {}
#[cfg(not(feature = "smr-hazard"))]
trait Hazard {}
#[cfg(not(feature = "smr-hazard"))]
impl<K> Hazard for K {}

fn test_map<'k, K: Workload>(key_set: &'k K, thread_count: usize, key_count: usize, shuffle: bool)
where
    for<'a> &'a <K::Key<'k> as Key>::Borrowed: Sync + core::fmt::Debug,
    <K::Value as arctic::concurrent::Value>::Borrowed: core::fmt::Debug,
    K::Key<'k>: Clone + Ord + core::fmt::Debug,
    K::Key<'k>: Hazard,
{
    let barrier = &Barrier::new(thread_count);

    let mut items = (0..key_count)
        .map(|index| (index, key_set.key(index)))
        .collect::<Vec<_>>();
    items.sort_unstable_by(|(_, key_a), (_, key_b)| key_a.cmp(key_b));
    items.dedup_by(|(_, key_a), (_, key_b)| key_a == key_b);

    if shuffle {
        let mut rng = Xoshiro256PlusPlus::seed_from_u64((thread_count * key_count) as u64);
        items.shuffle(&mut rng);
    } else {
        items.sort_unstable_by_key(|(index, _)| *index);
    }

    let map = &arctic::concurrent::Map::<
        K::Key<'_>,
        _,
        // Can't define type alias with unused type parameters or that uses outside types
        cfg_select! {
            feature = "smr-hazard" => {
                ::arctic::concurrent::smr::Hazard::<K::Key<'_>, K::Value>
            }
            feature = "smr-epoch" => {
                ::arctic::concurrent::smr::Epoch
            }
            feature = "smr-seize" => {
                ::arctic::concurrent::smr::Seize
            }
            _ => {
                ::arctic::concurrent::smr::NoOp
            }
        },
    >::default();

    std::thread::scope(|scope| {
        for chunk in items.chunks(key_count / thread_count) {
            scope.spawn(move || {
                barrier.wait();

                for (index, key) in chunk {
                    let value = key_set.value(*index);
                    map.insert(key.as_insert(), value)
                        .ok()
                        .as_deref()
                        .unwrap_or_else(|| panic!("Key {:?} should not be present", key.borrow()));

                    if map.get(key.borrow()).is_none() {
                        panic!("failed to find {:x?}", key);
                    }
                }

                barrier.wait();

                for (index, key) in chunk.iter().take(chunk.len() / 2) {
                    let value = map
                        .remove(key.borrow())
                        .unwrap_or_else(|| panic!("failed to find {:x?}", key));
                    key_set.validate(*index, key.borrow(), &value);
                }

                barrier.wait();

                for (index, key) in chunk.iter().skip(chunk.len() / 2) {
                    let value = map.get(key.borrow());
                    key_set.validate(*index, key.borrow(), value.as_deref().unwrap());
                }
            });
        }
    });
}

/// Identical to [`std::sync::Barrier`], but panics if the mutex is poisoned
/// or if the barrier times out.
///
/// Otherwise a panic within [`std::thread::scope`] causes other
/// threads to block indefinitely on the barrier.
struct Barrier {
    mutex: Mutex<BarrierState>,
    condition: Condvar,
    thread_count: usize,
}

impl Barrier {
    fn new(thread_count: usize) -> Self {
        Self {
            mutex: Mutex::new(BarrierState::default()),
            condition: Condvar::new(),
            thread_count,
        }
    }

    fn wait(&self) {
        let mut state = self.mutex.lock().expect("Poisoned mutex");
        let generation = state.generation;
        if state.count + 1 < self.thread_count {
            state.count += 1;
            let (_state, info) = self
                .condition
                .wait_timeout_while(state, Duration::from_secs(5), |state| {
                    state.generation == generation
                })
                .expect("Poisoned mutex");

            if info.timed_out() {
                panic!("Timed out on barrier")
            }
        } else {
            state.count = 0;
            state.generation = generation.wrapping_add(1);
            self.condition.notify_all();
        }
    }
}

#[derive(Default)]
struct BarrierState {
    count: usize,
    generation: usize,
}