1
use std::hint::black_box;
2
use std::sync::Arc;
3
use std::sync::Barrier;
4
use std::sync::atomic::AtomicBool;
5
use std::sync::atomic::Ordering;
6
use std::thread::{self};
7

            
8
use criterion::Criterion;
9
use dashmap::DashSet;
10
use rand::RngExt;
11
use rand::distr::Bernoulli;
12
use rand::distr::Distribution;
13

            
14
use merc_unsafety::ShardedHashMap;
15

            
16
/// The number of operations each thread performs per measured iteration.
17
pub const NUM_ITERATIONS: usize = 50000;
18
/// The number of threads to use for the benchmarks.
19
pub const THREADS: [usize; 6] = [1, 2, 4, 8, 16, 20];
20
/// The (average) number of `get` operations per `insert` operation.
21
pub const READ_RATIOS: [u32; 4] = [1, 10, 100, 1000];
22
/// The percentage of the value space that is pre-populated, controlling how
23
/// many lookups hit and how many insertions find an existing value.
24
pub const EXISTING_PERCENTAGES: [u64; 5] = [0, 50, 95, 99, 100];
25
/// The range `[0, VALUE_SPACE)` of `usize` values operations sample from. Kept
26
/// bounded so the table size stays bounded across criterion's repeated runs.
27
pub const VALUE_SPACE: u64 = 1 << 16;
28

            
29
/// A concurrent set of `usize` values exercised by the benchmark harness.
30
/// `clone` must share the underlying storage so all threads hit one instance.
31
pub trait ConcurrentSet: Clone + Send + 'static {
32
    /// Returns true if `value` is present.
33
    fn get(&self, value: usize) -> bool;
34
    /// Inserts `value`, returning true if it was freshly inserted.
35
    fn insert(&self, value: usize) -> bool;
36
}
37

            
38
/// [`ShardedHashMap`] storing the values directly via its convenience API.
39
#[derive(Clone)]
40
pub struct ShardedSet(Arc<ShardedHashMap<usize>>);
41

            
42
impl ShardedSet {
43
    pub fn new() -> ShardedSet {
44
        ShardedSet(Arc::new(ShardedHashMap::new()))
45
    }
46
}
47

            
48
impl Default for ShardedSet {
49
    fn default() -> ShardedSet {
50
        ShardedSet::new()
51
    }
52
}
53

            
54
impl ConcurrentSet for ShardedSet {
55
    fn get(&self, value: usize) -> bool {
56
        self.0.contains(&value)
57
    }
58

            
59
    fn insert(&self, value: usize) -> bool {
60
        self.0.insert(value).1
61
    }
62
}
63

            
64
/// [`DashSet`] baseline, the concurrent set the rest of the crate uses.
65
#[derive(Clone)]
66
pub struct DashSetWrapper(Arc<DashSet<usize>>);
67

            
68
impl DashSetWrapper {
69
    pub fn new() -> DashSetWrapper {
70
        DashSetWrapper(Arc::new(DashSet::new()))
71
    }
72
}
73

            
74
impl Default for DashSetWrapper {
75
    fn default() -> DashSetWrapper {
76
        DashSetWrapper::new()
77
    }
78
}
79

            
80
impl ConcurrentSet for DashSetWrapper {
81
    fn get(&self, value: usize) -> bool {
82
        self.0.contains(&value)
83
    }
84

            
85
    fn insert(&self, value: usize) -> bool {
86
        self.0.insert(value)
87
    }
88
}
89

            
90
/// Pre-populates `set` so that roughly `existing_percent` of [`VALUE_SPACE`] is
91
/// resident, then drives `num_threads` threads that each perform
92
/// `num_iterations` operations per measured iteration, choosing a `get` over an
93
/// `insert` with odds `read_ratio` to one. Values are sampled uniformly from
94
/// [`VALUE_SPACE`], so `existing_percent` sets the lookup hit rate and the share
95
/// of insertions that find an existing value.
96
pub fn benchmark<S: ConcurrentSet>(
97
    c: &mut Criterion,
98
    name: &str,
99
    set: S,
100
    num_threads: usize,
101
    num_iterations: usize,
102
    read_ratio: u32,
103
    existing_percent: u64,
104
) {
105
    // Pre-populate the lower part of the value space.
106
    let existing = VALUE_SPACE * existing_percent / 100;
107
    for value in 0..existing {
108
        set.insert(value as usize);
109
    }
110

            
111
    // Share threads to avoid spawn overhead between samples.
112
    let mut threads = vec![];
113

            
114
    #[derive(Clone)]
115
    struct ThreadInfo<S> {
116
        busy: Arc<AtomicBool>,
117
        begin_barrier: Arc<Barrier>,
118
        end_barrier: Arc<Barrier>,
119
        dist: Bernoulli,
120
        set: S,
121
    }
122

            
123
    let info = ThreadInfo {
124
        busy: Arc::new(AtomicBool::new(true)),
125
        begin_barrier: Arc::new(Barrier::new(num_threads + 1)),
126
        end_barrier: Arc::new(Barrier::new(num_threads + 1)),
127
        // Sampling true means an insertion, occurring once per `read_ratio` ops.
128
        dist: Bernoulli::from_ratio(1, read_ratio + 1).unwrap(),
129
        set,
130
    };
131

            
132
    for _ in 0..num_threads {
133
        let info = info.clone();
134
        threads.push(thread::spawn(move || {
135
            let mut rng = rand::rng();
136

            
137
            loop {
138
                info.begin_barrier.wait();
139

            
140
                if !info.busy.load(Ordering::SeqCst) {
141
                    // Quit the thread.
142
                    break;
143
                }
144

            
145
                for _ in 0..num_iterations {
146
                    let value = rng.random_range(0..VALUE_SPACE) as usize;
147
                    if info.dist.sample(&mut rng) {
148
                        black_box(info.set.insert(value));
149
                    } else {
150
                        black_box(info.set.get(value));
151
                    }
152
                }
153

            
154
                info.end_barrier.wait();
155
            }
156
        }));
157
    }
158

            
159
    c.bench_function(
160
        format!("{name} threads {num_threads} iterations {num_iterations} read-ratio {read_ratio} existing {existing_percent}%").as_str(),
161
        |bencher| {
162
            bencher.iter(|| {
163
                info.begin_barrier.wait();
164

            
165
                info.end_barrier.wait();
166
            });
167
        },
168
    );
169

            
170
    // Tell the threads to quit and wait for them to join.
171
    info.busy.store(false, Ordering::SeqCst);
172
    info.begin_barrier.wait();
173

            
174
    for thread in threads {
175
        thread.join().unwrap();
176
    }
177
}