dynamo_runtime/storage/kv/

nats.rs

// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use std::{collections::HashMap, pin::Pin, time::Duration};

use crate::{protocols::EndpointId, slug::Slug, storage::kv, transports::nats::Client};
use async_nats::jetstream::kv::Operation;
use async_trait::async_trait;
use futures::StreamExt;

use super::{Bucket, Store, StoreError, StoreOutcome};

#[derive(Clone)]
pub struct NATSStore {
    client: Client,
    endpoint: EndpointId,
}

pub struct NATSBucket {
    nats_store: async_nats::jetstream::kv::Store,
}

#[async_trait]
impl Store for NATSStore {
    type Bucket = NATSBucket;

    async fn get_or_create_bucket(
        &self,
        bucket_name: &str,
        ttl: Option<Duration>,
    ) -> Result<Self::Bucket, StoreError> {
        let name = Slug::slugify(bucket_name);
        let nats_store = self
            .get_or_create_key_value(&self.endpoint.namespace, &name, ttl)
            .await?;
        Ok(NATSBucket { nats_store })
    }

    async fn get_bucket(&self, bucket_name: &str) -> Result<Option<Self::Bucket>, StoreError> {
        let name = Slug::slugify(bucket_name);
        match self.get_key_value(&self.endpoint.namespace, &name).await? {
            Some(nats_store) => Ok(Some(NATSBucket { nats_store })),
            None => Ok(None),
        }
    }

    fn connection_id(&self) -> u64 {
        self.client.client().server_info().client_id
    }

    fn shutdown(&self) {
        // TODO: Track and delete any owned keys
        // The TTL should ensure NATS does it, but best we do it immediately
    }

    async fn check_connection(&self) -> Result<(), StoreError> {
        if self.client.client().connection_state() != async_nats::connection::State::Connected {
            return Err(StoreError::NATSError("NATS is disconnected".into()));
        }
        Ok(())
    }
}

impl NATSStore {
    pub fn new(client: Client, endpoint: EndpointId) -> Self {
        NATSStore { client, endpoint }
    }

    /// Get or create a key-value store (aka bucket) in NATS.
    ///
    /// ttl is only used if we are creating the bucket, so if that has
    /// changed first delete the bucket.
    async fn get_or_create_key_value(
        &self,
        namespace: &str,
        bucket_name: &Slug,
        // Delete entries older than this
        ttl: Option<Duration>,
    ) -> Result<async_nats::jetstream::kv::Store, StoreError> {
        if let Ok(Some(kv)) = self.get_key_value(namespace, bucket_name).await {
            return Ok(kv);
        }

        // It doesn't exist, create it

        let bucket_name = single_name(namespace, bucket_name);
        let js = self.client.jetstream();
        let create_result = js
            .create_key_value(
                // TODO: configure the bucket, probably need to pass some of these values in
                async_nats::jetstream::kv::Config {
                    bucket: bucket_name.clone(),
                    max_age: ttl.unwrap_or_default(),
                    ..Default::default()
                },
            )
            .await;
        let nats_store = create_result
            .map_err(|err| StoreError::KeyValueError(err.to_string(), bucket_name.clone()))?;
        tracing::debug!("Created bucket {bucket_name}");
        Ok(nats_store)
    }

    async fn get_key_value(
        &self,
        namespace: &str,
        bucket_name: &Slug,
    ) -> Result<Option<async_nats::jetstream::kv::Store>, StoreError> {
        let bucket_name = single_name(namespace, bucket_name);
        let js = self.client.jetstream();

        use async_nats::jetstream::context::KeyValueErrorKind;
        match js.get_key_value(&bucket_name).await {
            Ok(store) => Ok(Some(store)),
            Err(err) if err.kind() == KeyValueErrorKind::GetBucket => {
                // bucket doesn't exist
                Ok(None)
            }
            Err(err) => Err(StoreError::KeyValueError(err.to_string(), bucket_name)),
        }
    }
}

#[async_trait]
impl Bucket for NATSBucket {
    async fn insert(
        &self,
        key: &kv::Key,
        value: bytes::Bytes,
        revision: u64,
    ) -> Result<StoreOutcome, StoreError> {
        if revision == 0 {
            self.create(key, value).await
        } else {
            self.update(key, value, revision).await
        }
    }

    async fn compare_and_replace(
        &self,
        key: &kv::Key,
        expected: bytes::Bytes,
        value: bytes::Bytes,
    ) -> Result<StoreOutcome, StoreError> {
        let entry = self
            .nats_store
            .entry(key)
            .await
            .map_err(|error| StoreError::NATSError(error.to_string()))?
            .ok_or_else(|| StoreError::MissingKey(key.to_string()))?;
        if matches!(entry.operation, Operation::Delete | Operation::Purge) {
            return Err(StoreError::MissingKey(key.to_string()));
        }
        if entry.value != expected {
            return Err(StoreError::Retry);
        }

        match self.nats_store.update(key, value, entry.revision).await {
            Ok(revision) => Ok(StoreOutcome::Created(revision)),
            Err(error)
                if error.kind()
                    == async_nats::jetstream::kv::UpdateErrorKind::WrongLastRevision =>
            {
                match self.nats_store.entry(key).await {
                    Ok(None) => Err(StoreError::MissingKey(key.to_string())),
                    Ok(Some(entry))
                        if matches!(entry.operation, Operation::Delete | Operation::Purge) =>
                    {
                        Err(StoreError::MissingKey(key.to_string()))
                    }
                    Ok(Some(_)) => Err(StoreError::Retry),
                    Err(error) => Err(StoreError::NATSError(error.to_string())),
                }
            }
            Err(error) => Err(StoreError::NATSError(error.to_string())),
        }
    }

    async fn get(&self, key: &kv::Key) -> Result<Option<bytes::Bytes>, StoreError> {
        self.nats_store
            .get(key)
            .await
            .map_err(|e| StoreError::NATSError(e.to_string()))
    }

    async fn delete(&self, key: &kv::Key) -> Result<(), StoreError> {
        self.nats_store
            .delete(key)
            .await
            .map_err(|e| StoreError::NATSError(e.to_string()))
    }

    /// Replays the existing entries as individual [`kv::WatchEvent::Put`] events, not as the
    /// initial [`kv::WatchEvent::Resync`] the [`kv::Bucket::watch`] contract asks for. Discovery
    /// does not select this store, so nothing depends on that order here.
    async fn watch(
        &self,
    ) -> Result<Pin<Box<dyn futures::Stream<Item = kv::WatchEvent> + Send + 'life0>>, StoreError>
    {
        let watch_stream = self
            .nats_store
            .watch_with_history(">")
            .await
            .map_err(|e| StoreError::NATSError(e.to_string()))?;
        // Map the `Entry` to `Entry.value` which is Bytes of the stored value.
        Ok(Box::pin(
            watch_stream.filter_map(
                |maybe_entry: Result<
                    async_nats::jetstream::kv::Entry,
                    async_nats::error::Error<_>,
                >| async move {
                    match maybe_entry {
                        Ok(entry) => {
                            let key = kv::Key::new(entry.key);
                            Some(match entry.operation {
                                Operation::Put => {
                                    let item = kv::KeyValue::new(key, entry.value);
                                    kv::WatchEvent::Put(item)
                                }
                                Operation::Delete => kv::WatchEvent::Delete(key),
                                // TODO: What is Purge? Not urgent, NATS impl not used
                                Operation::Purge => kv::WatchEvent::Delete(key),
                            })
                        }
                        Err(e) => {
                            tracing::error!(error=%e, "watch fatal err");
                            None
                        }
                    }
                },
            ),
        ))
    }

    async fn entries(&self) -> Result<HashMap<kv::Key, bytes::Bytes>, StoreError> {
        let mut key_stream = self
            .nats_store
            .keys()
            .await
            .map_err(|e| StoreError::NATSError(e.to_string()))?;
        let mut out = HashMap::new();
        while let Some(Ok(key)) = key_stream.next().await {
            if let Ok(Some(entry)) = self.nats_store.entry(&key).await {
                out.insert(kv::Key::new(key), entry.value);
            }
        }
        Ok(out)
    }
}

impl NATSBucket {
    async fn create(&self, key: &kv::Key, value: bytes::Bytes) -> Result<StoreOutcome, StoreError> {
        match self.nats_store.create(&key, value).await {
            Ok(revision) => Ok(StoreOutcome::Created(revision)),
            Err(err) if err.kind() == async_nats::jetstream::kv::CreateErrorKind::AlreadyExists => {
                // key exists, get the revsion
                match self.nats_store.entry(key).await {
                    Ok(Some(entry)) => Ok(StoreOutcome::Exists(entry.revision)),
                    Ok(None) => {
                        tracing::error!(
                            %key,
                            "Race condition, key deleted between create and fetch. Retry."
                        );
                        Err(StoreError::Retry)
                    }
                    Err(err) => Err(StoreError::NATSError(err.to_string())),
                }
            }
            Err(err) => Err(StoreError::NATSError(err.to_string())),
        }
    }

    async fn update(
        &self,
        key: &kv::Key,
        value: bytes::Bytes,
        revision: u64,
    ) -> Result<StoreOutcome, StoreError> {
        match self.nats_store.update(key, value.clone(), revision).await {
            Ok(revision) => Ok(StoreOutcome::Created(revision)),
            Err(err)
                if err.kind() == async_nats::jetstream::kv::UpdateErrorKind::WrongLastRevision =>
            {
                tracing::warn!(revision, %key, "Update WrongLastRevision, resync");
                self.resync_update(key, value).await
            }
            Err(err) => Err(StoreError::NATSError(err.to_string())),
        }
    }

    /// We have the wrong revision for a key. Fetch it's entry to get the correct revision,
    /// and try the update again.
    async fn resync_update(
        &self,
        key: &kv::Key,
        value: bytes::Bytes,
    ) -> Result<StoreOutcome, StoreError> {
        match self.nats_store.entry(key).await {
            Ok(Some(entry)) => {
                // Re-try the update with new version number
                let next_rev = entry.revision + 1;
                match self.nats_store.update(key, value, next_rev).await {
                    Ok(correct_revision) => Ok(StoreOutcome::Created(correct_revision)),
                    Err(err) => Err(StoreError::NATSError(format!(
                        "Error during update of key {key} after resync: {err}"
                    ))),
                }
            }
            Ok(None) => {
                tracing::warn!(%key, "Entry does not exist during resync, creating.");
                self.create(key, value).await
            }
            Err(err) => {
                tracing::error!(%key, %err, "Failed fetching entry during resync");
                Err(StoreError::NATSError(err.to_string()))
            }
        }
    }
}

/// async-nats won't let us use a multi-part subject to create KV buckets (and probably many other
/// things).
fn single_name(namespace: &str, name: &Slug) -> String {
    format!("{namespace}_{name}")
}

#[cfg(feature = "integration")]
#[cfg(test)]
mod compare_and_replace_tests {
    use std::sync::Arc;

    use tokio::sync::Barrier;

    use super::*;
    use crate::storage::kv::{Bucket as _, Key, Store as _};

    #[tokio::test]
    async fn delete_wins_race_with_compare_and_replace() {
        let client = crate::transports::nats::ClientOptions::default()
            .connect()
            .await
            .unwrap();
        let endpoint = EndpointId {
            namespace: "test".to_string(),
            component: "storage".to_string(),
            name: "compare-and-replace".to_string(),
        };
        let store = NATSStore::new(client, endpoint);
        let bucket_name = format!("compare_and_replace_{}", uuid::Uuid::new_v4());
        let jetstream_bucket_name = single_name("test", &Slug::slugify(&bucket_name));
        let bucket = Arc::new(
            store
                .get_or_create_bucket(&bucket_name, None)
                .await
                .unwrap(),
        );
        let key = Key::new("model".to_string());
        bucket.insert(&key, "old".into(), 0).await.unwrap();

        let barrier = Arc::new(Barrier::new(3));
        let task_bucket = bucket.clone();
        let task_key = key.clone();
        let task_barrier = barrier.clone();
        let update = tokio::spawn(async move {
            task_barrier.wait().await;
            task_bucket
                .compare_and_replace(&task_key, "old".into(), "new".into())
                .await
        });
        let task_bucket = bucket.clone();
        let task_key = key.clone();
        let task_barrier = barrier.clone();
        let delete = tokio::spawn(async move {
            task_barrier.wait().await;
            task_bucket.delete(&task_key).await
        });

        barrier.wait().await;
        let update_result = update.await.unwrap();
        delete.await.unwrap().unwrap();
        assert!(update_result.is_ok() || matches!(update_result, Err(StoreError::MissingKey(_))));
        assert_eq!(bucket.get(&key).await.unwrap(), None);

        drop(bucket);
        store
            .client
            .jetstream()
            .delete_key_value(&jetstream_bucket_name)
            .await
            .unwrap();
    }
}