Cosmos: Add Basic Retry Policy for Throtteled Requests

sdk/cosmos/.dict.txt

@@ -5,6 +5,8 @@ euclidian
 pkranges
 sprocs
 udfs
+backoff
+pluggable
 
 # Cosmos' docs all use "Autoscale" as a single word, rather than a compound "AutoScale" or "Auto Scale"
 autoscale

sdk/cosmos/azure_data_cosmos/src/handler/mod.rs

@@ -0,0 +1,5 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+//! Handler types for request processing and retry logic.
+pub(crate) mod retry_handler;

sdk/cosmos/azure_data_cosmos/src/handler/retry_handler.rs

@@ -0,0 +1,142 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+use crate::retry_policies::resource_throttle_retry_policy::ResourceThrottleRetryPolicy;
+use crate::retry_policies::{RetryPolicy, RetryResult};
+use async_trait::async_trait;
+use azure_core::http::{request::Request, RawResponse};
+use std::sync::Arc;
+use typespec_client_core::async_runtime::get_async_runtime;
+
+// Helper trait to conditionally require Send on non-WASM targets
+#[cfg(not(target_arch = "wasm32"))]
+pub trait CosmosConditionalSend: Send {}
+#[cfg(not(target_arch = "wasm32"))]
+impl<T: Send> CosmosConditionalSend for T {}
+
+#[cfg(target_arch = "wasm32")]
+pub trait CosmosConditionalSend {}
+#[cfg(target_arch = "wasm32")]
+impl<T> CosmosConditionalSend for T {}
+
+/// Trait defining the interface for retry handlers in Cosmos DB operations
+///
+/// This trait provides a contract for implementing retry logic that wraps HTTP requests
+/// with automatic retry capabilities. Implementations can inject custom retry policies
+/// and handle both transient failures (errors) and non-success HTTP responses.
+#[allow(dead_code)]
+#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
+#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
+pub trait RetryHandler: Send + Sync {
+    /// Sends an HTTP request with automatic retry logic
+    ///
+    /// This method wraps the provided sender callback with retry logic, automatically
+    /// handling transient failures and implementing exponential backoff. The method
+    /// will continue retrying until either:
+    /// - The request succeeds (non-error 2xx status)
+    /// - The retry policy determines no more retries should be attempted
+    /// - Maximum retry attempts are exceeded
+    ///
+    /// # Arguments
+    /// * `request` - Mutable reference to the HTTP request to send (may be modified by retry policy)
+    /// * `sender` - Callback function that performs the actual HTTP request. This function
+    ///              takes a mutable request reference and returns a future that resolves to
+    ///              a `RawResponse` or error.
+    ///
+    /// # Type Parameters
+    /// * `Sender` - Function type that takes `&mut Request` and returns a Future
+    /// * `Fut` - Future type returned by the sender that resolves to `Result<RawResponse>`
+    ///
+    /// # Returns
+    /// `Result<RawResponse>` - The final response (success or failure after all retry attempts)
+    async fn send<Sender, Fut>(
+        &self,
+        request: &mut Request,
+        sender: Sender,
+    ) -> azure_core::Result<RawResponse>
+    where
+        Sender: Fn(&mut Request) -> Fut + Send + Sync,
+        Fut: std::future::Future<Output = azure_core::Result<RawResponse>> + CosmosConditionalSend;
+}
+
+/// Concrete retry handler implementation with exponential back off.
+/// This handler provides automatic retry capabilities for Cosmos DB operations using
+/// a pluggable retry policy system. It wraps HTTP requests with intelligent retry logic
+/// that handles both transient network errors and HTTP error responses.
+#[derive(Debug, Clone)]
+pub struct BackOffRetryHandler {}
+
+impl BackOffRetryHandler {
+    /// Creates a new instance of `BackOffRetryHandler`
+    ///
+    /// This constructor initializes a retry handler with exponential backoff capabilities.
+    /// The handler will dynamically select appropriate retry policies based on the request
+    /// characteristics (e.g., `ResourceThrottleRetryPolicy` for rate limiting scenarios).
+    ///
+    /// # Returns
+    ///
+    /// A new `BackOffRetryHandler` instance ready to handle request retries with
+    /// automatic policy selection and exponential backoff logic.
+    ///
+    /// # Example
+    ///
+    /// ```rust,ignore
+    /// use azure_data_cosmos::handler::retry_handler::BackOffRetryHandler;
+    ///
+    /// let retry_handler = BackOffRetryHandler::new();
+    /// // Use retry_handler.send() to make requests with automatic retry
+    /// ```
+    pub fn new() -> Self {
+        Self {}
+    }
+
+    /// Returns the appropriate retry policy based on the request
+    ///
+    /// This method examines the underlying operation and resource types and determines
+    /// retry policy should be used for this specific request.
+    /// # Arguments
+    /// * `request` - The HTTP request to analyze
+    pub fn get_policy_for_request(&self, _request: &Request) -> Arc<dyn RetryPolicy> {
+        // For now, always return ResourceThrottleRetryPolicy. Future implementation should check
+        // the request operation type and resource type and accordingly return the respective retry
+        // policy.
+        Arc::new(ResourceThrottleRetryPolicy::new(5, 200, 10))
+    }
+}
+
+#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
+#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
+impl RetryHandler for BackOffRetryHandler {
+    /// Sends an HTTP request with automatic retry and exponential back off
+    ///
+    /// This implementation of the `RetryHandler::send` method provides robust
+    /// retry logic.
+    ///
+    /// # Arguments
+    /// * `request` - Mutable HTTP request (may be modified by retry policy between attempts)
+    /// * `sender` - Callback that performs the actual HTTP request
+    async fn send<Sender, Fut>(
+        &self,
+        request: &mut Request,
+        sender: Sender,
+    ) -> azure_core::Result<RawResponse>
+    where
+        Sender: Fn(&mut Request) -> Fut + Send + Sync,
+        Fut: std::future::Future<Output = azure_core::Result<RawResponse>> + CosmosConditionalSend,
+    {
+        // Get the appropriate retry policy based on the request
+        let retry_policy = self.get_policy_for_request(request);
+        retry_policy.on_before_send_request(request);
+
+        loop {
+            // Invoke the provided sender callback instead of calling inner_send_async directly
+            let result = sender(request).await;
+            let retry_result = retry_policy.should_retry(&result).await;
+
+            match retry_result {
+                RetryResult::DoNotRetry => return result,
+                RetryResult::Retry { after } => get_async_runtime().sleep(after).await,
+            }
+        }
+    }
+}

sdk/cosmos/azure_data_cosmos/src/lib.rs

@@ -17,8 +17,6 @@ pub(crate) mod utils;
 
 pub mod models;
 
-mod location_cache;
-
 #[doc(inline)]
 pub use clients::CosmosClient;
 
@@ -28,3 +26,6 @@ pub use partition_key::*;
 pub use query::Query;
 
 pub use feed::{FeedPage, FeedPager};
+pub(crate) mod handler;
+pub mod retry_policies;
+pub mod routing;

sdk/cosmos/azure_data_cosmos/src/options/mod.rs

@@ -10,7 +10,7 @@ use std::fmt;
 use std::fmt::Display;
 
 /// Options used when creating a [`CosmosClient`](crate::CosmosClient).
-#[derive(Clone, Default)]
+#[derive(Clone, Default, Debug)]
 pub struct CosmosClientOptions {
     pub client_options: ClientOptions,
 }

sdk/cosmos/azure_data_cosmos/src/pipeline/mod.rs

@@ -17,6 +17,7 @@ use futures::TryStreamExt;
 use serde::de::DeserializeOwned;
 use url::Url;
 
+use crate::handler::retry_handler::{BackOffRetryHandler, RetryHandler};
 use crate::{
     constants,
     models::ThroughputProperties,
@@ -29,6 +30,7 @@ use crate::{
 pub struct CosmosPipeline {
     pub endpoint: Url,
     pipeline: azure_core::http::Pipeline,
+    retry_handler: BackOffRetryHandler,
 }
 
 impl CosmosPipeline {
@@ -37,16 +39,19 @@ impl CosmosPipeline {
         auth_policy: AuthorizationPolicy,
         client_options: ClientOptions,
     ) -> Self {
+        let pipeline = azure_core::http::Pipeline::new(
+            option_env!("CARGO_PKG_NAME"),
+            option_env!("CARGO_PKG_VERSION"),
+            client_options,
+            Vec::new(),
+            vec![Arc::new(auth_policy)],
+            None,
+        );
+
         CosmosPipeline {
             endpoint,
-            pipeline: azure_core::http::Pipeline::new(
-                option_env!("CARGO_PKG_NAME"),
-                option_env!("CARGO_PKG_VERSION"),
-                client_options,
-                Vec::new(),
-                vec![Arc::new(auth_policy)],
-                None,
-            ),
+            pipeline,
+            retry_handler: BackOffRetryHandler::new(),
         }
     }
 
@@ -65,9 +70,20 @@ impl CosmosPipeline {
         request: &mut Request,
         resource_link: ResourceLink,
     ) -> azure_core::Result<RawResponse> {
-        let ctx = ctx.with_value(resource_link);
-        let r = self.pipeline.send(&ctx, request, None).await?;
-        Ok(r)
+        // Clone pipeline and convert context to owned so the closure can be Fn
+        let pipeline = self.pipeline.clone();
+        let ctx_owned = ctx.with_value(resource_link).into_owned();
+
+        // Build a sender closure that forwards to the inner pipeline.send
+        let sender = move |req: &mut Request| {
+            let pipeline = pipeline.clone();
+            let ctx = ctx_owned.clone();
+            let mut req_clone = req.clone();
+            async move { pipeline.send(&ctx, &mut req_clone, None).await }
+        };
+
+        // Delegate to the retry handler, providing the sender callback
+        self.retry_handler.send(request, sender).await
     }
 
     pub async fn send<T>(

sdk/cosmos/azure_data_cosmos/src/retry_policies/mod.rs

@@ -0,0 +1,61 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+pub mod resource_throttle_retry_policy;
+use async_trait::async_trait;
+use azure_core::http::RawResponse;
+use azure_core::time::Duration;
+use typespec_client_core::http::Request;
+
+/// Result of a retry policy decision
+///
+/// This enum represents the outcome of evaluating whether an HTTP request should be retried
+/// after encountering an error or receiving a response that may warrant a retry (such as
+/// transient failures, rate limiting, or service unavailability).
+///
+/// # Variants
+///
+/// * `DoNotRetry` - The operation should not be retried. This is returned for successful
+///   responses, permanent failures, or when retry limits have been exhausted.
+///
+/// * `Retry { after }` - The operation should be retried after waiting for the specified
+///   duration. The delay allows for exponential backoff or respects server-provided retry
+///   hints (e.g., from `Retry-After` headers).
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum RetryResult {
+    /// Indicates that the operation should not be retried.
+    DoNotRetry,
+    /// Indicates that the operation should be retried after waiting for the duration specified in `after`.
+    Retry { after: Duration },
+}
+
+/// Trait defining the retry policy interface for Cosmos DB operations
+///
+/// This trait provides a contract for implementing retry logic for transient failures
+/// in Azure Cosmos DB operations. Implementers can define custom retry behavior for
+/// both exceptions (errors) and HTTP responses based on their specific requirements.
+#[async_trait]
+pub trait RetryPolicy: Send + Sync {
+    /// Called before sending a request to allow policy-specific modifications
+    ///
+    /// This method is invoked immediately before each request is sent (including retries).
+    /// # Arguments
+    /// * `request` - Mutable reference to the HTTP request being sent
+    fn on_before_send_request(&self, request: &mut Request);
+
+    /// Determines whether an HTTP request should be retried based on the response or error
+    ///
+    /// This method evaluates the result of an HTTP request attempt and decides whether
+    /// the operation should be retried, and if so, how long to wait before the next attempt.
+    ///
+    /// # Arguments
+    ///
+    /// * `response` - A reference to the result of the HTTP request attempt. This can be:
+    ///   - `Ok(RawResponse)` - A successful HTTP response (which may still indicate an error via status code)
+    ///   - `Err(azure_core::Error)` - A network or client-side error
+    ///
+    /// # Returns
+    ///
+    /// A `RetryResult` indicating the retry decision.
+    async fn should_retry(&self, response: &azure_core::Result<RawResponse>) -> RetryResult;
+}