Enhance Cosmos KV to support multi-region reads

Signed-off-by: Kate Goldenring <[email protected]>

Author: kate-goldenring

Cargo.lock

@@ -10,9 +10,9 @@ rust-version.workspace = true
 
 [dependencies]
 anyhow = { workspace = true }
-azure_data_cosmos = { git = "https://github.com/azure/azure-sdk-for-rust.git", rev = "8c4caa251c3903d5eae848b41bb1d02a4d65231c" }
-azure_identity = { git = "https://github.com/azure/azure-sdk-for-rust.git", rev = "8c4caa251c3903d5eae848b41bb1d02a4d65231c" }
-azure_core = { git = "https://github.com/azure/azure-sdk-for-rust.git", rev = "8c4caa251c3903d5eae848b41bb1d02a4d65231c" }
+azure_data_cosmos = { git = "https://github.com/azure/azure-sdk-for-rust.git", tag = "azure_data_cosmos-0.21.0"  }
+azure_identity = { git = "https://github.com/azure/azure-sdk-for-rust.git", tag = "azure_data_cosmos-0.21.0"  }
+azure_core = { git = "https://github.com/azure/azure-sdk-for-rust.git", tag = "azure_data_cosmos-0.21.0"  }
 futures = { workspace = true }
 serde = { workspace = true }
 spin-core = { path = "../core" }

crates/key-value-azure/Cargo.toml

@@ -2,6 +2,7 @@ use anyhow::{Context, Result};
 use azure_data_cosmos::{
     prelude::{
         AuthorizationToken, CollectionClient, CosmosClient, CosmosClientBuilder, Operation, Query,
+        TentativeWritesAllowance,
     },
     CosmosEntity,
 };
@@ -145,7 +146,7 @@ struct AzureCosmosStore {
 #[async_trait]
 impl Store for AzureCosmosStore {
     async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, Error> {
-        let pair = self.get_pair(key).await?;
+        let pair = self.get_pair::<Pair>(key).await?;
         Ok(pair.map(|p| p.value))
     }
 
@@ -158,24 +159,35 @@ impl Store for AzureCosmosStore {
         self.client
             .create_document(pair)
             .is_upsert(true)
+            .allow_tentative_writes(TentativeWritesAllowance::Allow)
             .await
             .map_err(log_error)?;
         Ok(())
     }
 
     async fn delete(&self, key: &str) -> Result<(), Error> {
-        if self.exists(key).await? {
-            let document_client = self
-                .client
-                .document_client(key, &self.store_id)
-                .map_err(log_error)?;
-            document_client.delete_document().await.map_err(log_error)?;
+        let document_client = self
+            .client
+            .document_client(key, &self.store_id.clone().unwrap_or(key.to_string()))
+            .map_err(log_error)?;
+        if let Err(e) = document_client
+            .delete_document()
+            .allow_tentative_writes(TentativeWritesAllowance::Allow)
+            .await
+        {
+            if e.as_http_error()
+                .map(|e| e.status())
+                .unwrap_or(azure_core::StatusCode::Continue)
+                != 404
+            {
+                return Err(log_error(e));
+            }
         }
         Ok(())
     }
 
     async fn exists(&self, key: &str) -> Result<bool, Error> {
-        Ok(self.get_pair(key).await?.is_some())
+        Ok(self.get_pair::<Key>(key).await?.is_some())
     }
 
     async fn get_keys(&self) -> Result<Vec<String>, Error> {
@@ -216,24 +228,67 @@ impl Store for AzureCosmosStore {
         Ok(())
     }
 
+    // WARNING: this function only works on the primary region because the
+    // `azure_data_cosmos-0.21.0` release of the Azure cosmos SDK does not
+    // support setting allow_tentative_writes on patch requests. The initial
+    // value for the item must be set through this interfaces, as this sets the
+    // number value if it does not exist. If the value was previously set using
+    // the `set` interface, this will fail due to a type mismatch.
+    //
+    // TODO: The function should parse the new value from the return response
+    // rather than sending an additional new request. However, the current SDK
+    // version does not support this.
     async fn increment(&self, key: String, delta: i64) -> Result<i64, Error> {
         let operations = vec![Operation::incr("/value", delta).map_err(log_error)?];
-        let _ = self
+        match self
             .client
-            .document_client(key.clone(), &self.store_id)
+            .document_client(&key, &self.store_id.clone().unwrap_or(key.to_string()))
             .map_err(log_error)?
             .patch_document(operations)
             .await
-            .map_err(log_error)?;
-        let pair = self.get_pair(key.as_ref()).await?;
-        match pair {
-            Some(p) => Ok(i64::from_le_bytes(
-                p.value.try_into().expect("incorrect length"),
-            )),
-            None => Err(Error::Other(
-                "increment returned an empty value after patching, which indicates a bug"
-                    .to_string(),
-            )),
+        {
+            Err(e) => {
+                if e.as_http_error()
+                    .map(|e| e.status())
+                    .unwrap_or(azure_core::StatusCode::Continue)
+                    == 404
+                {
+                    let counter = Counter {
+                        id: key.clone(),
+                        value: delta,
+                        store_id: self.store_id.clone(),
+                    };
+                    if let Err(e) = self
+                        .client
+                        .create_document(counter)
+                        .is_upsert(false)
+                        .allow_tentative_writes(TentativeWritesAllowance::Allow)
+                        .await
+                    {
+                        if e.as_http_error()
+                            .map(|e| e.status())
+                            .unwrap_or(azure_core::StatusCode::Continue)
+                            == 409
+                        {
+                            // Conflict trying to create counter, retry increment
+                            self.increment(key, delta).await?;
+                        } else {
+                            return Err(log_error(e));
+                        }
+                    }
+                    Ok(delta)
+                } else {
+                    Err(log_error(e))
+                }
+            }
+            Ok(_) => self
+                .get_pair::<Counter>(key.as_ref())
+                .await?
+                .map(|c| c.value)
+                .ok_or(Error::Other(
+                    "increment returned an empty value after patching, which indicates a bug"
+                        .to_string(),
+                )),
         }
     }
 
@@ -274,8 +329,9 @@ impl CompareAndSwap {
 
 #[async_trait]
 impl Cas for CompareAndSwap {
-    /// `current` will fetch the current value for the key and store the etag for the record. The
-    /// etag will be used to perform and optimistic concurrency update using the `if-match` header.
+    /// `current` will fetch the current value for the key and store the etag
+    /// for the record. The etag will be used to perform and optimistic
+    /// concurrency update using the `if-match` header.
     async fn current(&self) -> Result<Option<Vec<u8>>, Error> {
         let mut stream = self
             .client
@@ -307,8 +363,8 @@ impl Cas for CompareAndSwap {
         }
     }
 
-    /// `swap` updates the value for the key using the etag saved in the `current` function for
-    /// optimistic concurrency.
+    /// `swap` updates the value for the key using the etag saved in the
+    /// `current` function for optimistic concurrency.
     async fn swap(&self, value: Vec<u8>) -> Result<(), SwapError> {
         let pair = Pair {
             id: self.key.clone(),
@@ -327,15 +383,18 @@ impl Cas for CompareAndSwap {
                 // attempt to replace the document if the etag matches
                 doc_client
                     .replace_document(pair)
+                    .allow_tentative_writes(TentativeWritesAllowance::Allow)
                     .if_match_condition(azure_core::request_options::IfMatchCondition::Match(etag))
                     .await
                     .map_err(|e| SwapError::CasFailed(format!("{e:?}")))
                     .map(drop)
             }
             None => {
-                // if we have no etag, then we assume the document does not yet exist and must insert; no upserts.
+                // if we have no etag, then we assume the document does not yet
+                // exist and must insert; no upserts.
                 self.client
                     .create_document(pair)
+                    .allow_tentative_writes(TentativeWritesAllowance::Allow)
                     .await
                     .map_err(|e| SwapError::CasFailed(format!("{e:?}")))
                     .map(drop)
@@ -353,15 +412,18 @@ impl Cas for CompareAndSwap {
 }
 
 impl AzureCosmosStore {
-    async fn get_pair(&self, key: &str) -> Result<Option<Pair>, Error> {
+    async fn get_pair<F>(&self, key: &str) -> Result<Option<F>, Error>
+    where
+        F: CosmosEntity + Send + Sync + serde::de::DeserializeOwned + Clone,
+    {
         let query = self
             .client
             .query_documents(Query::new(self.get_query(key)))
             .query_cross_partition(true)
             .max_item_count(1);
 
         // There can be no duplicated keys, so we create the stream and only take the first result.
-        let mut stream = query.into_stream::<Pair>();
+        let mut stream = query.into_stream::<F>();
         let Some(res) = stream.next().await else {
             return Ok(None);
         };
@@ -379,10 +441,10 @@ impl AzureCosmosStore {
             .query_cross_partition(true);
         let mut res = Vec::new();
 
-        let mut stream = query.into_stream::<Pair>();
+        let mut stream = query.into_stream::<Key>();
         while let Some(resp) = stream.next().await {
             let resp = resp.map_err(log_error)?;
-            res.extend(resp.results.into_iter().map(|(pair, _)| pair.id));
+            res.extend(resp.results.into_iter().map(|(key, _)| key.id));
         }
 
         Ok(res)
@@ -435,6 +497,7 @@ fn append_store_id_condition(
     }
 }
 
+// Pair structure for key value operations
 #[derive(Serialize, Deserialize, Clone, Debug)]
 pub struct Pair {
     pub id: String,
@@ -450,3 +513,36 @@ impl CosmosEntity for Pair {
         self.store_id.clone().unwrap_or_else(|| self.id.clone())
     }
 }
+
+// Counter structure for increment operations
+#[derive(Serialize, Deserialize, Clone, Debug)]
+pub struct Counter {
+    pub id: String,
+    pub value: i64,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub store_id: Option<String>,
+}
+
+impl CosmosEntity for Counter {
+    type Entity = String;
+
+    fn partition_key(&self) -> Self::Entity {
+        self.store_id.clone().unwrap_or_else(|| self.id.clone())
+    }
+}
+
+// Key structure for operations with generic value types
+#[derive(Serialize, Deserialize, Clone, Debug)]
+pub struct Key {
+    pub id: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub store_id: Option<String>,
+}
+
+impl CosmosEntity for Key {
+    type Entity = String;
+
+    fn partition_key(&self) -> Self::Entity {
+        self.store_id.clone().unwrap_or_else(|| self.id.clone())
+    }
+}