fix: more optimal handling of resource paths

Author: analogrelay

sdk/cosmos/azure_data_cosmos/src/driver_bridge.rs

@@ -24,18 +24,20 @@ use crate::{
 /// Converts a driver [`DriverResponse`] into the SDK's typed [`CosmosResponse<T>`].
 ///
 /// This reconstructs an `azure_core::Response<T>` from the driver's raw bytes,
-/// status code, and headers, then wraps it in the SDK's response type.
+/// status code, and headers, then wraps it in the SDK's response type using
+/// the pre-parsed headers from the driver to avoid a redundant parse.
 pub(crate) fn driver_response_to_cosmos_response<T>(
     driver_response: DriverResponse,
 ) -> CosmosResponse<T> {
     let status_code: StatusCode = driver_response.status().status_code();
-    let headers = driver_response_headers_to_headers(driver_response.headers());
+    let cosmos_headers = driver_response.headers().clone();
+    let headers = driver_response_headers_to_headers(&cosmos_headers);
     let body = driver_response.into_body();
 
     let raw_response = RawResponse::from_bytes(status_code, headers, Bytes::from(body));
     let typed_response: Response<T> = raw_response.into();
 
-    CosmosResponse::from_response(typed_response)
+    CosmosResponse::from_driver_response(typed_response, cosmos_headers)
 }
 
 /// Converts driver [`CosmosResponseHeaders`] into raw [`Headers`] for the SDK response.

sdk/cosmos/azure_data_cosmos/src/models/cosmos_response.rs

@@ -50,11 +50,14 @@ impl<T> CosmosResponse<T> {
         }
     }
 
-    /// Creates a `CosmosResponse` from a typed response without a request.
+    /// Creates a `CosmosResponse` from a typed response and pre-parsed headers.
     ///
-    /// Used for driver-routed operations where no `CosmosRequest` is available.
-    pub(crate) fn from_response(response: Response<T>) -> Self {
-        let cosmos_headers = CosmosResponseHeaders::from_headers(response.headers());
+    /// Used by the driver bridge to avoid re-parsing headers that were already
+    /// parsed by the driver pipeline.
+    pub(crate) fn from_driver_response(
+        response: Response<T>,
+        cosmos_headers: CosmosResponseHeaders,
+    ) -> Self {
         let diagnostics = CosmosDiagnostics::from_headers(&cosmos_headers);
         Self {
             response,

sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/components.rs

@@ -19,7 +19,7 @@ use crate::{
         routing::{CosmosEndpoint, LocationIndex},
         transport::AuthorizationContext,
     },
-    models::CosmosStatus,
+    models::{CosmosResponseHeaders, CosmosStatus},
     options::Region,
 };
 
@@ -296,12 +296,17 @@ impl TransportResult {
     ///
     /// Successful status codes are mapped to `Success`; non-success status codes
     /// are mapped to `HttpError` with `request_sent` set to `Sent`.
-    pub fn from_http_response(status: CosmosStatus, headers: Headers, body: Vec<u8>) -> Self {
+    pub fn from_http_response(
+        status: CosmosStatus,
+        headers: Headers,
+        cosmos_headers: CosmosResponseHeaders,
+        body: Vec<u8>,
+    ) -> Self {
         if status.is_success() {
             Self {
                 outcome: TransportOutcome::Success {
                     status,
-                    headers,
+                    cosmos_headers,
                     body,
                 },
             }
@@ -310,37 +315,53 @@ impl TransportResult {
                 outcome: TransportOutcome::HttpError {
                     status,
                     headers,
+                    cosmos_headers,
                     body,
                     request_sent: RequestSentStatus::Sent,
                 },
             }
         }
     }
 
-    /// Returns the response headers if this is an HTTP response.
-    pub fn response_headers(&self) -> Option<&Headers> {
+    /// Returns the parsed Cosmos response headers if this is an HTTP response.
+    pub fn cosmos_headers(&self) -> Option<&CosmosResponseHeaders> {
         match &self.outcome {
-            TransportOutcome::Success { headers, .. } => Some(headers),
-            TransportOutcome::HttpError { headers, .. } => Some(headers),
+            TransportOutcome::Success { cosmos_headers, .. } => Some(cosmos_headers),
+            TransportOutcome::HttpError { cosmos_headers, .. } => Some(cosmos_headers),
             TransportOutcome::TransportError { .. } | TransportOutcome::DeadlineExceeded { .. } => {
                 None
             }
         }
     }
+
+    /// Returns the raw response headers for HTTP error responses.
+    ///
+    /// Raw headers are only retained for error responses (needed to build a `RawResponse`
+    /// for callers). For success responses, only parsed `CosmosResponseHeaders` are kept.
+    pub fn response_headers(&self) -> Option<&Headers> {
+        match &self.outcome {
+            TransportOutcome::HttpError { headers, .. } => Some(headers),
+            _ => None,
+        }
+    }
 }
 
 /// The outcome of a single transport attempt.
 pub(crate) enum TransportOutcome {
     /// Successful response (2xx).
     Success {
         status: CosmosStatus,
-        headers: Headers,
+        /// Parsed Cosmos-specific response headers.
+        cosmos_headers: CosmosResponseHeaders,
         body: Vec<u8>,
     },
     /// HTTP error response (non-2xx) that may be retryable at the operation level.
     HttpError {
         status: CosmosStatus,
+        /// Raw headers retained for building `RawResponse` in error reporting.
         headers: Headers,
+        /// Parsed Cosmos-specific response headers.
+        cosmos_headers: CosmosResponseHeaders,
         body: Vec<u8>,
         request_sent: RequestSentStatus,
     },
@@ -371,12 +392,9 @@ impl std::fmt::Display for TransportOutcome {
 impl std::fmt::Debug for TransportOutcome {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            TransportOutcome::Success {
-                status, headers, ..
-            } => f
+            TransportOutcome::Success { status, .. } => f
                 .debug_struct("Success")
                 .field("status", status)
-                .field("headers", headers)
                 .field("body", &"...")
                 .finish(),
             TransportOutcome::HttpError {

sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/operation_pipeline.rs

@@ -19,8 +19,8 @@ use crate::{
         LocationStateStore,
     },
     models::{
-        header_names, AccountEndpoint, ActivityId, CosmosOperation, CosmosResponse,
-        CosmosResponseHeaders, Credential, DefaultConsistencyLevel, SessionToken, SubStatusCode,
+        header_names, AccountEndpoint, ActivityId, CosmosOperation, CosmosResponse, Credential,
+        DefaultConsistencyLevel, SessionToken, SubStatusCode,
     },
     options::{OperationOptionsView, ReadConsistencyStrategy},
 };
@@ -215,13 +215,12 @@ pub(crate) async fn execute_operation_pipeline(
         // variant does not carry headers — capturing after evaluation
         // would silently drop tokens from those responses.
         if session_consistency_active {
-            if let Some(headers) = result.response_headers() {
-                let cosmos_headers = CosmosResponseHeaders::from_headers(headers);
+            if let Some(cosmos_headers) = result.cosmos_headers() {
                 if should_capture_session_token_from_status(
                     cosmos_headers.substatus.as_ref(),
                     &result.outcome,
                 ) {
-                    session_manager.capture_session_token(operation, &cosmos_headers);
+                    session_manager.capture_session_token(operation, cosmos_headers);
                 }
             }
         }
@@ -435,9 +434,11 @@ fn build_transport_request(
     resolved_session_token: Option<SessionToken>,
 ) -> azure_core::Result<TransportRequest> {
     let resource_ref = operation.resource_reference();
-    let request_path = resource_ref.request_path();
+    // Compute both paths in a single pass with a single allocation.
+    let paths = resource_ref.compute_paths();
     let url = {
         let mut base = routing.selected_url.clone();
+        let request_path = paths.request_path();
         let normalized = if request_path.starts_with('/') {
             request_path.to_string()
         } else if request_path.is_empty() {
@@ -451,10 +452,9 @@ fn build_transport_request(
 
     let method = operation.operation_type().http_method();
     let resource_type = operation.resource_type();
-    let resource_link = resource_ref.link_for_signing();
-    let signing_link = resource_link.trim_start_matches('/');
-
-    let auth_context = AuthorizationContext::new(method, resource_type, signing_link);
+    // Move `paths` into AuthorizationContext so the signing link is a zero-copy
+    // sub-slice of the path buffer — no additional string allocation needed.
+    let auth_context = AuthorizationContext::from_paths(method, resource_type, paths);
 
     // Build headers from the operation.
     // Custom headers are inserted first so that SDK-set headers below always
@@ -526,10 +526,9 @@ fn build_cosmos_response(
     match result.outcome {
         TransportOutcome::Success {
             status,
-            headers,
+            cosmos_headers,
             body,
         } => {
-            let cosmos_headers = CosmosResponseHeaders::from_headers(&headers);
             diagnostics.set_operation_status(status.status_code(), status.sub_status());
 
             let diagnostics_ctx = Arc::new(diagnostics.complete());
@@ -1002,15 +1001,15 @@ mod tests {
 
         use crate::{
             driver::pipeline::components::TransportOutcome,
-            models::{CosmosStatus, SubStatusCode},
+            models::{CosmosResponseHeaders, CosmosStatus, SubStatusCode},
         };
 
         use super::super::should_capture_session_token_from_status;
 
         fn success_outcome() -> TransportOutcome {
             TransportOutcome::Success {
                 status: CosmosStatus::new(StatusCode::Ok),
-                headers: Headers::new(),
+                cosmos_headers: CosmosResponseHeaders::default(),
                 body: Vec::new(),
             }
         }
@@ -1019,6 +1018,7 @@ mod tests {
             TransportOutcome::HttpError {
                 status: CosmosStatus::new(status),
                 headers: Headers::new(),
+                cosmos_headers: CosmosResponseHeaders::default(),
                 body: Vec::new(),
                 request_sent: crate::diagnostics::RequestSentStatus::Sent,
             }

sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/retry_evaluation.rs

@@ -46,6 +46,7 @@ pub(crate) fn evaluate_transport_result(
             headers,
             body,
             request_sent,
+            ..
         } => {
             let request_definitely_not_sent = request_sent.definitely_not_sent();
 
@@ -284,7 +285,10 @@ mod tests {
     use super::*;
     use crate::{
         diagnostics::RequestSentStatus,
-        models::{AccountReference, CosmosOperation, CosmosStatus, DatabaseReference},
+        models::{
+            AccountReference, CosmosOperation, CosmosResponseHeaders, CosmosStatus,
+            DatabaseReference,
+        },
     };
     use azure_core::http::StatusCode;
 
@@ -310,7 +314,7 @@ mod tests {
         TransportResult {
             outcome: TransportOutcome::Success {
                 status: CosmosStatus::new(StatusCode::Ok),
-                headers: azure_core::http::headers::Headers::new(),
+                cosmos_headers: CosmosResponseHeaders::default(),
                 body: b"{}".to_vec(),
             },
         }
@@ -334,6 +338,7 @@ mod tests {
             outcome: TransportOutcome::HttpError {
                 status: CosmosStatus::new(status_code),
                 headers: azure_core::http::headers::Headers::new(),
+                cosmos_headers: CosmosResponseHeaders::default(),
                 body: vec![],
                 request_sent: RequestSentStatus::Sent,
             },
@@ -482,6 +487,7 @@ mod tests {
             outcome: TransportOutcome::HttpError {
                 status: CosmosStatus::WRITE_FORBIDDEN,
                 headers: azure_core::http::headers::Headers::new(),
+                cosmos_headers: CosmosResponseHeaders::default(),
                 body: vec![],
                 request_sent: RequestSentStatus::Sent,
             },
@@ -505,6 +511,7 @@ mod tests {
             outcome: TransportOutcome::HttpError {
                 status: CosmosStatus::READ_SESSION_NOT_AVAILABLE,
                 headers: azure_core::http::headers::Headers::new(),
+                cosmos_headers: CosmosResponseHeaders::default(),
                 body: vec![],
                 request_sent: RequestSentStatus::Sent,
             },

sdk/cosmos/azure_data_cosmos_driver/src/driver/transport/authorization_policy.rs

@@ -10,22 +10,55 @@ use crate::models::{Credential, ResourceType};
 use azure_core::http::Method;
 use tracing::trace;
 
+use crate::models::ResourcePaths;
+
 /// Cosmos DB AAD scope for token authentication.
 const COSMOS_AAD_SCOPE: &str = "https://cosmos.azure.com/.default";
 
+/// The resource link used when signing a Cosmos DB request.
+///
+/// `Paths` owns a [`ResourcePaths`] so the signing link is derived as a
+/// zero-copy sub-slice of the pre-computed path buffer (the hot path).
+/// `Owned` holds an independently allocated `String` for call sites that
+/// construct an `AuthorizationContext` outside of the normal request pipeline.
+pub(crate) enum ResourceLink {
+    /// Signing link is derived from the pre-computed [`ResourcePaths`] buffer.
+    Paths(ResourcePaths),
+    /// Signing link is an independently owned string.
+    Owned(String),
+}
+
+impl ResourceLink {
+    pub(crate) fn as_str(&self) -> &str {
+        match self {
+            Self::Paths(p) => p.signing_link(),
+            Self::Owned(s) => s.as_str(),
+        }
+    }
+}
+
+impl std::fmt::Debug for ResourceLink {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(self.as_str())
+    }
+}
+
 /// Authorization context needed to build a Cosmos DB signature.
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub(crate) struct AuthorizationContext {
     /// The HTTP method of the request.
     pub(crate) method: Method,
     /// The resource type being accessed.
     pub(crate) resource_type: ResourceType,
     /// The resource link for signing (path without leading slash, unencoded).
-    pub(crate) resource_link: String,
+    pub(crate) resource_link: ResourceLink,
 }
 
 impl AuthorizationContext {
-    /// Creates a new authorization context.
+    /// Creates a new authorization context with an owned resource link string.
+    ///
+    /// Use [`AuthorizationContext::from_paths`] on the hot path to avoid copying
+    /// the signing link out of the pre-computed [`ResourcePaths`].
     pub(crate) fn new(
         method: Method,
         resource_type: ResourceType,
@@ -34,7 +67,21 @@ impl AuthorizationContext {
         Self {
             method,
             resource_type,
-            resource_link: resource_link.into(),
+            resource_link: ResourceLink::Owned(resource_link.into()),
+        }
+    }
+
+    /// Creates a new authorization context that derives the signing link directly
+    /// from `paths`, avoiding any additional string allocation.
+    pub(crate) fn from_paths(
+        method: Method,
+        resource_type: ResourceType,
+        paths: ResourcePaths,
+    ) -> Self {
+        Self {
+            method,
+            resource_type,
+            resource_link: ResourceLink::Paths(paths),
         }
     }
 }
@@ -82,7 +129,7 @@ fn build_string_to_sign(auth_ctx: &AuthorizationContext, date_string: &str) -> S
         "{}\n{}\n{}\n{}\n\n",
         method_str,
         auth_ctx.resource_type.path_segment(),
-        auth_ctx.resource_link,
+        auth_ctx.resource_link.as_str(),
         date_string,
     )
 }