concurrent_kv_store/
concurrent_kv_store.rs

1//! A tiny in-memory key-value store whose `increment` has a lost-update
2//! race — it reads the current value, releases the lock, and writes
3//! back the incremented value, so two workers incrementing the same key
4//! at once can each write over the other.
5//!
6//! Run it with:
7//!
8//! ```text
9//! cargo test --example concurrent_kv_store
10//! ```
11
12use hegel::TestCase;
13use hegel::generators as gs;
14use hegel::stateful::{ConcurrentPool, concurrent_pool, run_concurrent};
15use std::collections::HashMap;
16use std::sync::Mutex;
17use std::sync::atomic::{AtomicI64, Ordering};
18
19struct KvStore {
20    map: Mutex<HashMap<u64, i64>>,
21}
22
23impl KvStore {
24    fn new() -> Self {
25        KvStore {
26            map: Mutex::new(HashMap::new()),
27        }
28    }
29
30    fn lock(&self) -> std::sync::MutexGuard<'_, HashMap<u64, i64>> {
31        self.map.lock().unwrap_or_else(|e| e.into_inner())
32    }
33
34    fn get(&self, key: u64) -> Option<i64> {
35        self.lock().get(&key).copied()
36    }
37
38    fn put(&self, key: u64, value: i64) {
39        self.lock().insert(key, value);
40    }
41
42    fn put_if_absent(&self, key: u64, value: i64) -> bool {
43        let mut map = self.lock();
44        if map.contains_key(&key) {
45            return false;
46        }
47        map.insert(key, value);
48        true
49    }
50
51    fn increment(&self, key: u64) {
52        let value = self.get(key).unwrap_or(0);
53        std::thread::yield_now(); // Makes it easier to find the race.
54        self.put(key, value + 1);
55    }
56
57    fn snapshot(&self) -> HashMap<u64, i64> {
58        self.lock().clone()
59    }
60}
61
62struct KvTest {
63    store: KvStore,
64    keys: ConcurrentPool<u64>,
65    increments: AtomicI64,
66}
67
68#[hegel::concurrent_state_machine]
69impl KvTest {
70    #[rule(group = "rw")]
71    fn register(&self, tc: TestCase) {
72        let key = tc.draw(gs::integers::<u64>().max_value(3));
73        if self.store.put_if_absent(key, 0) {
74            self.keys.add(&tc, key);
75        }
76    }
77
78    #[rule(group = "rw")]
79    fn increment(&self, tc: TestCase) {
80        let key = tc.draw(self.keys.values_reusable());
81        self.store.increment(key);
82        self.increments.fetch_add(1, Ordering::SeqCst);
83    }
84
85    #[rule(group = "rw")]
86    fn read(&self, tc: TestCase) {
87        let key = tc.draw(self.keys.values_reusable());
88        let value = self.store.get(key);
89        tc.note(&format!("read {key} -> {value:?}"));
90    }
91
92    #[rule(group = "snapshot")]
93    fn snapshot(&self, tc: TestCase) {
94        let snapshot = self.store.snapshot();
95        tc.note(&format!("snapshot holds {} keys", snapshot.len()));
96    }
97
98    #[invariant]
99    fn no_lost_updates(&self, _: TestCase) {
100        let stored: i64 = self.store.snapshot().values().sum();
101        let performed = self.increments.load(Ordering::SeqCst);
102        assert_eq!(
103            stored, performed,
104            "increments were lost: the store sums to {stored} after {performed} increments"
105        );
106    }
107}
108
109#[hegel::test]
110fn test_concurrent_kv_store(tc: TestCase) {
111    let test = KvTest {
112        store: KvStore::new(),
113        keys: concurrent_pool(&tc),
114        increments: AtomicI64::new(0),
115    };
116    run_concurrent(test, tc, 1, 4);
117}
118
119fn main() {}
concurrent_kv_store/concurrent_kv_store.rs

concurrent_kv_store/
concurrent_kv_store.rs
