Cosmos: Container Metadata Cache

Cargo.toml

@@ -109,6 +109,7 @@ getrandom = { version = "0.3" }
 gloo-timers = { version = "0.3" }
 hmac = { version = "0.12" }
 litemap = "0.7.4"
+moka = { version = "0.12.11", features = ["future"] }
 openssl = { version = "0.10.72" }
 opentelemetry = { version = "0.30", features = ["trace"] }
 opentelemetry_sdk = "0.30"

eng/dict/crates.txt

@@ -1,10 +1,16 @@
 async-lock
 async-stream
 async-trait
+azure_canary
+azure_canary_core
 azure_core
 azure_core
 azure_core_amqp
 azure_core_amqp
+azure_core_macros
+azure_core_macros
+azure_core_opentelemetry
+azure_core_opentelemetry
 azure_core_test
 azure_core_test
 azure_core_test_macros
@@ -13,16 +19,21 @@ azure_data_cosmos
 azure_identity
 azure_identity
 azure_messaging_eventhubs
+azure_messaging_eventhubs
+azure_messaging_eventhubs_checkpointstore_blob
+azure_messaging_servicebus
+azure_security_keyvault_certificates
 azure_security_keyvault_keys
 azure_security_keyvault_secrets
 azure_storage_blob
-azure_canary
-azure_canary_core
+azure_storage_blob
+azure_storage_common
+azure_storage_queue
 base64
 bytes
 cargo_metadata
 clap
-dotenvy
+criterion
 dyn-clone
 fe2o3-amqp
 fe2o3-amqp-cbs
@@ -32,18 +43,22 @@ fe2o3-amqp-types
 flate2
 futures
 getrandom
-gloo
+gloo-timers
 hmac
 litemap
-log
+moka
 openssl
+opentelemetry
+opentelemetry-http
+opentelemetry_sdk
 pin-project
 proc-macro2
 quick-xml
 quote
 rand
 rand_chacha
 reqwest
+rust_decimal
 rustc_version
 serde
 serde_amqp
@@ -52,7 +67,6 @@ serde_json
 serde_test
 serial_test
 sha2
-storage
 syn
 tar
 thiserror
@@ -66,8 +80,9 @@ typespec_client_core
 typespec_client_core
 typespec_macros
 typespec_macros
-ureq
 url
 uuid
+wasm-bindgen-futures
+wasm-bindgen-test
 zerofrom
 zip

sdk/cosmos/assets.json

@@ -1,6 +1,6 @@
 {
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "rust",
-  "Tag": "rust/azure_data_cosmos_a39b424a5b",
+  "Tag": "rust/azure_data_cosmos_69ad1e4995",
   "TagPrefix": "rust/azure_data_cosmos"
 }

sdk/cosmos/azure_data_cosmos/Cargo.toml

@@ -22,6 +22,7 @@ serde.workspace = true
 tracing.workspace = true
 typespec_client_core = { workspace = true, features = ["derive"] }
 url.workspace = true
+moka.workspace = true
 
 [dev-dependencies]
 azure_identity.workspace = true

sdk/cosmos/azure_data_cosmos/src/cache.rs

@@ -0,0 +1,124 @@
+use std::sync::Arc;
+
+use moka::future::Cache;
+
+use crate::{models::ContainerProperties, resource_context::ResourceLink, ResourceId};
+
+#[derive(Debug)]
+pub enum CacheError {
+    FetchError(Arc<azure_core::Error>),
+}
+
+impl From<Arc<azure_core::Error>> for CacheError {
+    fn from(e: Arc<azure_core::Error>) -> Self {
+        CacheError::FetchError(e)
+    }
+}
+
+impl From<CacheError> for azure_core::Error {
+    fn from(e: CacheError) -> Self {
+        match e {
+            CacheError::FetchError(e) => {
+                let message = format!("error updating Container Metadata Cache: {}", e);
+                azure_core::Error::with_error(azure_core::error::ErrorKind::Other, e, message)
+            }
+        }
+    }
+}
+
+impl std::fmt::Display for CacheError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            CacheError::FetchError(e) => write!(f, "error fetching latest value: {}", e),
+        }
+    }
+}
+
+impl std::error::Error for CacheError {
+    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        match self {
+            CacheError::FetchError(e) => Some(&**e),
+        }
+    }
+}
+
+/// A subset of container properties that are stable and suitable for caching.
+pub(crate) struct ContainerMetadata {
+    pub resource_id: ResourceId,
+    pub container_link: ResourceLink,
+}
+
+impl ContainerMetadata {
+    // We can't use From<ContainerProperties> because we also want the container link.
+    pub fn from_properties(
+        properties: &ContainerProperties,
+        container_link: ResourceLink,
+    ) -> azure_core::Result<Self> {
+        let resource_id = properties
+            .system_properties
+            .resource_id
+            .clone()
+            .ok_or_else(|| {
+                azure_core::Error::new(
+                    azure_core::error::ErrorKind::Other,
+                    "container properties is missing expected value 'resource_id'",
+                )
+            })?;
+        Ok(Self {
+            resource_id,
+            container_link,
+        })
+    }
+}
+
+/// A cache for container metadata, including properties and routing information.
+///
+/// The cache can be cloned cheaply, and all clones share the same underlying cache data.
+#[derive(Clone)]
+pub struct ContainerMetadataCache {
+    /// Caches stable container metadata, mapping from container link to metadata.
+    container_properties_cache: Cache<ResourceLink, Arc<ContainerMetadata>>,
+}
+
+// TODO: Review this value.
+// Cosmos has a backend limit of 500 databases and containers per account by default.
+// This value affects when Moka will start evicting entries from the cache.
+// It could probably be much lower without much impact, but we need to do the research to be sure.
+const MAX_CACHE_CAPACITY: u64 = 500;
+
+impl ContainerMetadataCache {
+    /// Creates a new `ContainerMetadataCache` with default settings.
+    ///
+    /// Since the cache is designed to be shared, it is returned inside an `Arc`.
+    pub fn new() -> Self {
+        let container_properties_cache = Cache::new(MAX_CACHE_CAPACITY);
+        Self {
+            container_properties_cache,
+        }
+    }
+
+    /// Unconditionally updates the cache with the provided container metadata.
+    pub async fn set_container_metadata(&self, metadata: ContainerMetadata) {
+        let metadata = Arc::new(metadata);
+
+        self.container_properties_cache
+            .insert(metadata.container_link.clone(), metadata)
+            .await;
+    }
+
+    /// Gets the container metadata from the cache, or initializes it using the provided async function if not present.
+    pub async fn get_container_metadata(
+        &self,
+        key: &ResourceLink,
+        init: impl std::future::Future<Output = azure_core::Result<ContainerMetadata>>,
+    ) -> Result<Arc<ContainerMetadata>, CacheError> {
+        // TODO: Background refresh. We can do background refresh by storing an expiry time in the cache entry.
+        // Then, if the entry is stale, we can return the stale entry and spawn a background task to refresh it.
+        // There's a little trickiness here in that we can't directly spawn a task because that depends on a specific Async Runtime (tokio, smol, etc).
+        // The core SDK has an AsyncRuntime abstraction that we can use to spawn the task.
+        Ok(self
+            .container_properties_cache
+            .try_get_with_by_ref(key, async { init.await.map(Arc::new) })
+            .await?)
+    }
+}