hightower_kv/
state.rs

1use hashbrown::HashMap;
2use hashbrown::hash_map::Entry;
3use parking_lot::Mutex;
4use std::collections::hash_map::DefaultHasher;
5use std::hash::{Hash, Hasher};
6use std::sync::atomic::{AtomicUsize, Ordering};
7
8use crate::command::Command;
9
10const DEFAULT_SHARDS: usize = 64;
11
12/// Single-threaded key-value state with versioned entries
13#[derive(Debug, Default)]
14pub struct KvState {
15    entries: HashMap<Vec<u8>, (Vec<u8>, u64)>,
16}
17
18/// Outcome of applying a command to the state
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum ApplyOutcome {
21    /// Command was successfully applied
22    Applied,
23    /// Command was ignored because it was stale
24    IgnoredStale,
25    /// Key was removed from the state
26    Removed,
27}
28
29impl KvState {
30    /// Creates a new empty key-value state
31    pub fn new() -> Self {
32        Self::default()
33    }
34
35    /// Returns the number of entries in the state
36    pub fn len(&self) -> usize {
37        self.entries.len()
38    }
39
40    /// Gets the value for the given key
41    pub fn get(&self, key: &[u8]) -> Option<&[u8]> {
42        self.entries
43            .get(key)
44            .map(|(value, _version)| value.as_slice())
45    }
46
47    /// Evaluates what would happen if the command were applied without mutating state
48    pub fn evaluate(&self, command: &Command) -> ApplyOutcome {
49        match command {
50            Command::Set { key, version, .. } => match self.entries.get(key) {
51                Some((_, current_version)) if *version <= *current_version => {
52                    ApplyOutcome::IgnoredStale
53                }
54                _ => ApplyOutcome::Applied,
55            },
56            Command::Delete { key, version, .. } => match self.entries.get(key) {
57                Some((_, current_version)) if *version > *current_version => ApplyOutcome::Removed,
58                Some(_) => ApplyOutcome::IgnoredStale,
59                None => ApplyOutcome::Removed,
60            },
61        }
62    }
63
64    /// Applies a command to the state and returns the outcome
65    pub fn apply(&mut self, command: &Command) -> ApplyOutcome {
66        match command {
67            Command::Set {
68                key,
69                value,
70                version,
71                ..
72            } => {
73                let outcome = self.evaluate(command);
74                if matches!(outcome, ApplyOutcome::Applied) {
75                    self.entries.insert(key.clone(), (value.clone(), *version));
76                }
77                outcome
78            }
79            Command::Delete { key, .. } => {
80                let outcome = self.evaluate(command);
81                if matches!(outcome, ApplyOutcome::Removed) {
82                    self.entries.remove(key);
83                }
84                outcome
85            }
86        }
87    }
88
89    /// Inserts a snapshot entry directly into the state without version checking
90    pub fn insert_snapshot(&mut self, key: Vec<u8>, value: Vec<u8>, version: u64) {
91        self.entries.insert(key, (value, version));
92    }
93
94    /// Returns an iterator over all entries in the state
95    pub fn iter(&self) -> impl Iterator<Item = (&Vec<u8>, &(Vec<u8>, u64))> {
96        self.entries.iter()
97    }
98
99    /// Consumes the state and returns the underlying entries map
100    pub fn into_entries(self) -> HashMap<Vec<u8>, (Vec<u8>, u64)> {
101        self.entries
102    }
103}
104
105/// Thread-safe sharded key-value state for concurrent access
106#[derive(Debug)]
107pub struct ConcurrentKvState {
108    shards: Vec<Shard>,
109}
110
111impl Default for ConcurrentKvState {
112    fn default() -> Self {
113        Self::new()
114    }
115}
116
117impl ConcurrentKvState {
118    /// Creates a new concurrent key-value state with default shard count
119    pub fn new() -> Self {
120        Self::with_shard_count(DEFAULT_SHARDS)
121    }
122
123    fn with_shard_count(count: usize) -> Self {
124        let shard_count = count.max(1);
125        let mut shards = Vec::with_capacity(shard_count);
126        for _ in 0..shard_count {
127            shards.push(Shard::new());
128        }
129        Self { shards }
130    }
131
132    /// Returns the total number of entries across all shards
133    pub fn len(&self) -> usize {
134        self.shards
135            .iter()
136            .map(|shard| shard.len.load(Ordering::Relaxed))
137            .sum()
138    }
139
140    /// Gets a cloned value for the given key
141    pub fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
142        let shard = self.shard_for(key);
143        let guard = shard.entries.lock();
144        guard.get(key).map(|(value, _)| value.clone())
145    }
146
147    /// Locks the shard containing the given key and returns an entry guard
148    pub fn lock_entry<'a>(&'a self, key: &[u8]) -> EntryGuard<'a> {
149        let shard = self.shard_for(key);
150        let guard = shard.entries.lock();
151        EntryGuard { shard, guard }
152    }
153
154    /// Executes a read operation on a snapshot of the entire state
155    pub fn read_with<F, R>(&self, reader: F) -> R
156    where
157        F: FnOnce(&KvState) -> R,
158    {
159        let snapshot = self.to_snapshot();
160        reader(&snapshot)
161    }
162
163    /// Creates a consistent snapshot of the entire state
164    pub fn to_snapshot(&self) -> KvState {
165        let mut snapshot = KvState::new();
166        for shard in &self.shards {
167            let guard = shard.entries.lock();
168            for (key, (value, version)) in guard.iter() {
169                snapshot.insert_snapshot(key.clone(), value.clone(), *version);
170            }
171        }
172        snapshot
173    }
174
175    /// Inserts a snapshot entry directly into the state
176    pub fn insert_snapshot(&self, key: Vec<u8>, value: Vec<u8>, version: u64) {
177        let shard = self.shard_for(&key);
178        let mut guard = shard.entries.lock();
179        if guard.insert(key, (value, version)).is_none() {
180            shard.len.fetch_add(1, Ordering::Relaxed);
181        }
182    }
183
184    fn shard_for(&self, key: &[u8]) -> &Shard {
185        let mut hasher = DefaultHasher::new();
186        key.hash(&mut hasher);
187        let idx = (hasher.finish() as usize) % self.shards.len();
188        &self.shards[idx]
189    }
190
191    #[cfg(test)]
192    pub(crate) fn clear_for_test(&self) {
193        for shard in &self.shards {
194            let mut guard = shard.entries.lock();
195            guard.clear();
196            shard.len.store(0, Ordering::Relaxed);
197        }
198    }
199}
200
201impl From<KvState> for ConcurrentKvState {
202    fn from(state: KvState) -> Self {
203        let concurrent = ConcurrentKvState::new();
204        for (key, (value, version)) in state.into_entries() {
205            concurrent.insert_snapshot(key, value, version);
206        }
207        concurrent
208    }
209}
210
211/// Guard for a locked shard entry that allows atomic operations
212pub struct EntryGuard<'a> {
213    shard: &'a Shard,
214    guard: parking_lot::MutexGuard<'a, HashMap<Vec<u8>, (Vec<u8>, u64)>>,
215}
216
217impl<'a> EntryGuard<'a> {
218    /// Evaluates what would happen if the command were applied without mutating state
219    pub fn evaluate(&self, command: &Command) -> ApplyOutcome {
220        match command {
221            Command::Set { key, version, .. } => match self.guard.get(key) {
222                Some((_, current_version)) if *version <= *current_version => {
223                    ApplyOutcome::IgnoredStale
224                }
225                _ => ApplyOutcome::Applied,
226            },
227            Command::Delete { key, version, .. } => match self.guard.get(key) {
228                Some((_, current_version)) if *version > *current_version => ApplyOutcome::Removed,
229                Some(_) => ApplyOutcome::IgnoredStale,
230                None => ApplyOutcome::Removed,
231            },
232        }
233    }
234
235    /// Applies a command to the locked entry and returns the outcome
236    pub fn apply(&mut self, command: &Command) -> ApplyOutcome {
237        match command {
238            Command::Set {
239                key,
240                value,
241                version,
242                ..
243            } => {
244                let outcome = self.evaluate(command);
245                if matches!(outcome, ApplyOutcome::Applied) {
246                    match self.guard.entry(key.clone()) {
247                        Entry::Occupied(mut entry) => {
248                            entry.insert((value.clone(), *version));
249                        }
250                        Entry::Vacant(entry) => {
251                            entry.insert((value.clone(), *version));
252                            self.shard.len.fetch_add(1, Ordering::Relaxed);
253                        }
254                    }
255                }
256                outcome
257            }
258            Command::Delete { key, .. } => {
259                let outcome = self.evaluate(command);
260                if matches!(outcome, ApplyOutcome::Removed) {
261                    if self.guard.remove(key).is_some() {
262                        self.shard.len.fetch_sub(1, Ordering::Relaxed);
263                    }
264                }
265                outcome
266            }
267        }
268    }
269}
270
271#[derive(Debug)]
272struct Shard {
273    entries: Mutex<HashMap<Vec<u8>, (Vec<u8>, u64)>>,
274    len: AtomicUsize,
275}
276
277impl Shard {
278    fn new() -> Self {
279        Self {
280            entries: Mutex::new(HashMap::new()),
281            len: AtomicUsize::new(0),
282        }
283    }
284}
285
286#[cfg(test)]
287mod tests {
288    use super::*;
289
290    #[test]
291    fn set_and_get_roundtrip() {
292        let mut state = KvState::new();
293        let cmd = Command::Set {
294            key: b"key".to_vec(),
295            value: b"value".to_vec(),
296            version: 1,
297            timestamp: 1,
298        };
299        assert_eq!(state.apply(&cmd), ApplyOutcome::Applied);
300        assert_eq!(state.get(b"key"), Some(&b"value"[..]));
301        assert_eq!(state.len(), 1);
302    }
303
304    #[test]
305    fn stale_write_is_ignored() {
306        let mut state = KvState::new();
307        let first = Command::Set {
308            key: b"key".to_vec(),
309            value: b"v1".to_vec(),
310            version: 5,
311            timestamp: 5,
312        };
313        let stale = Command::Set {
314            key: b"key".to_vec(),
315            value: b"v0".to_vec(),
316            version: 4,
317            timestamp: 4,
318        };
319        assert_eq!(state.apply(&first), ApplyOutcome::Applied);
320        assert_eq!(state.apply(&stale), ApplyOutcome::IgnoredStale);
321        assert_eq!(state.get(b"key"), Some(&b"v1"[..]));
322    }
323
324    #[test]
325    fn delete_removes_newer_versions() {
326        let mut state = KvState::new();
327        let set = Command::Set {
328            key: b"key".to_vec(),
329            value: b"v1".to_vec(),
330            version: 5,
331            timestamp: 5,
332        };
333        let delete = Command::Delete {
334            key: b"key".to_vec(),
335            version: 6,
336            timestamp: 6,
337        };
338        assert_eq!(state.apply(&set), ApplyOutcome::Applied);
339        assert_eq!(state.apply(&delete), ApplyOutcome::Removed);
340        assert!(state.get(b"key").is_none());
341    }
342
343    #[test]
344    fn concurrent_state_handles_set_and_get() {
345        let state = ConcurrentKvState::new();
346        let command = Command::Set {
347            key: b"k".to_vec(),
348            value: b"v".to_vec(),
349            version: 1,
350            timestamp: 1,
351        };
352        {
353            let mut guard = state.lock_entry(command.key());
354            assert_eq!(guard.evaluate(&command), ApplyOutcome::Applied);
355            assert_eq!(guard.apply(&command), ApplyOutcome::Applied);
356        }
357        assert_eq!(state.get(b"k"), Some(b"v".to_vec()));
358    }
359
360    #[test]
361    fn concurrent_state_respects_versions() {
362        let state = ConcurrentKvState::new();
363        let first = Command::Set {
364            key: b"k".to_vec(),
365            value: b"one".to_vec(),
366            version: 2,
367            timestamp: 1,
368        };
369        let stale = Command::Set {
370            key: b"k".to_vec(),
371            value: b"zero".to_vec(),
372            version: 1,
373            timestamp: 0,
374        };
375        {
376            let mut guard = state.lock_entry(first.key());
377            assert_eq!(guard.apply(&first), ApplyOutcome::Applied);
378        }
379        {
380            let mut guard = state.lock_entry(stale.key());
381            assert_eq!(guard.apply(&stale), ApplyOutcome::IgnoredStale);
382        }
383        assert_eq!(state.get(b"k"), Some(b"one".to_vec()));
384    }
385
386    #[test]
387    fn snapshot_round_trip_between_states() {
388        let mut snapshot = KvState::new();
389        snapshot.insert_snapshot(b"a".to_vec(), b"1".to_vec(), 1);
390        snapshot.insert_snapshot(b"b".to_vec(), b"2".to_vec(), 2);
391
392        let concurrent: ConcurrentKvState = snapshot.into();
393        let reflect = concurrent.to_snapshot();
394        assert_eq!(reflect.len(), 2);
395        assert_eq!(reflect.get(b"a"), Some(&b"1"[..]));
396        assert_eq!(reflect.get(b"b"), Some(&b"2"[..]));
397    }
398}
399
400#[cfg(test)]
401impl KvState {
402    pub(crate) fn clear_for_test(&mut self) {
403        self.entries.clear();
404    }
405}
hightower_kv/state.rs

hightower_kv/
state.rs
