dynamo_runtime/storage/kv/
mem.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::collections::HashMap;
5use std::collections::hash_map::Entry;
6use std::pin::Pin;
7use std::sync::Arc;
8use std::time::Duration;
9
10use async_trait::async_trait;
11use rand::Rng as _;
12use tokio::sync::broadcast;
13
14use super::{Bucket, Key, KeyValue, Store, StoreError, StoreOutcome, WatchEvent};
15
16pub(crate) const MEMORY_EVENT_BUFFER_CAPACITY: usize = 16_384;
17
18#[derive(Clone, Debug)]
19enum MemoryEvent {
20    Put {
21        bucket: String,
22        key: String,
23        value: bytes::Bytes,
24    },
25    Delete {
26        bucket: String,
27        key: String,
28    },
29}
30
31#[derive(Clone)]
32pub struct MemoryStore {
33    inner: Arc<MemoryStoreInner>,
34    connection_id: u64,
35}
36
37impl Default for MemoryStore {
38    fn default() -> Self {
39        Self::new()
40    }
41}
42
43struct MemoryStoreInner {
44    data: parking_lot::Mutex<HashMap<String, MemoryBucket>>,
45    change_sender: broadcast::Sender<MemoryEvent>,
46}
47
48pub struct MemoryBucketRef {
49    name: String,
50    inner: Arc<MemoryStoreInner>,
51}
52
53struct MemoryBucket {
54    data: HashMap<String, (u64, bytes::Bytes)>,
55}
56
57impl MemoryBucket {
58    fn new() -> Self {
59        MemoryBucket {
60            data: HashMap::new(),
61        }
62    }
63}
64
65impl MemoryStore {
66    pub(super) fn new() -> Self {
67        let (tx, _) = broadcast::channel(MEMORY_EVENT_BUFFER_CAPACITY);
68        MemoryStore {
69            inner: Arc::new(MemoryStoreInner {
70                data: parking_lot::Mutex::new(HashMap::new()),
71                change_sender: tx,
72            }),
73            connection_id: rand::rng().random(),
74        }
75    }
76}
77
78#[async_trait]
79impl Store for MemoryStore {
80    type Bucket = MemoryBucketRef;
81
82    async fn get_or_create_bucket(
83        &self,
84        bucket_name: &str,
85        // MemoryStore doesn't respect TTL yet
86        _ttl: Option<Duration>,
87    ) -> Result<Self::Bucket, StoreError> {
88        let mut locked_data = self.inner.data.lock();
89        // Ensure the bucket exists
90        locked_data
91            .entry(bucket_name.to_string())
92            .or_insert_with(MemoryBucket::new);
93        // Return an object able to access it
94        Ok(MemoryBucketRef {
95            name: bucket_name.to_string(),
96            inner: self.inner.clone(),
97        })
98    }
99
100    /// This operation cannot fail on MemoryStore. Always returns Ok.
101    async fn get_bucket(&self, bucket_name: &str) -> Result<Option<Self::Bucket>, StoreError> {
102        let locked_data = self.inner.data.lock();
103        match locked_data.get(bucket_name) {
104            Some(_) => Ok(Some(MemoryBucketRef {
105                name: bucket_name.to_string(),
106                inner: self.inner.clone(),
107            })),
108            None => Ok(None),
109        }
110    }
111
112    fn connection_id(&self) -> u64 {
113        self.connection_id
114    }
115
116    fn shutdown(&self) {}
117}
118
119#[async_trait]
120impl Bucket for MemoryBucketRef {
121    async fn insert(
122        &self,
123        key: &Key,
124        value: bytes::Bytes,
125        revision: u64,
126    ) -> Result<StoreOutcome, StoreError> {
127        let mut locked_data = self.inner.data.lock();
128        let mut b = locked_data.get_mut(&self.name);
129        let Some(bucket) = b.as_mut() else {
130            return Err(StoreError::MissingBucket(self.name.to_string()));
131        };
132        let outcome = match bucket.data.entry(key.to_string()) {
133            Entry::Vacant(e) => {
134                e.insert((revision, value.clone()));
135                let _ = self.inner.change_sender.send(MemoryEvent::Put {
136                    bucket: self.name.clone(),
137                    key: key.to_string(),
138                    value,
139                });
140                StoreOutcome::Created(revision)
141            }
142            Entry::Occupied(mut entry) => {
143                let (rev, _v) = entry.get();
144                if revision == 0 || *rev == revision {
145                    StoreOutcome::Exists(*rev)
146                } else {
147                    entry.insert((revision, value.clone()));
148                    let _ = self.inner.change_sender.send(MemoryEvent::Put {
149                        bucket: self.name.clone(),
150                        key: key.to_string(),
151                        value,
152                    });
153                    StoreOutcome::Created(revision)
154                }
155            }
156        };
157        Ok(outcome)
158    }
159
160    async fn compare_and_replace(
161        &self,
162        key: &Key,
163        expected: bytes::Bytes,
164        value: bytes::Bytes,
165    ) -> Result<StoreOutcome, StoreError> {
166        let mut locked_data = self.inner.data.lock();
167        let bucket = locked_data
168            .get_mut(&self.name)
169            .ok_or_else(|| StoreError::MissingBucket(self.name.to_string()))?;
170        let entry = bucket
171            .data
172            .get_mut(&key.0)
173            .ok_or_else(|| StoreError::MissingKey(key.to_string()))?;
174        if entry.1 != expected {
175            return Err(StoreError::Retry);
176        }
177        let next_revision = entry.0.saturating_add(1).max(1);
178        *entry = (next_revision, value.clone());
179        let _ = self.inner.change_sender.send(MemoryEvent::Put {
180            bucket: self.name.clone(),
181            key: key.to_string(),
182            value,
183        });
184        Ok(StoreOutcome::Created(next_revision))
185    }
186
187    async fn get(&self, key: &Key) -> Result<Option<bytes::Bytes>, StoreError> {
188        let locked_data = self.inner.data.lock();
189        let Some(bucket) = locked_data.get(&self.name) else {
190            return Ok(None);
191        };
192        Ok(bucket.data.get(&key.0).map(|(_, v)| v.clone()))
193    }
194
195    async fn delete(&self, key: &Key) -> Result<(), StoreError> {
196        let mut locked_data = self.inner.data.lock();
197        let Some(bucket) = locked_data.get_mut(&self.name) else {
198            return Err(StoreError::MissingBucket(self.name.to_string()));
199        };
200        if bucket.data.remove(&key.0).is_some() {
201            let _ = self.inner.change_sender.send(MemoryEvent::Delete {
202                bucket: self.name.clone(),
203                key: key.to_string(),
204            });
205        }
206        Ok(())
207    }
208
209    async fn watch(
210        &self,
211    ) -> Result<Pin<Box<dyn futures::Stream<Item = WatchEvent> + Send + 'life0>>, StoreError> {
212        // Subscribe while holding the data lock so the snapshot and incremental stream have no
213        // race: mutations take the same lock before broadcasting their update.
214        let data_lock = self.inner.data.lock();
215        let Some(bucket) = data_lock.get(&self.name) else {
216            return Err(StoreError::MissingBucket(self.name.to_string()));
217        };
218        let mut changes = self.inner.change_sender.subscribe();
219        let snapshot: HashMap<Key, bytes::Bytes> = bucket
220            .data
221            .iter()
222            .map(|(key, (_revision, value))| (Key::new(key.clone()), value.clone()))
223            .collect();
224        drop(data_lock);
225        let bucket_name = self.name.clone();
226        let inner = self.inner.clone();
227
228        Ok(Box::pin(async_stream::stream! {
229            yield WatchEvent::Resync(snapshot);
230            loop {
231                match changes.recv().await {
232                    Ok(MemoryEvent::Put { bucket, key, value }) => {
233                        if bucket != bucket_name {
234                            continue;
235                        }
236                        let item = KeyValue::new(Key::new(key), value);
237                        yield WatchEvent::Put(item);
238                    },
239                    Ok(MemoryEvent::Delete { bucket, key }) => {
240                        if bucket != bucket_name {
241                            continue;
242                        }
243                        yield WatchEvent::Delete(Key::new(key));
244                    },
245                    Err(broadcast::error::RecvError::Lagged(_)) => {
246                        let snapshot = {
247                            let data = inner.data.lock();
248                            // Discard retained events that predate this authoritative snapshot.
249                            // Mutations take the same lock before broadcasting, so the replacement
250                            // receiver observes every update that follows the snapshot.
251                            changes = inner.change_sender.subscribe();
252                            data.get(&bucket_name)
253                                .map(|bucket| {
254                                    bucket.data
255                                        .iter()
256                                        .map(|(key, (_revision, value))| {
257                                            (Key::new(key.clone()), value.clone())
258                                        })
259                                        .collect()
260                                })
261                                .unwrap_or_default()
262                        };
263                        yield WatchEvent::Resync(snapshot);
264                    },
265                    Err(broadcast::error::RecvError::Closed) => break,
266                }
267            }
268        }))
269    }
270
271    async fn entries(&self) -> Result<HashMap<Key, bytes::Bytes>, StoreError> {
272        let locked_data = self.inner.data.lock();
273        match locked_data.get(&self.name) {
274            Some(bucket) => {
275                let mut out = HashMap::new();
276                for (k, (_rev, v)) in bucket.data.iter() {
277                    let key = Key::new([self.name.clone(), k.to_string()].join("/"));
278                    let value = v.clone();
279                    out.insert(key, value);
280                }
281                Ok(out)
282            }
283            None => Err(StoreError::MissingBucket(self.name.clone())),
284        }
285    }
286}
287
288#[cfg(test)]
289mod tests {
290    use super::MEMORY_EVENT_BUFFER_CAPACITY;
291    use crate::storage::kv::{
292        Bucket as _, Key, MemoryStore, Store as _, StoreError, StoreOutcome, WatchEvent,
293    };
294    use futures::StreamExt;
295    use std::collections::{HashMap, HashSet};
296    use std::sync::Arc;
297    use std::time::Duration;
298    use tokio::sync::Barrier;
299
300    #[tokio::test]
301    async fn delete_wins_race_with_compare_and_replace() {
302        let store = MemoryStore::new();
303        let bucket = Arc::new(store.get_or_create_bucket("bucket", None).await.unwrap());
304        let key = Key::new("model".to_string());
305        bucket.insert(&key, "old".into(), 0).await.unwrap();
306
307        let barrier = Arc::new(Barrier::new(3));
308        let update_bucket = bucket.clone();
309        let update_key = key.clone();
310        let update_barrier = barrier.clone();
311        let update = tokio::spawn(async move {
312            update_barrier.wait().await;
313            update_bucket
314                .compare_and_replace(&update_key, "old".into(), "new".into())
315                .await
316        });
317        let delete_bucket = bucket.clone();
318        let delete_key = key.clone();
319        let delete_barrier = barrier.clone();
320        let delete = tokio::spawn(async move {
321            delete_barrier.wait().await;
322            delete_bucket.delete(&delete_key).await
323        });
324
325        barrier.wait().await;
326        let update_result = update.await.unwrap();
327        delete.await.unwrap().unwrap();
328        assert!(update_result.is_ok() || matches!(update_result, Err(StoreError::MissingKey(_))));
329        assert_eq!(bucket.get(&key).await.unwrap(), None);
330    }
331
332    #[tokio::test]
333    async fn revision_zero_replay_preserves_cas_update() {
334        let store = MemoryStore::new();
335        let bucket = store.get_or_create_bucket("bucket", None).await.unwrap();
336        let key = Key::new("model".to_string());
337
338        assert_eq!(
339            bucket.insert(&key, "initial".into(), 0).await.unwrap(),
340            StoreOutcome::Created(0)
341        );
342        assert_eq!(
343            bucket
344                .compare_and_replace(&key, "initial".into(), "updated".into())
345                .await
346                .unwrap(),
347            StoreOutcome::Created(1)
348        );
349        assert_eq!(
350            bucket.insert(&key, "initial".into(), 0).await.unwrap(),
351            StoreOutcome::Exists(1)
352        );
353        assert_eq!(
354            bucket.get(&key).await.unwrap().unwrap().as_ref(),
355            b"updated"
356        );
357    }
358
359    #[tokio::test]
360    async fn multiple_watchers_receive_updates_without_cross_bucket_stealing() {
361        let store = MemoryStore::new();
362        let bucket_a = store.get_or_create_bucket("bucket-a", None).await.unwrap();
363        let bucket_b = store.get_or_create_bucket("bucket-b", None).await.unwrap();
364        let mut a_first = bucket_a.watch().await.unwrap();
365        let mut a_second = bucket_a.watch().await.unwrap();
366        let mut b = bucket_b.watch().await.unwrap();
367
368        bucket_a
369            .insert(&Key::new("shared-key".to_string()), "a".into(), 1)
370            .await
371            .unwrap();
372        bucket_b
373            .insert(&Key::new("shared-key".to_string()), "b".into(), 1)
374            .await
375            .unwrap();
376
377        for watcher in [&mut a_first, &mut a_second, &mut b] {
378            assert_eq!(
379                watcher.next().await.unwrap(),
380                WatchEvent::Resync(HashMap::new())
381            );
382        }
383        for watcher in [&mut a_first, &mut a_second] {
384            let WatchEvent::Put(item) = watcher.next().await.unwrap() else {
385                panic!("expected bucket-a put");
386            };
387            assert_eq!(item.value(), b"a");
388        }
389        let WatchEvent::Put(item) = b.next().await.unwrap() else {
390            panic!("expected bucket-b put");
391        };
392        assert_eq!(item.value(), b"b");
393    }
394
395    #[tokio::test]
396    async fn watcher_observes_updates_to_an_existing_key() {
397        let store = MemoryStore::new();
398        let bucket = store.get_or_create_bucket("bucket", None).await.unwrap();
399        bucket
400            .insert(&Key::new("key".to_string()), "old".into(), 1)
401            .await
402            .unwrap();
403        let mut watcher = bucket.watch().await.unwrap();
404        assert_eq!(
405            watcher.next().await.unwrap(),
406            WatchEvent::Resync(HashMap::from([(
407                Key::new("key".to_string()),
408                bytes::Bytes::from("old")
409            )]))
410        );
411
412        bucket
413            .insert(&Key::new("key".to_string()), "new".into(), 2)
414            .await
415            .unwrap();
416        let WatchEvent::Put(item) = watcher.next().await.unwrap() else {
417            panic!("expected updated put");
418        };
419        assert_eq!(item.value(), b"new");
420    }
421
422    #[tokio::test]
423    async fn lag_resync_discards_retained_pre_snapshot_events() {
424        let store = MemoryStore::new();
425        let bucket = store.get_or_create_bucket("bucket", None).await.unwrap();
426        let mut watcher = bucket.watch().await.unwrap();
427        // Consume the startup snapshot so the resync below can only come from the lag.
428        assert_eq!(
429            watcher.next().await.unwrap(),
430            WatchEvent::Resync(HashMap::new())
431        );
432        let key = Key::new("key".to_string());
433
434        for revision in 1..=MEMORY_EVENT_BUFFER_CAPACITY + 1 {
435            bucket
436                .insert(&key, format!("value-{revision}").into(), revision as u64)
437                .await
438                .unwrap();
439        }
440
441        let WatchEvent::Resync(snapshot) = watcher.next().await.unwrap() else {
442            panic!("expected authoritative resync after lag");
443        };
444        assert_eq!(snapshot.len(), 1);
445        let latest_value = format!("value-{}", MEMORY_EVENT_BUFFER_CAPACITY + 1);
446        assert_eq!(snapshot.get(&key).unwrap(), latest_value.as_bytes());
447
448        bucket
449            .insert(
450                &key,
451                "post-resync".into(),
452                (MEMORY_EVENT_BUFFER_CAPACITY + 2) as u64,
453            )
454            .await
455            .unwrap();
456        let next = tokio::time::timeout(Duration::from_secs(1), watcher.next())
457            .await
458            .expect("post-resync update timed out")
459            .unwrap();
460        let WatchEvent::Put(item) = next else {
461            panic!("expected post-resync put");
462        };
463        assert_eq!(item.value(), b"post-resync");
464    }
465
466    #[tokio::test]
467    async fn test_entries_full_path() {
468        let m = MemoryStore::new();
469        let bucket = m.get_or_create_bucket("bucket1", None).await.unwrap();
470        let _ = bucket
471            .insert(&Key::new("key1".to_string()), "value1".into(), 0)
472            .await
473            .unwrap();
474        let _ = bucket
475            .insert(&Key::new("key2".to_string()), "value2".into(), 0)
476            .await
477            .unwrap();
478        let entries = bucket.entries().await.unwrap();
479        let keys: HashSet<Key> = entries.into_keys().collect();
480        assert!(keys.contains(&Key::new("bucket1/key1".to_string())));
481        assert!(keys.contains(&Key::new("bucket1/key2".to_string())));
482    }
483}
dynamo_runtime/storage/kv/mem.rs

dynamo_runtime/storage/kv/
mem.rs
