dynamo_runtime/storage/
kv.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Interface to a traditional key-value store such as etcd.
5//! "key_value_store" spelt out because in AI land "KV" means something else.
6
7use std::borrow::Cow;
8use std::pin::Pin;
9use std::str::FromStr;
10use std::sync::Arc;
11use std::time::Duration;
12use std::{collections::HashMap, path::PathBuf};
13use std::{env, fmt};
14
15use crate::CancellationToken;
16use crate::transports::etcd as etcd_transport;
17use async_trait::async_trait;
18use futures::StreamExt;
19use percent_encoding::{NON_ALPHANUMERIC, percent_decode_str, percent_encode};
20use serde::{Deserialize, Serialize};
21
22mod mem;
23pub use mem::MemoryStore;
24mod nats;
25pub use nats::NATSStore;
26mod etcd;
27pub use etcd::EtcdStore;
28mod file;
29pub use file::FileStore;
30
31const WATCH_SEND_TIMEOUT: Duration = Duration::from_millis(1000);
32
33/// String we use as the Key in a key-value storage operation. Simple String wrapper
34/// that can encode / decode a string.
35#[derive(Debug, Clone, PartialEq, Eq, Hash)]
36pub struct Key(String);
37
38impl Key {
39    pub fn new(s: String) -> Key {
40        Key(s)
41    }
42
43    /// Takes a URL-safe percent-encoded string and creates a Key from it by decoding first.
44    /// dynamo%2Fbackend%2Fgenerate%2F17216e63492ef21f becomes dynamo/backend/generate/17216e63492ef21f
45    pub fn from_url_safe(s: &str) -> Key {
46        Key(percent_decode_str(s).decode_utf8_lossy().to_string())
47    }
48
49    /// A URL-safe percent-encoded representation of this key.
50    /// e.g.  dynamo/backend/generate/17216e63492ef21f becomes dynamo%2Fbackend%2Fgenerate%2F17216e63492ef21f
51    pub fn url_safe(&self) -> Cow<'_, str> {
52        percent_encode(self.0.as_bytes(), NON_ALPHANUMERIC).into()
53    }
54}
55
56impl From<&str> for Key {
57    fn from(s: &str) -> Key {
58        Key::new(s.to_string())
59    }
60}
61
62impl fmt::Display for Key {
63    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
64        write!(f, "{}", self.0)
65    }
66}
67
68impl AsRef<str> for Key {
69    fn as_ref(&self) -> &str {
70        &self.0
71    }
72}
73
74impl From<&Key> for String {
75    fn from(k: &Key) -> String {
76        k.0.clone()
77    }
78}
79
80#[derive(Debug, Clone, PartialEq)]
81pub struct KeyValue {
82    key: Key,
83    value: bytes::Bytes,
84}
85
86impl KeyValue {
87    pub fn new(key: Key, value: bytes::Bytes) -> Self {
88        KeyValue { key, value }
89    }
90
91    pub fn key(&self) -> String {
92        self.key.clone().to_string()
93    }
94
95    pub fn key_str(&self) -> &str {
96        self.key.as_ref()
97    }
98
99    pub fn value(&self) -> &[u8] {
100        &self.value
101    }
102
103    pub fn value_str(&self) -> anyhow::Result<&str> {
104        std::str::from_utf8(self.value()).map_err(From::from)
105    }
106}
107
108#[derive(Debug, Clone, PartialEq)]
109pub enum WatchEvent {
110    Put(KeyValue),
111    Delete(Key),
112    Resync(HashMap<Key, bytes::Bytes>),
113}
114
115#[async_trait]
116pub trait Store: Send + Sync {
117    type Bucket: Bucket + Send + Sync + 'static;
118
119    async fn get_or_create_bucket(
120        &self,
121        bucket_name: &str,
122        // auto-delete items older than this
123        ttl: Option<Duration>,
124    ) -> Result<Self::Bucket, StoreError>;
125
126    async fn get_bucket(&self, bucket_name: &str) -> Result<Option<Self::Bucket>, StoreError>;
127
128    fn connection_id(&self) -> u64;
129
130    fn shutdown(&self);
131}
132
133#[derive(Clone, Debug, Default)]
134pub enum Selector {
135    // Box it because it is significantly bigger than the other variants
136    Etcd(Box<etcd_transport::ClientOptions>),
137    File(PathBuf),
138    #[default]
139    Memory,
140    // Nats not listed because likely we want to remove that impl. It is not currently used and not well tested.
141}
142
143impl fmt::Display for Selector {
144    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
145        match self {
146            Selector::Etcd(opts) => {
147                let urls = opts.etcd_url.join(",");
148                write!(f, "Etcd({urls})")
149            }
150            Selector::File(path) => write!(f, "File({})", path.display()),
151            Selector::Memory => write!(f, "Memory"),
152        }
153    }
154}
155
156impl FromStr for Selector {
157    type Err = anyhow::Error;
158
159    fn from_str(s: &str) -> anyhow::Result<Selector> {
160        match s {
161            "etcd" => Ok(Self::Etcd(Box::default())),
162            "file" => {
163                let root = env::var("DYN_FILE_KV")
164                    .map(PathBuf::from)
165                    .unwrap_or_else(|_| env::temp_dir().join("dynamo_store_kv"));
166                Ok(Self::File(root))
167            }
168            "mem" => Ok(Self::Memory),
169            x => anyhow::bail!("Unknown key-value store type '{x}'"),
170        }
171    }
172}
173
174impl TryFrom<String> for Selector {
175    type Error = anyhow::Error;
176
177    fn try_from(s: String) -> anyhow::Result<Selector> {
178        s.parse()
179    }
180}
181
182#[allow(clippy::large_enum_variant)]
183enum KeyValueStoreEnum {
184    Memory(MemoryStore),
185    Nats(NATSStore),
186    Etcd(EtcdStore),
187    File(FileStore),
188}
189
190impl KeyValueStoreEnum {
191    async fn get_or_create_bucket(
192        &self,
193        bucket_name: &str,
194        // auto-delete items older than this
195        ttl: Option<Duration>,
196    ) -> Result<Box<dyn Bucket>, StoreError> {
197        use KeyValueStoreEnum::*;
198        Ok(match self {
199            Memory(x) => Box::new(x.get_or_create_bucket(bucket_name, ttl).await?),
200            Nats(x) => Box::new(x.get_or_create_bucket(bucket_name, ttl).await?),
201            Etcd(x) => Box::new(x.get_or_create_bucket(bucket_name, ttl).await?),
202            File(x) => Box::new(x.get_or_create_bucket(bucket_name, ttl).await?),
203        })
204    }
205
206    async fn get_bucket(&self, bucket_name: &str) -> Result<Option<Box<dyn Bucket>>, StoreError> {
207        use KeyValueStoreEnum::*;
208        let maybe_bucket: Option<Box<dyn Bucket>> = match self {
209            Memory(x) => x
210                .get_bucket(bucket_name)
211                .await?
212                .map(|b| Box::new(b) as Box<dyn Bucket>),
213            Nats(x) => x
214                .get_bucket(bucket_name)
215                .await?
216                .map(|b| Box::new(b) as Box<dyn Bucket>),
217            Etcd(x) => x
218                .get_bucket(bucket_name)
219                .await?
220                .map(|b| Box::new(b) as Box<dyn Bucket>),
221            File(x) => x
222                .get_bucket(bucket_name)
223                .await?
224                .map(|b| Box::new(b) as Box<dyn Bucket>),
225        };
226        Ok(maybe_bucket)
227    }
228
229    fn connection_id(&self) -> u64 {
230        use KeyValueStoreEnum::*;
231        match self {
232            Memory(x) => x.connection_id(),
233            Etcd(x) => x.connection_id(),
234            Nats(x) => x.connection_id(),
235            File(x) => x.connection_id(),
236        }
237    }
238
239    fn shutdown(&self) {
240        use KeyValueStoreEnum::*;
241        match self {
242            Memory(x) => x.shutdown(),
243            Etcd(x) => x.shutdown(),
244            Nats(x) => x.shutdown(),
245            File(x) => x.shutdown(),
246        }
247    }
248}
249
250#[derive(Clone)]
251pub struct Manager(Arc<KeyValueStoreEnum>);
252
253impl Default for Manager {
254    fn default() -> Self {
255        Manager::memory()
256    }
257}
258
259impl Manager {
260    /// In-memory KeyValueStoreManager for testing
261    pub fn memory() -> Self {
262        Self::new(KeyValueStoreEnum::Memory(MemoryStore::new()))
263    }
264
265    pub fn etcd(etcd_client: crate::transports::etcd::Client) -> Self {
266        Self::new(KeyValueStoreEnum::Etcd(EtcdStore::new(etcd_client)))
267    }
268
269    pub fn file<P: Into<PathBuf>>(cancel_token: CancellationToken, root: P) -> Self {
270        Self::new(KeyValueStoreEnum::File(FileStore::new(cancel_token, root)))
271    }
272
273    fn new(s: KeyValueStoreEnum) -> Manager {
274        Manager(Arc::new(s))
275    }
276
277    pub async fn get_or_create_bucket(
278        &self,
279        bucket_name: &str,
280        // auto-delete items older than this
281        ttl: Option<Duration>,
282    ) -> Result<Box<dyn Bucket>, StoreError> {
283        self.0.get_or_create_bucket(bucket_name, ttl).await
284    }
285
286    pub async fn get_bucket(
287        &self,
288        bucket_name: &str,
289    ) -> Result<Option<Box<dyn Bucket>>, StoreError> {
290        self.0.get_bucket(bucket_name).await
291    }
292
293    pub fn connection_id(&self) -> u64 {
294        self.0.connection_id()
295    }
296
297    pub async fn load<T: for<'a> Deserialize<'a>>(
298        &self,
299        bucket: &str,
300        key: &Key,
301    ) -> Result<Option<T>, StoreError> {
302        let Some(bucket) = self.0.get_bucket(bucket).await? else {
303            // No bucket means no cards
304            return Ok(None);
305        };
306        Ok(match bucket.get(key).await? {
307            Some(card_bytes) => {
308                let card: T = serde_json::from_slice(card_bytes.as_ref())?;
309                Some(card)
310            }
311            None => None,
312        })
313    }
314
315    /// Returns a receiver that will receive all the existing keys, and
316    /// then block and receive new keys as they are created.
317    /// Starts a task that runs forever, watches the store.
318    pub fn watch(
319        self: Arc<Self>,
320        bucket_name: &str,
321        bucket_ttl: Option<Duration>,
322        cancel_token: CancellationToken,
323    ) -> (
324        tokio::task::JoinHandle<Result<(), StoreError>>,
325        tokio::sync::mpsc::Receiver<WatchEvent>,
326    ) {
327        let bucket_name = bucket_name.to_string();
328        // Use a larger channel capacity to reduce the likelihood that a slow consumer
329        // during the initial KV-store replay phase triggers send timeouts. Events may
330        // still be dropped if the consumer cannot keep up within `WATCH_SEND_TIMEOUT`.
331        let (tx, rx) = tokio::sync::mpsc::channel(16384);
332        let watch_task = tokio::spawn(async move {
333            // Start listening for changes but don't poll this yet
334            let bucket = self
335                .0
336                .get_or_create_bucket(&bucket_name, bucket_ttl)
337                .await?;
338            let mut stream = bucket.watch().await?;
339
340            // Send all the existing keys
341            for (key, bytes) in bucket.entries().await? {
342                if let Err(err) = tx
343                    .send_timeout(
344                        WatchEvent::Put(KeyValue::new(key, bytes)),
345                        WATCH_SEND_TIMEOUT,
346                    )
347                    .await
348                {
349                    tracing::error!(bucket_name, %err, "KeyValueStoreManager.watch failed adding existing key to channel");
350                }
351            }
352
353            // Now block waiting for new entries
354            loop {
355                let event = tokio::select! {
356                    _ = cancel_token.cancelled() => break,
357                    result = stream.next() => match result {
358                        Some(event) => event,
359                        None => break,
360                    }
361                };
362                match event {
363                    WatchEvent::Resync(_) => {
364                        // Resync is an authoritative snapshot; do not drop it on a
365                        // timeout like incremental events. If the receiver closes,
366                        // the watch has no consumer and can stop.
367                        tokio::select! {
368                            _ = cancel_token.cancelled() => break,
369                            result = tx.send(event) => {
370                                if let Err(err) = result {
371                                    tracing::error!(bucket_name, %err, "KeyValueStoreManager.watch failed adding resync to channel");
372                                    break;
373                                }
374                            }
375                        }
376                    }
377                    event => {
378                        if let Err(err) = tx.send_timeout(event, WATCH_SEND_TIMEOUT).await {
379                            tracing::error!(bucket_name, %err, "KeyValueStoreManager.watch failed adding new key to channel");
380                        }
381                    }
382                }
383            }
384
385            Ok::<(), StoreError>(())
386        });
387        (watch_task, rx)
388    }
389
390    pub async fn publish<T: Serialize + Versioned + Send + Sync>(
391        &self,
392        bucket_name: &str,
393        bucket_ttl: Option<Duration>,
394        key: &Key,
395        obj: &mut T,
396    ) -> anyhow::Result<StoreOutcome> {
397        let obj_json = serde_json::to_vec(obj)?;
398        let bucket = self.0.get_or_create_bucket(bucket_name, bucket_ttl).await?;
399
400        let outcome = bucket.insert(key, obj_json.into(), obj.revision()).await?;
401
402        match outcome {
403            StoreOutcome::Created(revision) | StoreOutcome::Exists(revision) => {
404                obj.set_revision(revision);
405            }
406        }
407        Ok(outcome)
408    }
409
410    /// Cleanup any temporary state.
411    /// TODO: Should this be async? Take &mut self?
412    pub fn shutdown(&self) {
413        self.0.shutdown()
414    }
415}
416
417/// An online storage for key-value config values.
418#[async_trait]
419pub trait Bucket: Send + Sync {
420    /// A bucket is a collection of key/value pairs.
421    /// Insert a value into a bucket, if it doesn't exist already
422    /// The Key should be the name of the item, not including the bucket name.
423    async fn insert(
424        &self,
425        key: &Key,
426        value: bytes::Bytes,
427        revision: u64,
428    ) -> Result<StoreOutcome, StoreError>;
429
430    /// Fetch an item from the key-value storage
431    /// The Key should be the name of the item, not including the bucket name.
432    async fn get(&self, key: &Key) -> Result<Option<bytes::Bytes>, StoreError>;
433
434    /// Delete an item from the bucket
435    /// The Key should be the name of the item, not including the bucket name.
436    async fn delete(&self, key: &Key) -> Result<(), StoreError>;
437
438    /// A stream of items inserted into the bucket.
439    /// Every time the stream is polled it will either return a newly created entry, or block until
440    /// such time.
441    async fn watch(
442        &self,
443    ) -> Result<Pin<Box<dyn futures::Stream<Item = WatchEvent> + Send + '_>>, StoreError>;
444
445    /// The entries in this bucket.
446    /// The Key includes the full path including the bucket name.
447    /// That means you cannot directory get a Key from `entries` and pass it to `get` or `delete`.
448    async fn entries(&self) -> Result<HashMap<Key, bytes::Bytes>, StoreError>;
449}
450
451#[derive(Debug, Copy, Clone, Eq, PartialEq)]
452pub enum StoreOutcome {
453    /// The operation succeeded and created a new entry with this revision.
454    /// Note that "create" also means update, because each new revision is a "create".
455    Created(u64),
456    /// The operation did not do anything, the value was already present, with this revision.
457    Exists(u64),
458}
459impl fmt::Display for StoreOutcome {
460    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
461        match self {
462            StoreOutcome::Created(revision) => write!(f, "Created at {revision}"),
463            StoreOutcome::Exists(revision) => write!(f, "Exists at {revision}"),
464        }
465    }
466}
467
468#[derive(thiserror::Error, Debug)]
469pub enum StoreError {
470    #[error("Could not find bucket '{0}'")]
471    MissingBucket(String),
472
473    #[error("Could not find key '{0}'")]
474    MissingKey(String),
475
476    #[error("Internal storage error: '{0}'")]
477    ProviderError(String),
478
479    #[error("Internal NATS error: {0}")]
480    NATSError(String),
481
482    #[error("Internal etcd error: {0}")]
483    EtcdError(String),
484
485    #[error("Internal filesystem error: {0}")]
486    FilesystemError(String),
487
488    #[error("Key Value Error: {0} for bucket '{1}'")]
489    KeyValueError(String, String),
490
491    #[error("Error decoding bytes: {0}")]
492    JSONDecodeError(#[from] serde_json::error::Error),
493
494    #[error("Race condition, retry the call")]
495    Retry,
496}
497
498/// A trait allowing to get/set a revision on an object.
499/// NATS uses this to ensure atomic updates.
500pub trait Versioned {
501    fn revision(&self) -> u64;
502    fn set_revision(&mut self, r: u64);
503}
504
505#[cfg(test)]
506mod tests {
507    use std::sync::Arc;
508
509    use super::*;
510    use futures::{StreamExt, pin_mut};
511
512    const BUCKET_NAME: &str = "v1/mdc";
513
514    /// Convert the value returned by `watch()` into a broadcast stream that multiple
515    /// clients can listen to.
516    #[allow(dead_code)]
517    pub struct TappableStream {
518        tx: tokio::sync::broadcast::Sender<WatchEvent>,
519    }
520
521    #[allow(dead_code)]
522    impl TappableStream {
523        async fn new<T>(stream: T, max_size: usize) -> Self
524        where
525            T: futures::Stream<Item = WatchEvent> + Send + 'static,
526        {
527            let (tx, _) = tokio::sync::broadcast::channel(max_size);
528            let tx2 = tx.clone();
529            tokio::spawn(async move {
530                pin_mut!(stream);
531                while let Some(x) = stream.next().await {
532                    let _ = tx2.send(x);
533                }
534            });
535            TappableStream { tx }
536        }
537
538        fn subscribe(&self) -> tokio::sync::broadcast::Receiver<WatchEvent> {
539            self.tx.subscribe()
540        }
541    }
542
543    fn init() {
544        crate::logging::init();
545    }
546
547    #[tokio::test]
548    async fn test_memory_storage() -> anyhow::Result<()> {
549        init();
550
551        let s = Arc::new(MemoryStore::new());
552        let s2 = Arc::clone(&s);
553
554        let bucket = s.get_or_create_bucket(BUCKET_NAME, None).await?;
555        let res = bucket.insert(&"test1".into(), "value1".into(), 0).await?;
556        assert_eq!(res, StoreOutcome::Created(0));
557
558        let expected = [
559            WatchEvent::Put(KeyValue::new(Key::new("test1".into()), "value1".into())),
560            WatchEvent::Put(KeyValue::new(Key::new("test2".into()), "value2".into())),
561            WatchEvent::Put(KeyValue::new(
562                Key::new("test2".into()),
563                "value2-updated".into(),
564            )),
565            WatchEvent::Put(KeyValue::new(Key::new("test3".into()), "value3".into())),
566        ];
567
568        let (got_first_tx, got_first_rx) = tokio::sync::oneshot::channel();
569        let ingress = tokio::spawn(async move {
570            let b2 = s2.get_or_create_bucket(BUCKET_NAME, None).await?;
571            let mut stream = b2.watch().await?;
572
573            // Put in before starting the watch-all
574            let v = stream.next().await.unwrap();
575            assert_eq!(v, expected[0]);
576
577            got_first_tx.send(()).unwrap();
578
579            // Put in after
580            let v = stream.next().await.unwrap();
581            assert_eq!(v, expected[1]);
582
583            let v = stream.next().await.unwrap();
584            assert_eq!(v, expected[2]);
585
586            let v = stream.next().await.unwrap();
587            assert_eq!(v, expected[3]);
588
589            Ok::<_, StoreError>(())
590        });
591
592        // MemoryStore uses a HashMap with no inherent ordering, so we must ensure test1 is
593        // fetched before test2 is inserted, otherwise they can come out in any order, and we
594        // wouldn't be testing the watch behavior.
595        got_first_rx.await?;
596
597        let res = bucket.insert(&"test2".into(), "value2".into(), 0).await?;
598        assert_eq!(res, StoreOutcome::Created(0));
599
600        // Repeat a key and revision. Ignored.
601        let res = bucket.insert(&"test2".into(), "value2".into(), 0).await?;
602        assert_eq!(res, StoreOutcome::Exists(0));
603
604        // Increment revision
605        let res = bucket
606            .insert(&"test2".into(), "value2-updated".into(), 1)
607            .await?;
608        assert_eq!(res, StoreOutcome::Created(1));
609
610        let res = bucket.insert(&"test3".into(), "value3".into(), 0).await?;
611        assert_eq!(res, StoreOutcome::Created(0));
612
613        // ingress exits once it has received all values
614        let _ = ingress.await?;
615
616        Ok(())
617    }
618
619    #[tokio::test]
620    async fn test_broadcast_stream() -> anyhow::Result<()> {
621        init();
622
623        let s: &'static _ = Box::leak(Box::new(MemoryStore::new()));
624        let bucket: &'static _ =
625            Box::leak(Box::new(s.get_or_create_bucket(BUCKET_NAME, None).await?));
626
627        let res = bucket.insert(&"test1".into(), "value1".into(), 0).await?;
628        assert_eq!(res, StoreOutcome::Created(0));
629
630        let stream = bucket.watch().await?;
631        let tap = TappableStream::new(stream, 10).await;
632
633        let mut rx1 = tap.subscribe();
634        let mut rx2 = tap.subscribe();
635
636        let item = WatchEvent::Put(KeyValue::new(Key::new("test1".to_string()), "GK".into()));
637        let item_clone = item.clone();
638        let handle1 = tokio::spawn(async move {
639            let b = rx1.recv().await.unwrap();
640            assert_eq!(b, item_clone);
641        });
642        let handle2 = tokio::spawn(async move {
643            let b = rx2.recv().await.unwrap();
644            assert_eq!(b, item);
645        });
646
647        bucket.insert(&"test1".into(), "GK".into(), 1).await?;
648
649        let _ = futures::join!(handle1, handle2);
650        Ok(())
651    }
652}
dynamo_runtime/storage/kv.rs

dynamo_runtime/storage/
kv.rs
