From 72ce3f9fa1a025227f058ee3321ec316e5ad3851 Mon Sep 17 00:00:00 2001 From: vincenttran-msft <101599632+vincenttran-msft@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:19:52 -0700 Subject: [PATCH 1/6] data locality ported, WIP --- .../src/clients/blob_client.rs | 437 +++++++++- .../generated/clients/append_blob_client.rs | 22 +- .../src/generated/clients/blob_client.rs | 44 +- .../clients/blob_container_client.rs | 48 +- .../generated/clients/blob_service_client.rs | 26 +- .../generated/clients/block_blob_client.rs | 26 +- .../src/generated/clients/page_blob_client.rs | 10 +- .../src/models/method_options.rs | 32 + .../azure_storage_blob/src/models/mod.rs | 2 +- .../src/partitioned_transfer/download.rs | 13 + .../src/partitioned_transfer/layout.rs | 825 ++++++++++++++++++ .../src/partitioned_transfer/mod.rs | 2 + .../azure_storage_blob/tests/blob_client.rs | 411 ++++++++- 13 files changed, 1788 insertions(+), 110 deletions(-) create mode 100644 sdk/storage/azure_storage_blob/src/partitioned_transfer/layout.rs diff --git a/sdk/storage/azure_storage_blob/src/clients/blob_client.rs b/sdk/storage/azure_storage_blob/src/clients/blob_client.rs index 5bf60b8db1e..a8cd54b4239 100644 --- a/sdk/storage/azure_storage_blob/src/clients/blob_client.rs +++ b/sdk/storage/azure_storage_blob/src/clients/blob_client.rs @@ -5,27 +5,36 @@ pub use crate::generated::clients::{BlobClient, BlobClientOptions}; use crate::{ generated::{ - clients::BlobClient as GeneratedBlobClient, models::BlobClientDownloadInternalOptions, + clients::BlobClient as GeneratedBlobClient, + models::{BlobClientDownloadInternalOptions, BlobClientListLayoutOptions}, }, models::{ BlobClientDownloadIntoResult, BlobClientDownloadOptions, BlobClientDownloadResult, BlobClientUploadOptions, BlobClientUploadResult, BlobDownloadProperties, HttpRange, - StorageErrorCode, + LayoutAwareRouting, StorageErrorCode, }, - partitioned_transfer::{self, PartitionedDownloadBehavior}, + partitioned_transfer::{self, Layout, LayoutEndpoint, PartitionedDownloadBehavior}, AppendBlobClient, BlockBlobClient, PageBlobClient, }; use async_trait::async_trait; use azure_core::{ + async_runtime::get_async_runtime, credentials::TokenCredential, error::ErrorKind, http::{ + headers::Headers, policies::{auth::BearerTokenAuthorizationPolicy, Policy}, - AsyncRawResponse, Etag, NoFormat, Pipeline, RequestContent, StatusCode, Url, UrlExt, + AsyncRawResponse, Context, Etag, NoFormat, Pipeline, RequestContent, StatusCode, Url, + UrlExt, }, tracing, Bytes, Result, }; -use std::{ops::Range, sync::Arc}; +use futures::lock::Mutex; +use std::{ + ops::Range, + sync::{Arc, OnceLock}, + time::{Duration, Instant}, +}; impl BlobClient { /// Creates a new BlobClient from a blob URL. @@ -71,7 +80,7 @@ impl BlobClient { option_env!("CARGO_PKG_NAME"), option_env!("CARGO_PKG_VERSION"), options.client_options.clone(), - Vec::default(), + vec![Arc::new(partitioned_transfer::LayoutRoutingPolicy)], per_retry_policies, None, ); @@ -194,14 +203,16 @@ impl BlobClient { let partition_size = options .partition_size .unwrap_or(crate::partitioned_transfer::defaults::DEFAULT_DOWNLOAD_PARTITION_SIZE); - let range = options.range.clone(); let inner_client = GeneratedBlobClient { endpoint: self.endpoint.clone(), pipeline: self.pipeline.clone(), version: self.version.clone(), tracer: self.tracer.clone(), }; - let behavior = BlobClientDownloadBehavior::new(inner_client, options.into()); + let layout_request = layout_request_from_options(&options); + let range = options.range.clone(); + let behavior = + BlobClientDownloadBehavior::new(inner_client, options.into(), layout_request); let response = partitioned_transfer::download(range, parallel, partition_size, Arc::new(behavior)) .await?; @@ -238,14 +249,16 @@ impl BlobClient { let partition_size = options .partition_size .unwrap_or(crate::partitioned_transfer::defaults::DEFAULT_DOWNLOAD_PARTITION_SIZE); - let range = options.range.clone(); let inner_client = GeneratedBlobClient { endpoint: self.endpoint.clone(), pipeline: self.pipeline.clone(), version: self.version.clone(), tracer: self.tracer.clone(), }; - let behavior = BlobClientDownloadBehavior::new(inner_client, options.into()); + let layout_request = layout_request_from_options(&options); + let range = options.range.clone(); + let behavior = + BlobClientDownloadBehavior::new(inner_client, options.into(), layout_request); let (_, headers, len) = partitioned_transfer::download_into( buffer, range, @@ -304,11 +317,22 @@ impl BlobClient { struct BlobClientDownloadBehavior<'a> { client: GeneratedBlobClient, options: BlobClientDownloadInternalOptions<'a>, + layout_request: Option>, + cache: OnceLock>, } impl<'a> BlobClientDownloadBehavior<'a> { - fn new(client: GeneratedBlobClient, options: BlobClientDownloadInternalOptions<'a>) -> Self { - Self { client, options } + fn new( + client: GeneratedBlobClient, + options: BlobClientDownloadInternalOptions<'a>, + layout_request: Option>, + ) -> Self { + Self { + client, + options, + layout_request, + cache: OnceLock::new(), + } } } @@ -320,6 +344,19 @@ impl PartitionedDownloadBehavior for BlobClientDownloadBehavior<'_> { etag_lock: Option, ) -> Result { let mut opt = self.options.clone(); + if let Some(Some(cache)) = self.cache.get() { + if let Some(range) = &range { + if let Some(layout) = cache.current(&self.client).await { + if let Some(endpoint) = layout.ideal_endpoint(range.start as i64) { + opt.method_options.context = opt + .method_options + .context + .clone() + .with_value(LayoutEndpoint(endpoint.to_owned())); + } + } + } + } opt.range = range.map(HttpRange::from); if let Some(etag) = etag_lock { opt.if_match = Some(etag); @@ -333,4 +370,380 @@ impl PartitionedDownloadBehavior for BlobClientDownloadBehavior<'_> { .await .map(AsyncRawResponse::from) } + + async fn prepare(&self, initial_headers: &Headers, etag_lock: Option<&Etag>) -> Result<()> { + let Some(layout_request) = &self.layout_request else { + return Ok(()); + }; + if initial_headers.get_optional_str(&"x-ms-download-hint".into()) != Some("layout") { + let _ = self.cache.set(None); + return Ok(()); + } + let mut request = layout_request.clone(); + if request.if_match.is_none() { + request.if_match = etag_lock.cloned(); + } + let context = self.options.method_options.context.clone(); + let cache = partitioned_transfer::fetch_layout(&self.client, &context, &request) + .await? + .map(|prefetch| LayoutCache::new(request, Arc::new(prefetch.layout))); + let _ = self.cache.set(cache); + Ok(()) + } +} + +const LAYOUT_TTL: Duration = Duration::from_secs(300); +const LAYOUT_REFRESH_BUFFER: Duration = Duration::from_secs(30); +const LAYOUT_REFRESH_BACKOFF: Duration = Duration::from_secs(30); + +struct LayoutCache { + request: BlobClientListLayoutOptions<'static>, + state: Arc>, +} + +struct CachedLayout { + layout: Arc, + refresh_at: Instant, + expires_at: Instant, + refreshing: bool, + retry_at: Option, +} + +impl CachedLayout { + fn new(layout: Arc) -> Self { + let now = Instant::now(); + Self { + layout, + refresh_at: now + (LAYOUT_TTL - LAYOUT_REFRESH_BUFFER), + expires_at: now + LAYOUT_TTL, + refreshing: false, + retry_at: None, + } + } +} + +impl LayoutCache { + fn new(request: BlobClientListLayoutOptions<'static>, layout: Arc) -> Self { + Self { + request, + state: Arc::new(Mutex::new(CachedLayout::new(layout))), + } + } + + async fn current(&self, client: &GeneratedBlobClient) -> Option> { + let mut state = self.state.lock().await; + let now = Instant::now(); + let backing_off = matches!(state.retry_at, Some(at) if now < at); + if now >= state.refresh_at && !state.refreshing && !backing_off { + state.refreshing = true; + let _refresh = get_async_runtime().spawn(Box::pin(Self::refresh( + clone_blob_client(client), + self.request.clone(), + self.state.clone(), + ))); + } + (now < state.expires_at).then(|| state.layout.clone()) + } + + async fn refresh( + client: GeneratedBlobClient, + request: BlobClientListLayoutOptions<'static>, + state: Arc>, + ) { + let result = partitioned_transfer::fetch_layout(&client, &Context::new(), &request).await; + let mut state = state.lock().await; + match result { + Ok(Some(prefetch)) => { + *state = CachedLayout::new(Arc::new(prefetch.layout)); + } + Ok(None) | Err(_) => { + state.refreshing = false; + state.retry_at = Some(Instant::now() + LAYOUT_REFRESH_BACKOFF); + } + } + } +} + +fn clone_blob_client(client: &GeneratedBlobClient) -> GeneratedBlobClient { + GeneratedBlobClient { + endpoint: client.endpoint.clone(), + pipeline: client.pipeline.clone(), + version: client.version.clone(), + tracer: client.tracer.clone(), + } +} + +fn layout_request_from_options( + options: &BlobClientDownloadOptions<'_>, +) -> Option> { + matches!(options.layout_aware_routing, LayoutAwareRouting::Enabled) + .then(|| layout_options_from_download(options)) +} + +fn layout_options_from_download( + options: &BlobClientDownloadOptions<'_>, +) -> BlobClientListLayoutOptions<'static> { + BlobClientListLayoutOptions { + encryption_algorithm: options.encryption_algorithm, + encryption_key: options.encryption_key.clone(), + encryption_key_sha256: options.encryption_key_sha256.clone(), + if_match: options.if_match.clone(), + if_modified_since: options.if_modified_since, + if_none_match: options.if_none_match.clone(), + if_tags: options.if_tags.clone(), + if_unmodified_since: options.if_unmodified_since, + lease_id: options.lease_id.clone(), + range: options.range.clone(), + snapshot: options.snapshot.clone(), + timeout: options.timeout, + version_id: options.version_id.clone(), + ..Default::default() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use azure_core::http::{ClientOptions, FixedRetryOptions, HttpClient, RetryOptions, Transport}; + use azure_core_test::http::MockHttpClient; + use futures::FutureExt as _; + use std::sync::atomic::{AtomicUsize, Ordering}; + + const LAYOUT_V1: &[u8] = br#" + + + + + + + +"#; + + const LAYOUT_V2: &[u8] = br#" + + + + + + + +"#; + + fn mock_client(transport: Arc) -> GeneratedBlobClient { + BlobClient::new( + "https://acct.blob.core.windows.net/container/blob" + .parse() + .unwrap(), + None, + Some(BlobClientOptions { + client_options: ClientOptions { + transport: Some(Transport::new(transport)), + retry: RetryOptions::fixed(FixedRetryOptions { + max_retries: 0, + ..Default::default() + }), + ..Default::default() + }, + ..Default::default() + }), + ) + .unwrap() + } + + fn etag_headers() -> Headers { + let mut headers = Headers::new(); + headers.insert("etag", "etag-1".to_owned()); + headers + } + + async fn seed_layout( + client: &GeneratedBlobClient, + request: &BlobClientListLayoutOptions<'static>, + ) -> Layout { + partitioned_transfer::fetch_layout(client, &Context::new(), request) + .await + .unwrap() + .unwrap() + .layout + } + + #[tokio::test] + async fn refresh_success_updates_layout_and_resets_deadlines() { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let mock: Arc = Arc::new(MockHttpClient::new(move |_req| { + let count = counter.fetch_add(1, Ordering::SeqCst); + async move { + let body = if count == 0 { LAYOUT_V1 } else { LAYOUT_V2 }; + Ok(AsyncRawResponse::from_bytes( + StatusCode::Ok, + etag_headers(), + Bytes::from_static(body), + )) + } + .boxed() + })); + let client = mock_client(mock); + let request = BlobClientListLayoutOptions::default(); + let state = Arc::new(Mutex::new(CachedLayout { + layout: Arc::new(seed_layout(&client, &request).await), + refresh_at: Instant::now(), + expires_at: Instant::now(), + refreshing: true, + retry_at: None, + })); + + LayoutCache::refresh(client, request, state.clone()).await; + let state = state.lock().await; + assert_eq!( + state.layout.ideal_endpoint(0), + Some("epv2.blob.storage.azure.net:443") + ); + assert!(!state.refreshing); + assert!(state.retry_at.is_none()); + assert!(state.expires_at > Instant::now()); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn current_routes_while_valid_and_suspends_once_expired() { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let mock: Arc = Arc::new(MockHttpClient::new(move |_req| { + counter.fetch_add(1, Ordering::SeqCst); + async move { + Ok(AsyncRawResponse::from_bytes( + StatusCode::Ok, + etag_headers(), + Bytes::from_static(LAYOUT_V1), + )) + } + .boxed() + })); + let client = mock_client(mock); + let request = BlobClientListLayoutOptions::default(); + let cache = LayoutCache::new( + request.clone(), + Arc::new(seed_layout(&client, &request).await), + ); + + assert_eq!( + cache.current(&client).await.unwrap().ideal_endpoint(0), + Some("epv1.blob.storage.azure.net:443") + ); + assert_eq!(calls.load(Ordering::SeqCst), 1); + + { + let mut state = cache.state.lock().await; + let now = Instant::now(); + state.refresh_at = now - Duration::from_secs(1); + state.expires_at = now - Duration::from_secs(1); + state.retry_at = Some(now + Duration::from_secs(300)); + } + assert!(cache.current(&client).await.is_none()); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn current_refreshes_in_background_without_blocking() { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let mock: Arc = Arc::new(MockHttpClient::new(move |_req| { + let count = counter.fetch_add(1, Ordering::SeqCst); + async move { + let body = if count == 0 { LAYOUT_V1 } else { LAYOUT_V2 }; + Ok(AsyncRawResponse::from_bytes( + StatusCode::Ok, + etag_headers(), + Bytes::from_static(body), + )) + } + .boxed() + })); + let client = mock_client(mock); + let request = BlobClientListLayoutOptions::default(); + let cache = LayoutCache::new( + request.clone(), + Arc::new(seed_layout(&client, &request).await), + ); + + { + let mut state = cache.state.lock().await; + let now = Instant::now(); + state.refresh_at = now - Duration::from_secs(1); + state.expires_at = now + Duration::from_secs(100); + } + + assert_eq!( + cache.current(&client).await.unwrap().ideal_endpoint(0), + Some("epv1.blob.storage.azure.net:443") + ); + + let mut spins = 0; + loop { + { + let state = cache.state.lock().await; + if !state.refreshing + && state.layout.ideal_endpoint(0) == Some("epv2.blob.storage.azure.net:443") + { + assert!(state.expires_at > Instant::now() + Duration::from_secs(200)); + assert!(state.retry_at.is_none()); + break; + } + } + assert!(spins < 10_000, "background refresh did not complete"); + spins += 1; + tokio::task::yield_now().await; + } + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn refresh_failure_keeps_layout_and_backs_off() { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let mock: Arc = Arc::new(MockHttpClient::new(move |_req| { + let count = counter.fetch_add(1, Ordering::SeqCst); + async move { + if count == 0 { + Ok(AsyncRawResponse::from_bytes( + StatusCode::Ok, + etag_headers(), + Bytes::from_static(LAYOUT_V1), + )) + } else { + Ok(AsyncRawResponse::from_bytes( + StatusCode::InternalServerError, + Headers::new(), + Bytes::from_static(b""), + )) + } + } + .boxed() + })); + let client = mock_client(mock); + let request = BlobClientListLayoutOptions::default(); + let expires_at = Instant::now() + Duration::from_secs(100); + let state = Arc::new(Mutex::new(CachedLayout { + layout: Arc::new(seed_layout(&client, &request).await), + refresh_at: Instant::now(), + expires_at, + refreshing: true, + retry_at: None, + })); + + LayoutCache::refresh(client, request, state.clone()).await; + let state = state.lock().await; + assert_eq!( + state.layout.ideal_endpoint(0), + Some("epv1.blob.storage.azure.net:443") + ); + assert!(!state.refreshing); + assert_eq!(state.expires_at, expires_at); + match state.retry_at { + Some(retry_at) => assert!(retry_at > Instant::now()), + None => panic!("expected a backoff to be set"), + } + assert_eq!(calls.load(Ordering::SeqCst), 2); + } } diff --git a/sdk/storage/azure_storage_blob/src/generated/clients/append_blob_client.rs b/sdk/storage/azure_storage_blob/src/generated/clients/append_blob_client.rs index a154867532f..29c9a75b47c 100644 --- a/sdk/storage/azure_storage_blob/src/generated/clients/append_blob_client.rs +++ b/sdk/storage/azure_storage_blob/src/generated/clients/append_blob_client.rs @@ -3,15 +3,13 @@ // // Code generated by Microsoft (R) Rust Code Generator. DO NOT EDIT. -use crate::{ - generated::models::{ - AppendBlobClientAppendBlockFromUrlOptions, AppendBlobClientAppendBlockFromUrlResult, - AppendBlobClientAppendBlockOptions, AppendBlobClientAppendBlockResult, - AppendBlobClientCreateOptions, AppendBlobClientCreateResult, AppendBlobClientSealOptions, - AppendBlobClientSealResult, - }, - SessionOptions, +use crate::generated::models::{ + AppendBlobClientAppendBlockFromUrlOptions, AppendBlobClientAppendBlockFromUrlResult, + AppendBlobClientAppendBlockOptions, AppendBlobClientAppendBlockResult, + AppendBlobClientCreateOptions, AppendBlobClientCreateResult, AppendBlobClientSealOptions, + AppendBlobClientSealResult, }; +// use crate::SessionOptions; use azure_core::{ base64, error::CheckSuccessOptions, @@ -28,7 +26,7 @@ use azure_core::{ pub struct AppendBlobClient { pub(crate) endpoint: Url, pub(crate) pipeline: Pipeline, - pub(crate) session_options: Option, + // pub(crate) session_options: Option, pub(crate) version: String, } @@ -37,8 +35,8 @@ pub struct AppendBlobClient { pub struct AppendBlobClientOptions { /// Allows customization of the client. pub client_options: ClientOptions, - /// Options for session token authentication. - pub session_options: Option, + // /// Options for session token authentication. + // pub session_options: Option, /// Specifies the version of the operation to use for this request. pub version: String, } @@ -593,7 +591,7 @@ impl Default for AppendBlobClientOptions { fn default() -> Self { Self { client_options: ClientOptions::default(), - session_options: None, + // session_options: None, version: String::from(DEFAULT_VERSION), } } diff --git a/sdk/storage/azure_storage_blob/src/generated/clients/blob_client.rs b/sdk/storage/azure_storage_blob/src/generated/clients/blob_client.rs index 77e440c2414..312daab9723 100644 --- a/sdk/storage/azure_storage_blob/src/generated/clients/blob_client.rs +++ b/sdk/storage/azure_storage_blob/src/generated/clients/blob_client.rs @@ -3,25 +3,22 @@ // // Code generated by Microsoft (R) Rust Code Generator. DO NOT EDIT. -use crate::{ - generated::models::{ - AccessTier, BlobClientAbortCopyOptions, BlobClientAcquireLeaseOptions, - BlobClientAcquireLeaseResult, BlobClientBreakLeaseOptions, BlobClientBreakLeaseResult, - BlobClientChangeLeaseOptions, BlobClientChangeLeaseResult, BlobClientCreateSnapshotOptions, - BlobClientCreateSnapshotResult, BlobClientDeleteImmutabilityPolicyOptions, - BlobClientDeleteOptions, BlobClientDownloadInternalOptions, - BlobClientDownloadInternalResult, BlobClientGetAccountInfoOptions, - BlobClientGetAccountInfoResult, BlobClientGetPropertiesOptions, - BlobClientGetPropertiesResult, BlobClientGetTagsOptions, BlobClientListLayoutOptions, - BlobClientReleaseLeaseOptions, BlobClientReleaseLeaseResult, BlobClientRenewLeaseOptions, - BlobClientRenewLeaseResult, BlobClientSetImmutabilityPolicyOptions, - BlobClientSetLegalHoldOptions, BlobClientSetMetadataOptions, - BlobClientSetPropertiesOptions, BlobClientSetTagsOptions, BlobClientSetTierOptions, - BlobClientStartCopyFromUrlOptions, BlobClientStartCopyFromUrlResult, - BlobClientUndeleteOptions, BlobLayout, BlobTags, - }, - SessionOptions, +use crate::generated::models::{ + AccessTier, BlobClientAbortCopyOptions, BlobClientAcquireLeaseOptions, + BlobClientAcquireLeaseResult, BlobClientBreakLeaseOptions, BlobClientBreakLeaseResult, + BlobClientChangeLeaseOptions, BlobClientChangeLeaseResult, BlobClientCreateSnapshotOptions, + BlobClientCreateSnapshotResult, BlobClientDeleteImmutabilityPolicyOptions, + BlobClientDeleteOptions, BlobClientDownloadInternalOptions, BlobClientDownloadInternalResult, + BlobClientGetAccountInfoOptions, BlobClientGetAccountInfoResult, + BlobClientGetPropertiesOptions, BlobClientGetPropertiesResult, BlobClientGetTagsOptions, + BlobClientListLayoutOptions, BlobClientReleaseLeaseOptions, BlobClientReleaseLeaseResult, + BlobClientRenewLeaseOptions, BlobClientRenewLeaseResult, + BlobClientSetImmutabilityPolicyOptions, BlobClientSetLegalHoldOptions, + BlobClientSetMetadataOptions, BlobClientSetPropertiesOptions, BlobClientSetTagsOptions, + BlobClientSetTierOptions, BlobClientStartCopyFromUrlOptions, BlobClientStartCopyFromUrlResult, + BlobClientUndeleteOptions, BlobLayout, BlobTags, }; +// use crate::SessionOptions; use azure_core::{ base64, error::CheckSuccessOptions, @@ -41,7 +38,7 @@ use std::collections::HashMap; pub struct BlobClient { pub(crate) endpoint: Url, pub(crate) pipeline: Pipeline, - pub(crate) session_options: Option, + // pub(crate) session_options: Option, pub(crate) version: String, } @@ -50,8 +47,8 @@ pub struct BlobClient { pub struct BlobClientOptions { /// Allows customization of the client. pub client_options: ClientOptions, - /// Options for session token authentication. - pub session_options: Option, + // /// Options for session token authentication. + // pub session_options: Option, /// Specifies the version of the operation to use for this request. pub version: String, } @@ -1220,6 +1217,9 @@ impl BlobClient { ) .await?; let (status, headers, body) = rsp.deconstruct(); + // TODO: [Emitter Fix Needed] This pageable operation accepts HTTP 204 as + // success, but the generated code still attempts to deserialize its empty + // body as XML, so a valid 204 fails before returning `PagerResult::Done`. let res: BlobClientListLayoutPage = xml::from_xml(&body)?; let rsp = RawResponse::from_bytes(status, headers, body).into(); Ok(match res.next_marker { @@ -1978,7 +1978,7 @@ impl Default for BlobClientOptions { fn default() -> Self { Self { client_options: ClientOptions::default(), - session_options: None, + // session_options: None, version: String::from(DEFAULT_VERSION), } } diff --git a/sdk/storage/azure_storage_blob/src/generated/clients/blob_container_client.rs b/sdk/storage/azure_storage_blob/src/generated/clients/blob_container_client.rs index fc3f871eb22..1632f7d1681 100644 --- a/sdk/storage/azure_storage_blob/src/generated/clients/blob_container_client.rs +++ b/sdk/storage/azure_storage_blob/src/generated/clients/blob_container_client.rs @@ -3,28 +3,26 @@ // // Code generated by Microsoft (R) Rust Code Generator. DO NOT EDIT. -use crate::{ - generated::models::{ - BlobContainerClientAcquireLeaseOptions, BlobContainerClientAcquireLeaseResult, - BlobContainerClientBreakLeaseOptions, BlobContainerClientBreakLeaseResult, - BlobContainerClientChangeLeaseOptions, BlobContainerClientChangeLeaseResult, - BlobContainerClientCreateOptions, BlobContainerClientCreateSessionOptions, - BlobContainerClientDeleteOptions, BlobContainerClientFindBlobsByTagsOptions, - BlobContainerClientGetAccessPolicyOptions, BlobContainerClientGetAccountInfoOptions, - BlobContainerClientGetAccountInfoResult, BlobContainerClientGetPropertiesOptions, - BlobContainerClientGetPropertiesResult, - BlobContainerClientListBlobsHierarchicalInternalOptions, - BlobContainerClientListBlobsHierarchicalInternalResult, - BlobContainerClientListBlobsHierarchicalXmlOptions, - BlobContainerClientListBlobsInternalOptions, BlobContainerClientListBlobsInternalResult, - BlobContainerClientListBlobsXmlOptions, BlobContainerClientReleaseLeaseOptions, - BlobContainerClientReleaseLeaseResult, BlobContainerClientRenewLeaseOptions, - BlobContainerClientRenewLeaseResult, BlobContainerClientSetAccessPolicyOptions, - BlobContainerClientSetMetadataOptions, CreateSessionConfiguration, CreateSessionResponse, - FilteredBlobResponse, ListBlobsHierarchicalResponse, ListBlobsResponse, SignedIdentifiers, - }, - SessionOptions, +use crate::generated::models::{ + BlobContainerClientAcquireLeaseOptions, BlobContainerClientAcquireLeaseResult, + BlobContainerClientBreakLeaseOptions, BlobContainerClientBreakLeaseResult, + BlobContainerClientChangeLeaseOptions, BlobContainerClientChangeLeaseResult, + BlobContainerClientCreateOptions, BlobContainerClientCreateSessionOptions, + BlobContainerClientDeleteOptions, BlobContainerClientFindBlobsByTagsOptions, + BlobContainerClientGetAccessPolicyOptions, BlobContainerClientGetAccountInfoOptions, + BlobContainerClientGetAccountInfoResult, BlobContainerClientGetPropertiesOptions, + BlobContainerClientGetPropertiesResult, + BlobContainerClientListBlobsHierarchicalInternalOptions, + BlobContainerClientListBlobsHierarchicalInternalResult, + BlobContainerClientListBlobsHierarchicalXmlOptions, + BlobContainerClientListBlobsInternalOptions, BlobContainerClientListBlobsInternalResult, + BlobContainerClientListBlobsXmlOptions, BlobContainerClientReleaseLeaseOptions, + BlobContainerClientReleaseLeaseResult, BlobContainerClientRenewLeaseOptions, + BlobContainerClientRenewLeaseResult, BlobContainerClientSetAccessPolicyOptions, + BlobContainerClientSetMetadataOptions, CreateSessionConfiguration, CreateSessionResponse, + FilteredBlobResponse, ListBlobsHierarchicalResponse, ListBlobsResponse, SignedIdentifiers, }; +// use crate::SessionOptions; use azure_core::{ error::CheckSuccessOptions, fmt::SafeDebug, @@ -43,7 +41,7 @@ use std::collections::HashMap; pub struct BlobContainerClient { pub(crate) endpoint: Url, pub(crate) pipeline: Pipeline, - pub(crate) session_options: Option, + // pub(crate) session_options: Option, pub(crate) version: String, } @@ -52,8 +50,8 @@ pub struct BlobContainerClient { pub struct BlobContainerClientOptions { /// Allows customization of the client. pub client_options: ClientOptions, - /// Options for session token authentication. - pub session_options: Option, + // /// Options for session token authentication. + // pub session_options: Option, /// Specifies the version of the operation to use for this request. pub version: String, } @@ -1407,7 +1405,7 @@ impl Default for BlobContainerClientOptions { fn default() -> Self { Self { client_options: ClientOptions::default(), - session_options: None, + // session_options: None, version: String::from(DEFAULT_VERSION), } } diff --git a/sdk/storage/azure_storage_blob/src/generated/clients/blob_service_client.rs b/sdk/storage/azure_storage_blob/src/generated/clients/blob_service_client.rs index f2f3d683f94..ec400c3bb66 100644 --- a/sdk/storage/azure_storage_blob/src/generated/clients/blob_service_client.rs +++ b/sdk/storage/azure_storage_blob/src/generated/clients/blob_service_client.rs @@ -3,17 +3,15 @@ // // Code generated by Microsoft (R) Rust Code Generator. DO NOT EDIT. -use crate::{ - generated::models::{ - BlobServiceClientFindBlobsByTagsOptions, BlobServiceClientGetAccountInfoOptions, - BlobServiceClientGetAccountInfoResult, BlobServiceClientGetPropertiesOptions, - BlobServiceClientGetStatisticsOptions, BlobServiceClientGetUserDelegationKeyOptions, - BlobServiceClientListContainersOptions, BlobServiceClientSetPropertiesOptions, - BlobServiceProperties, FilteredBlobResponse, KeyInfo, ListContainersResponse, - StorageServiceStats, - }, - SessionOptions, +use crate::generated::models::{ + BlobServiceClientFindBlobsByTagsOptions, BlobServiceClientGetAccountInfoOptions, + BlobServiceClientGetAccountInfoResult, BlobServiceClientGetPropertiesOptions, + BlobServiceClientGetStatisticsOptions, BlobServiceClientGetUserDelegationKeyOptions, + BlobServiceClientListContainersOptions, BlobServiceClientSetPropertiesOptions, + BlobServiceProperties, FilteredBlobResponse, KeyInfo, ListContainersResponse, + StorageServiceStats, }; +// use crate::SessionOptions; use azure_core::{ error::CheckSuccessOptions, fmt::SafeDebug, @@ -30,7 +28,7 @@ use azure_storage_common::models::UserDelegationKey; pub struct BlobServiceClient { pub(crate) endpoint: Url, pub(crate) pipeline: Pipeline, - pub(crate) session_options: Option, + // pub(crate) session_options: Option, pub(crate) version: String, } @@ -39,8 +37,8 @@ pub struct BlobServiceClient { pub struct BlobServiceClientOptions { /// Allows customization of the client. pub client_options: ClientOptions, - /// Options for session token authentication. - pub session_options: Option, + // /// Options for session token authentication. + // pub session_options: Option, /// Specifies the version of the operation to use for this request. pub version: String, } @@ -466,7 +464,7 @@ impl Default for BlobServiceClientOptions { fn default() -> Self { Self { client_options: ClientOptions::default(), - session_options: None, + // session_options: None, version: String::from(DEFAULT_VERSION), } } diff --git a/sdk/storage/azure_storage_blob/src/generated/clients/block_blob_client.rs b/sdk/storage/azure_storage_blob/src/generated/clients/block_blob_client.rs index b282ac58342..9d1452f372e 100644 --- a/sdk/storage/azure_storage_blob/src/generated/clients/block_blob_client.rs +++ b/sdk/storage/azure_storage_blob/src/generated/clients/block_blob_client.rs @@ -3,17 +3,15 @@ // // Code generated by Microsoft (R) Rust Code Generator. DO NOT EDIT. -use crate::{ - generated::models::{ - BlockBlobClientCommitBlockListOptions, BlockBlobClientCommitBlockListResult, - BlockBlobClientGetBlockListOptions, BlockBlobClientStageBlockFromUrlOptions, - BlockBlobClientStageBlockFromUrlResult, BlockBlobClientStageBlockOptions, - BlockBlobClientStageBlockResult, BlockBlobClientUploadBlobFromUrlOptions, - BlockBlobClientUploadBlobFromUrlResult, BlockBlobClientUploadInternalOptions, - BlockBlobClientUploadInternalResult, BlockList, BlockListType, BlockLookupList, - }, - SessionOptions, +use crate::generated::models::{ + BlockBlobClientCommitBlockListOptions, BlockBlobClientCommitBlockListResult, + BlockBlobClientGetBlockListOptions, BlockBlobClientStageBlockFromUrlOptions, + BlockBlobClientStageBlockFromUrlResult, BlockBlobClientStageBlockOptions, + BlockBlobClientStageBlockResult, BlockBlobClientUploadBlobFromUrlOptions, + BlockBlobClientUploadBlobFromUrlResult, BlockBlobClientUploadInternalOptions, + BlockBlobClientUploadInternalResult, BlockList, BlockListType, BlockLookupList, }; +// use crate::SessionOptions; use azure_core::{ base64, error::CheckSuccessOptions, @@ -30,7 +28,7 @@ use azure_core::{ pub struct BlockBlobClient { pub(crate) endpoint: Url, pub(crate) pipeline: Pipeline, - pub(crate) session_options: Option, + // pub(crate) session_options: Option, pub(crate) version: String, } @@ -39,8 +37,8 @@ pub struct BlockBlobClient { pub struct BlockBlobClientOptions { /// Allows customization of the client. pub client_options: ClientOptions, - /// Options for session token authentication. - pub session_options: Option, + // /// Options for session token authentication. + // pub session_options: Option, /// Specifies the version of the operation to use for this request. pub version: String, } @@ -944,7 +942,7 @@ impl Default for BlockBlobClientOptions { fn default() -> Self { Self { client_options: ClientOptions::default(), - session_options: None, + // session_options: None, version: String::from(DEFAULT_VERSION), } } diff --git a/sdk/storage/azure_storage_blob/src/generated/clients/page_blob_client.rs b/sdk/storage/azure_storage_blob/src/generated/clients/page_blob_client.rs index 591316cf4ea..c076f84a28f 100644 --- a/sdk/storage/azure_storage_blob/src/generated/clients/page_blob_client.rs +++ b/sdk/storage/azure_storage_blob/src/generated/clients/page_blob_client.rs @@ -14,8 +14,8 @@ use crate::{ PageBlobClientUploadPagesResult, PageList, SequenceNumberActionType, }, models::HttpRange, - SessionOptions, }; +// use crate::SessionOptions; use azure_core::{ base64, error::CheckSuccessOptions, @@ -33,7 +33,7 @@ use azure_core::{ pub struct PageBlobClient { pub(crate) endpoint: Url, pub(crate) pipeline: Pipeline, - pub(crate) session_options: Option, + // pub(crate) session_options: Option, pub(crate) version: String, } @@ -42,8 +42,8 @@ pub struct PageBlobClient { pub struct PageBlobClientOptions { /// Allows customization of the client. pub client_options: ClientOptions, - /// Options for session token authentication. - pub session_options: Option, + // /// Options for session token authentication. + // pub session_options: Option, /// Specifies the version of the operation to use for this request. pub version: String, } @@ -1015,7 +1015,7 @@ impl Default for PageBlobClientOptions { fn default() -> Self { Self { client_options: ClientOptions::default(), - session_options: None, + // session_options: None, version: String::from(DEFAULT_VERSION), } } diff --git a/sdk/storage/azure_storage_blob/src/models/method_options.rs b/sdk/storage/azure_storage_blob/src/models/method_options.rs index a43712cb80d..19162346b75 100644 --- a/sdk/storage/azure_storage_blob/src/models/method_options.rs +++ b/sdk/storage/azure_storage_blob/src/models/method_options.rs @@ -16,6 +16,31 @@ use crate::models::{ ImmutabilityPolicyMode, ListBlobsIncludeItem, }; +/// Determines whether locality-aware routing is used for the parallel range +/// requests issued by a download. +/// +/// This is a performance optimization only - the bytes returned are identical +/// regardless of the mode used. When routing is enabled, the blob's layout is +/// fetched and each range request is sent to the endpoint that serves it, falling +/// back to the client's configured endpoint when no layout is available. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +#[non_exhaustive] +pub enum LayoutAwareRouting { + /// The locality-aware routing behavior is determined by the client library and + /// may change in future releases. Currently equivalent to + /// [`LayoutAwareRouting::Disabled`]. + #[default] + Auto, + + /// Never route range requests based on the blob's layout. All requests are sent + /// to the client's configured endpoint. + Disabled, + + /// Opt in to locality-aware routing. The blob's layout is fetched and each range + /// download is routed to the endpoint that serves it. + Enabled, +} + /// Options to be passed to `BlobClient::download()` #[derive(Clone, Default, SafeDebug)] pub struct BlobClientDownloadOptions<'a> { @@ -46,6 +71,13 @@ pub struct BlobClientDownloadOptions<'a> { /// If specified, the operation only succeeds if the resource's lease is active and matches this ID. pub lease_id: Option, + /// Determines whether locality-aware routing is used for the parallel range + /// requests issued by this download. This is a performance optimization only: + /// the bytes returned are identical regardless of the mode. + /// + /// Defaults to [`LayoutAwareRouting::Auto`]. + pub layout_aware_routing: LayoutAwareRouting, + /// Allows customization of the method call. pub method_options: ClientMethodOptions<'a>, diff --git a/sdk/storage/azure_storage_blob/src/models/mod.rs b/sdk/storage/azure_storage_blob/src/models/mod.rs index 1a3214deab3..9ccc9f69f7e 100644 --- a/sdk/storage/azure_storage_blob/src/models/mod.rs +++ b/sdk/storage/azure_storage_blob/src/models/mod.rs @@ -21,9 +21,9 @@ pub use download_result::{ }; pub(crate) use format::decode_next_marker; pub use format::AutoFormat; -pub use method_options::BlobClientDownloadOptions; pub use method_options::BlockBlobClientUploadOptions; pub use method_options::BlockBlobClientUploadOptions as BlobClientUploadOptions; +pub use method_options::{BlobClientDownloadOptions, LayoutAwareRouting}; pub use method_options::{ BlobContainerClientListBlobsHierarchicalOptions, BlobContainerClientListBlobsOptions, }; diff --git a/sdk/storage/azure_storage_blob/src/partitioned_transfer/download.rs b/sdk/storage/azure_storage_blob/src/partitioned_transfer/download.rs index 578e436a1e8..b0957ba7cbe 100644 --- a/sdk/storage/azure_storage_blob/src/partitioned_transfer/download.rs +++ b/sdk/storage/azure_storage_blob/src/partitioned_transfer/download.rs @@ -38,6 +38,16 @@ pub(crate) trait PartitionedDownloadBehavior { range: Option>, etag_lock: Option, ) -> AzureResult; + + /// Called once after the initial response and before any subsequent range + /// requests are issued, so the behavior can prepare per-download state. + async fn prepare( + &self, + _initial_headers: &Headers, + _etag_lock: Option<&Etag>, + ) -> AzureResult<()> { + Ok(()) + } } /// Returns a stream that runs up to parallel-many ranged downloads at a time. @@ -88,6 +98,7 @@ where Box::pin(initial_response.into_body()), )); } + client.prepare(&headers, etag_lock.as_ref()).await?; let total_chunks = remaining_ranges.len() + 1; // channel for download workers to send results to their coordinator. @@ -201,6 +212,8 @@ where return Err(insufficient_buffer_err()); } + client.prepare(&headers, etag_lock.as_ref()).await?; + // if no real parallelism, just sequentially go through the ranges and write to buffer if parallel == 1 { let mut total_read = initial_response.into_body().collect_into(buffer).await?; diff --git a/sdk/storage/azure_storage_blob/src/partitioned_transfer/layout.rs b/sdk/storage/azure_storage_blob/src/partitioned_transfer/layout.rs new file mode 100644 index 00000000000..ca3d9e6e93a --- /dev/null +++ b/sdk/storage/azure_storage_blob/src/partitioned_transfer/layout.rs @@ -0,0 +1,825 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Blob layout model and endpoint resolution for locality-aware downloads. +//! +//! The service's Get Blob Layout API describes which byte ranges of a blob are +//! served by which endpoints. This module turns those paginated responses into a +//! flat, ascending list of [`LayoutSegment`]s and resolves the serving endpoint +//! for a given byte offset via binary search. + +use std::{collections::HashMap, sync::Arc}; + +use async_trait::async_trait; +use azure_core::{ + error::ErrorKind, + http::{ + policies::{Policy, PolicyResult}, + Context, Etag, Request, StatusCode, Url, + }, + Error, Result, +}; +use futures::StreamExt as _; + +use crate::generated::{ + clients::BlobClient, + models::{BlobClientListLayoutOptions, BlobLayout}, +}; + +/// A contiguous byte range of a blob and the endpoint that serves it. +/// +/// `end` is the inclusive last byte offset, matching the service's layout ranges. +// Internal-only helper; plain `Debug` is intentional so test assertions can print +// segment contents (endpoints are host:port, not secrets). +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct LayoutSegment { + /// Inclusive start byte offset of the range. + pub start: i64, + /// Inclusive end byte offset of the range. + pub end: i64, + /// The `host:port` endpoint serving this range, or `None` when the service did + /// not provide one; such ranges download from the client's configured endpoint. + pub endpoint: Option, +} + +/// The resolved layout of a blob: non-overlapping, ascending byte-range segments. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub(crate) struct Layout { + segments: Vec, +} + +impl Layout { + /// Returns `true` when the layout contains no segments (no routing applies). + pub fn is_empty(&self) -> bool { + self.segments.is_empty() + } + + /// Appends the segments described by a single Get Blob Layout page. + /// + /// Endpoint indices are scoped to their page, so each page's `Endpoints` are + /// mapped independently before its `Ranges` are resolved and appended. + pub fn extend_from_page(&mut self, page: &BlobLayout) { + let ranges = &page.ranges.range; + if ranges.is_empty() { + return; + } + + let index_to_endpoint: HashMap = page + .endpoints + .as_ref() + .and_then(|endpoints| endpoints.endpoint.as_ref()) + .map(|endpoints| { + endpoints + .iter() + .filter_map(|endpoint| Some((endpoint.index?, endpoint.value.as_deref()?))) + .collect() + }) + .unwrap_or_default(); + + self.segments.reserve(ranges.len()); + for range in ranges { + let endpoint = range + .endpoint_index + .and_then(|index| index_to_endpoint.get(&index).copied()) + .filter(|value| !value.is_empty()) + .map(str::to_owned); + self.segments.push(LayoutSegment { + start: range.start.unwrap_or_default(), + end: range.end.unwrap_or_default(), + endpoint, + }); + } + } + + /// Resolves the serving endpoint for the segment covering `offset`. + /// + /// Uses binary search to find the first segment whose inclusive `end` is at or + /// beyond `offset`. Returns `None` when no segment covers the offset or the + /// covering segment has no endpoint; callers then fall back to the client's + /// configured endpoint. The bytes returned are identical regardless. + pub fn ideal_endpoint(&self, offset: i64) -> Option<&str> { + if self.segments.is_empty() { + return None; + } + + let mut lo: isize = 0; + let mut hi: isize = self.segments.len() as isize - 1; + let mut overlap: Option = None; + while lo <= hi { + let mid = lo + (hi - lo) / 2; + if self.segments[mid as usize].end >= offset { + overlap = Some(mid as usize); + hi = mid - 1; + } else { + lo = mid + 1; + } + } + + overlap.and_then(|index| self.segments[index].endpoint.as_deref()) + } +} + +/// Context value carrying the layout endpoint (`host:port` or an absolute URL) that +/// a single download range request should be routed to. +#[derive(Clone, Debug)] +pub(crate) struct LayoutEndpoint(pub String); + +/// Per-call pipeline policy that reroutes a request to a layout endpoint while +/// preserving the original account authority as the `Host` header. +/// +/// It is a no-op unless a [`LayoutEndpoint`] is present in the request context, so +/// it is safe to install on every request. A malformed endpoint fails the request +/// so that a bad layout response surfaces loudly rather than silently degrading. +#[derive(Debug)] +pub(crate) struct LayoutRoutingPolicy; + +#[async_trait] +impl Policy for LayoutRoutingPolicy { + async fn send( + &self, + ctx: &Context, + request: &mut Request, + next: &[Arc], + ) -> PolicyResult { + if let Some(LayoutEndpoint(endpoint)) = ctx.value::() { + apply_layout_endpoint(request, endpoint)?; + } + next[0].send(ctx, request, &next[1..]).await + } +} + +/// Rewrites `request`'s URL authority to `endpoint`, preserving the original +/// authority as an explicit `Host` header. +/// +/// The rewrite is computed on a scratch copy and only committed once fully +/// successful, so a malformed endpoint returns an error without leaving the +/// request in a half-rewritten state. +fn apply_layout_endpoint(request: &mut Request, endpoint: &str) -> azure_core::Result<()> { + let (host, port) = parse_endpoint_authority(endpoint).ok_or_else(|| { + Error::with_message( + ErrorKind::Other, + format!("invalid layout endpoint {endpoint:?}"), + ) + })?; + let original_host = original_host_header(request.url()); + let mut rewritten = request.url().clone(); + rewritten + .set_host(Some(&host)) + .map_err(|e| Error::with_error(ErrorKind::Other, e, "invalid layout endpoint host"))?; + rewritten + .set_port(port) + .map_err(|()| Error::with_message(ErrorKind::Other, "invalid layout endpoint port"))?; + request.insert_header("host", original_host); + *request.url_mut() = rewritten; + Ok(()) +} + +/// Derives the `Host` header value the client would send for `url`, i.e. the host +/// plus the port when a non-default port is explicitly present. +fn original_host_header(url: &Url) -> String { + match (url.host_str(), url.port()) { + (Some(host), Some(port)) => format!("{host}:{port}"), + (Some(host), None) => host.to_owned(), + (None, _) => String::new(), + } +} + +/// Parses a layout endpoint value into a host and optional port. +/// +/// Accepts both the documented `host:port` form and an absolute URL form +/// (`scheme://host[:port]`). Returns `None` for values that cannot be interpreted +/// as an authority; the caller turns that into a failed request. +fn parse_endpoint_authority(endpoint: &str) -> Option<(String, Option)> { + let endpoint = endpoint.trim(); + if endpoint.is_empty() { + return None; + } + + if endpoint.contains("://") { + let url = Url::parse(endpoint).ok()?; + let host = url.host_str()?.to_owned(); + return Some((host, url.port())); + } + + match endpoint.rsplit_once(':') { + Some((host, port)) if !host.is_empty() => Some((host.to_owned(), Some(port.parse().ok()?))), + Some(_) => None, + None => Some((endpoint.to_owned(), None)), + } +} + +/// The outcome of prefetching a blob's layout for a locality-aware download. +pub(crate) struct LayoutPrefetch { + /// The resolved, non-empty layout used to route range requests. + pub layout: Layout, + /// The ETag pinning the download to a single blob version: the caller-supplied + /// condition when present, otherwise the ETag from the first layout page. + pub etag: Option, +} + +/// Fetches a blob's layout for locality-aware routing, following pagination. +/// +/// Returns `Ok(Some(_))` when a non-empty layout is available; `Ok(None)` when the +/// download should proceed normally without routing (no layout, HTTP 400, or 5xx); +/// and `Err(_)` when the download must fail (HTTP 403/404/409/412, a transport +/// error, or a response deserialization error). +/// +/// The generated pager owns request construction, response handling, and +/// continuation. The first page's ETag pins the subsequent blob download when the +/// caller did not already supply an `If-Match` condition. +pub(crate) async fn fetch_layout( + client: &BlobClient, + context: &Context<'_>, + options: &BlobClientListLayoutOptions<'_>, +) -> Result> { + let mut layout = Layout::default(); + let mut locked_etag = options.if_match.clone(); + let mut options = options.clone(); + options.method_options.context = context.clone().into_owned(); + let mut pages = client.list_layout(Some(options))?; + + while let Some(response) = pages.next().await { + let response = match response { + Ok(response) => response, + Err(err) => return classify_layout_error(err), + }; + if locked_etag.is_none() { + locked_etag = response + .headers() + .get_optional_str(&"etag".into()) + .map(Etag::from); + } + let page = response.into_model()?; + layout.extend_from_page(&page); + } + + if layout.is_empty() { + return Ok(None); + } + Ok(Some(LayoutPrefetch { + layout, + etag: locked_etag, + })) +} + +/// Maps a Get Blob Layout failure to a graceful fall-back (`Ok(None)`) or a hard +/// failure (`Err`). +/// +/// HTTP 400 (layout unsupported) and 5xx (transient) fall back to a normal +/// download; every other failure (403/404/409/412 and transport errors) fails. +fn classify_layout_error(err: Error) -> Result> { + match err.http_status() { + Some(status) if status == StatusCode::BadRequest || status.is_server_error() => Ok(None), + _ => Err(err), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::generated::{ + clients::BlobClientOptions, + models::{BlobLayoutEndpoint, BlobLayoutEndpoints, BlobLayoutRange, BlobLayoutRanges}, + }; + use azure_core::{ + http::{ + headers::Headers, AsyncRawResponse, ClientOptions, FixedRetryOptions, HttpClient, + Method, RetryOptions, Transport, + }, + Bytes, + }; + use azure_core_test::http::MockHttpClient; + use futures::FutureExt as _; + + fn segment(start: i64, end: i64, endpoint: Option<&str>) -> LayoutSegment { + LayoutSegment { + start, + end, + endpoint: endpoint.map(str::to_owned), + } + } + + fn layout(segments: Vec) -> Layout { + Layout { segments } + } + + fn endpoint(index: i32, value: &str) -> BlobLayoutEndpoint { + BlobLayoutEndpoint { + index: Some(index), + value: Some(value.to_owned()), + } + } + + fn range(start: i64, end: i64, endpoint_index: i32) -> BlobLayoutRange { + BlobLayoutRange { + start: Some(start), + end: Some(end), + endpoint_index: Some(endpoint_index), + } + } + + fn page(endpoints: Vec, ranges: Vec) -> BlobLayout { + BlobLayout { + endpoints: Some(BlobLayoutEndpoints { + endpoint: Some(endpoints), + }), + ranges: BlobLayoutRanges { range: ranges }, + ..Default::default() + } + } + + #[test] + fn empty_layout_resolves_to_none() { + let layout = Layout::default(); + assert!(layout.is_empty()); + assert_eq!(layout.ideal_endpoint(0), None); + assert_eq!(layout.ideal_endpoint(1_000), None); + } + + #[test] + fn single_segment_covers_all_contained_offsets() { + let layout = layout(vec![segment(0, 99, Some("a"))]); + assert_eq!(layout.ideal_endpoint(0), Some("a")); + assert_eq!(layout.ideal_endpoint(50), Some("a")); + assert_eq!(layout.ideal_endpoint(99), Some("a")); + // Beyond the final segment's end: no segment covers it. + assert_eq!(layout.ideal_endpoint(100), None); + } + + #[test] + fn binary_search_selects_covering_segment() { + let layout = layout(vec![ + segment(0, 9, Some("a")), + segment(10, 19, Some("b")), + segment(20, 29, Some("c")), + ]); + assert_eq!(layout.ideal_endpoint(0), Some("a")); + assert_eq!(layout.ideal_endpoint(9), Some("a")); + assert_eq!(layout.ideal_endpoint(10), Some("b")); + assert_eq!(layout.ideal_endpoint(15), Some("b")); + assert_eq!(layout.ideal_endpoint(19), Some("b")); + assert_eq!(layout.ideal_endpoint(20), Some("c")); + assert_eq!(layout.ideal_endpoint(29), Some("c")); + assert_eq!(layout.ideal_endpoint(30), None); + } + + #[test] + fn inclusive_end_boundaries_route_correctly() { + let layout = layout(vec![segment(0, 9, Some("a")), segment(10, 19, Some("b"))]); + // end is inclusive: offset == first segment's end stays on the first segment. + assert_eq!(layout.ideal_endpoint(9), Some("a")); + // one past it moves to the next segment. + assert_eq!(layout.ideal_endpoint(10), Some("b")); + } + + #[test] + fn segment_without_endpoint_resolves_to_none() { + let layout = layout(vec![segment(0, 9, None), segment(10, 19, Some("b"))]); + assert_eq!(layout.ideal_endpoint(5), None); + assert_eq!(layout.ideal_endpoint(15), Some("b")); + } + + #[test] + fn extend_from_page_maps_endpoint_indices() { + let mut layout = Layout::default(); + layout.extend_from_page(&page( + vec![endpoint(0, "h0:443"), endpoint(1, "h1:443")], + vec![range(0, 9, 1), range(10, 19, 0)], + )); + assert_eq!( + layout.segments, + vec![ + segment(0, 9, Some("h1:443")), + segment(10, 19, Some("h0:443")), + ] + ); + } + + #[test] + fn extend_from_page_handles_unordered_endpoint_indices() { + let mut layout = Layout::default(); + layout.extend_from_page(&page( + vec![endpoint(2, "h2:443"), endpoint(0, "h0:443")], + vec![range(0, 9, 2), range(10, 19, 0)], + )); + assert_eq!( + layout.segments, + vec![ + segment(0, 9, Some("h2:443")), + segment(10, 19, Some("h0:443")), + ] + ); + } + + #[test] + fn extend_from_page_missing_endpoint_index_yields_none() { + let mut layout = Layout::default(); + layout.extend_from_page(&page(vec![endpoint(0, "h0:443")], vec![range(0, 9, 7)])); + assert_eq!(layout.segments, vec![segment(0, 9, None)]); + } + + #[test] + fn extend_from_page_accumulates_across_pages_with_independent_indices() { + let mut layout = Layout::default(); + // Page 1: index 0 -> "p1:443". + layout.extend_from_page(&page(vec![endpoint(0, "p1:443")], vec![range(0, 9, 0)])); + // Page 2: index 0 -> "p2:443" (different endpoint space). + layout.extend_from_page(&page(vec![endpoint(0, "p2:443")], vec![range(10, 19, 0)])); + assert_eq!( + layout.segments, + vec![ + segment(0, 9, Some("p1:443")), + segment(10, 19, Some("p2:443")), + ] + ); + } + + #[test] + fn extend_from_page_with_empty_ranges_adds_nothing() { + let mut layout = Layout::default(); + layout.extend_from_page(&page(vec![endpoint(0, "h0:443")], vec![])); + assert!(layout.is_empty()); + } + + fn request_to(url: &str) -> Request { + Request::new(url.parse().unwrap(), Method::Get) + } + + fn host_header(request: &Request) -> Option { + request + .headers() + .get_optional_str(&"host".into()) + .map(str::to_owned) + } + + #[test] + fn parse_endpoint_authority_forms() { + assert_eq!( + parse_endpoint_authority("host.example.net:443"), + Some(("host.example.net".to_owned(), Some(443))) + ); + assert_eq!( + parse_endpoint_authority("host.example.net:8443"), + Some(("host.example.net".to_owned(), Some(8443))) + ); + assert_eq!( + parse_endpoint_authority("https://host.example.net:8443"), + Some(("host.example.net".to_owned(), Some(8443))) + ); + // Absolute URL with default port normalizes the port away. + assert_eq!( + parse_endpoint_authority("https://host.example.net"), + Some(("host.example.net".to_owned(), None)) + ); + // Bare host, no port. + assert_eq!( + parse_endpoint_authority("host.example.net"), + Some(("host.example.net".to_owned(), None)) + ); + // Surrounding whitespace is tolerated. + assert_eq!( + parse_endpoint_authority(" host.example.net:443 "), + Some(("host.example.net".to_owned(), Some(443))) + ); + // Malformed inputs yield None so the caller skips routing. + assert_eq!(parse_endpoint_authority(""), None); + assert_eq!(parse_endpoint_authority(":443"), None); + assert_eq!(parse_endpoint_authority("host.example.net:port"), None); + } + + #[test] + fn apply_layout_endpoint_rewrites_url_and_preserves_host() { + let mut request = request_to("https://acct.blob.core.windows.net/container/blob"); + apply_layout_endpoint(&mut request, "target.blob.storage.azure.net:443").unwrap(); + + assert_eq!( + request.url().host_str(), + Some("target.blob.storage.azure.net") + ); + assert_eq!(request.url().path(), "/container/blob"); + assert_eq!( + host_header(&request).as_deref(), + Some("acct.blob.core.windows.net") + ); + } + + #[test] + fn apply_layout_endpoint_preserves_original_non_default_port_in_host() { + let mut request = request_to("https://acct.blob.core.windows.net:10000/container/blob"); + apply_layout_endpoint(&mut request, "target.blob.storage.azure.net:443").unwrap(); + + assert_eq!( + request.url().host_str(), + Some("target.blob.storage.azure.net") + ); + assert_eq!( + host_header(&request).as_deref(), + Some("acct.blob.core.windows.net:10000") + ); + } + + #[test] + fn apply_layout_endpoint_absolute_url_form() { + let mut request = request_to("https://acct.blob.core.windows.net/container/blob"); + apply_layout_endpoint(&mut request, "https://target.blob.storage.azure.net:8443").unwrap(); + + assert_eq!( + request.url().host_str(), + Some("target.blob.storage.azure.net") + ); + assert_eq!(request.url().port(), Some(8443)); + assert_eq!( + host_header(&request).as_deref(), + Some("acct.blob.core.windows.net") + ); + } + + #[test] + fn apply_layout_endpoint_malformed_fails_and_leaves_request_untouched() { + let mut request = request_to("https://acct.blob.core.windows.net/container/blob"); + let result = apply_layout_endpoint(&mut request, ""); + + // Malformed endpoint fails the request and does not mutate it. + assert!(result.is_err()); + assert_eq!(request.url().host_str(), Some("acct.blob.core.windows.net")); + assert_eq!(host_header(&request), None); + } + + const LAYOUT_SINGLE_PAGE: &[u8] = br#" + + + + + + + + + +"#; + + const LAYOUT_PAGE_1: &[u8] = br#" + + + + + + + + m2 +"#; + + const LAYOUT_PAGE_2: &[u8] = br#" + + + + + + + +"#; + + const LAYOUT_EMPTY: &[u8] = br#" +"#; + + fn layout_client(transport: Arc) -> BlobClient { + BlobClient::new( + "https://acct.blob.core.windows.net/container/blob" + .parse() + .unwrap(), + None, + Some(BlobClientOptions { + client_options: ClientOptions { + transport: Some(Transport::new(transport)), + // Keep error-path tests fast: don't retry mocked 5xx responses. + retry: RetryOptions::fixed(FixedRetryOptions { + max_retries: 0, + ..Default::default() + }), + ..Default::default() + }, + ..Default::default() + }), + ) + .unwrap() + } + + fn headers_with_etag(etag: &str) -> Headers { + let mut headers = Headers::new(); + headers.insert("etag", etag.to_owned()); + headers + } + + #[tokio::test] + async fn fetch_layout_single_page_builds_layout() { + let mock: Arc = Arc::new(MockHttpClient::new(|req| { + assert!(req + .url() + .query() + .is_some_and(|query| query.contains("comp=layout"))); + async move { + Ok(AsyncRawResponse::from_bytes( + StatusCode::Ok, + headers_with_etag("etag-1"), + Bytes::from_static(LAYOUT_SINGLE_PAGE), + )) + } + .boxed() + })); + + let client = layout_client(mock); + let prefetch = fetch_layout( + &client, + &Context::new(), + &BlobClientListLayoutOptions::default(), + ) + .await + .unwrap() + .expect("routing should be available"); + + assert_eq!(prefetch.etag, Some(Etag::from("etag-1"))); + assert_eq!( + prefetch.layout.ideal_endpoint(0), + Some("ep0.blob.storage.azure.net:443") + ); + assert_eq!( + prefetch.layout.ideal_endpoint(4_194_304), + Some("ep1.blob.storage.azure.net:443") + ); + } + + #[tokio::test] + async fn fetch_layout_paginates_with_caller_etag() { + let mock: Arc = Arc::new(MockHttpClient::new(|req| { + let query = req.url().query().unwrap_or_default().to_owned(); + let if_match = req + .headers() + .get_optional_str(&"if-match".into()) + .map(str::to_owned); + let body: &[u8] = if query.contains("marker=m2") { + assert_eq!(if_match.as_deref(), Some("etag-1")); + LAYOUT_PAGE_2 + } else { + assert_eq!(if_match.as_deref(), Some("etag-1")); + LAYOUT_PAGE_1 + }; + async move { + Ok(AsyncRawResponse::from_bytes( + StatusCode::Ok, + headers_with_etag("etag-1"), + Bytes::copy_from_slice(body), + )) + } + .boxed() + })); + + let client = layout_client(mock); + let prefetch = fetch_layout( + &client, + &Context::new(), + &BlobClientListLayoutOptions { + if_match: Some(Etag::from("etag-1")), + ..Default::default() + }, + ) + .await + .unwrap() + .expect("routing should be available"); + + assert_eq!(prefetch.etag, Some(Etag::from("etag-1"))); + assert_eq!( + prefetch.layout.ideal_endpoint(0), + Some("ep0.blob.storage.azure.net:443") + ); + assert_eq!( + prefetch.layout.ideal_endpoint(4_194_304), + Some("ep1.blob.storage.azure.net:443") + ); + } + + // This is a characterization test for the current emitter bug: 204 is an accepted + // success response, but the generated pager attempts to deserialize its empty body as XML. + // Once the emitter is fixed, this test should be changed to assert the corrected empty or + // optional result generated for a successful 204 response. + #[tokio::test] + async fn generated_list_layout_fails_to_deserialize_valid_no_content_response() { + let mock: Arc = Arc::new(MockHttpClient::new(|_req| { + async move { + Ok(AsyncRawResponse::from_bytes( + StatusCode::NoContent, + Headers::new(), + Bytes::new(), + )) + } + .boxed() + })); + + let client = layout_client(mock); + let mut pages = client.list_layout(None).unwrap(); + + let error = pages + .next() + .await + .expect("the service returned one response") + .expect_err("HTTP 204 should expose the generated XML deserialization bug"); + + assert_eq!(*error.kind(), ErrorKind::DataConversion); + assert!(error + .to_string() + .contains("failed to deserialize the following xml")); + } + + #[tokio::test] + async fn fetch_layout_empty_ranges_falls_back() { + let mock: Arc = Arc::new(MockHttpClient::new(|_req| { + async move { + Ok(AsyncRawResponse::from_bytes( + StatusCode::Ok, + headers_with_etag("etag-1"), + Bytes::from_static(LAYOUT_EMPTY), + )) + } + .boxed() + })); + + let client = layout_client(mock); + let result = fetch_layout( + &client, + &Context::new(), + &BlobClientListLayoutOptions::default(), + ) + .await + .unwrap(); + assert!(result.is_none()); + } + + #[tokio::test] + async fn fetch_layout_bad_request_falls_back() { + let mock: Arc = Arc::new(MockHttpClient::new(|_req| { + async move { + Ok(AsyncRawResponse::from_bytes( + StatusCode::BadRequest, + Headers::new(), + Bytes::new(), + )) + } + .boxed() + })); + + let client = layout_client(mock); + let result = fetch_layout( + &client, + &Context::new(), + &BlobClientListLayoutOptions::default(), + ) + .await + .unwrap(); + assert!(result.is_none()); + } + + #[tokio::test] + async fn fetch_layout_server_error_falls_back() { + let mock: Arc = Arc::new(MockHttpClient::new(|_req| { + async move { + Ok(AsyncRawResponse::from_bytes( + StatusCode::InternalServerError, + Headers::new(), + Bytes::new(), + )) + } + .boxed() + })); + + let client = layout_client(mock); + let result = fetch_layout( + &client, + &Context::new(), + &BlobClientListLayoutOptions::default(), + ) + .await + .unwrap(); + assert!(result.is_none()); + } + + #[tokio::test] + async fn fetch_layout_forbidden_fails() { + let mock: Arc = Arc::new(MockHttpClient::new(|_req| { + async move { + Ok(AsyncRawResponse::from_bytes( + StatusCode::Forbidden, + Headers::new(), + Bytes::new(), + )) + } + .boxed() + })); + + let client = layout_client(mock); + let result = fetch_layout( + &client, + &Context::new(), + &BlobClientListLayoutOptions::default(), + ) + .await; + assert!(result.is_err()); + } +} diff --git a/sdk/storage/azure_storage_blob/src/partitioned_transfer/mod.rs b/sdk/storage/azure_storage_blob/src/partitioned_transfer/mod.rs index 716bc4ce9ea..f7ee6ce14eb 100644 --- a/sdk/storage/azure_storage_blob/src/partitioned_transfer/mod.rs +++ b/sdk/storage/azure_storage_blob/src/partitioned_transfer/mod.rs @@ -3,9 +3,11 @@ pub(crate) mod defaults; mod download; +mod layout; mod upload; pub(crate) use download::*; +pub(crate) use layout::{fetch_layout, Layout, LayoutEndpoint, LayoutRoutingPolicy}; pub(crate) use upload::*; use std::{cmp::max, future::Future, num::NonZero, pin::Pin}; diff --git a/sdk/storage/azure_storage_blob/tests/blob_client.rs b/sdk/storage/azure_storage_blob/tests/blob_client.rs index 59295cd1561..263fd30f36e 100644 --- a/sdk/storage/azure_storage_blob/tests/blob_client.rs +++ b/sdk/storage/azure_storage_blob/tests/blob_client.rs @@ -3,17 +3,21 @@ mod common; +use async_trait::async_trait; use azure_core::{ + credentials::TokenCredential, error::ErrorKind, http::{ headers::{HeaderName, Headers, CONTENT_TYPE}, - AsyncRawResponse, ClientOptions, Method, RequestContent, StatusCode, Transport, Url, - XmlFormat, + policies::{Policy, PolicyResult}, + AsyncRawResponse, ClientOptions, Context, Method, Request, RequestContent, StatusCode, + Transport, Url, XmlFormat, }, time::{parse_rfc3339, to_rfc3339, OffsetDateTime}, Bytes, }; use azure_core_test::{http::MockHttpClient, recorded, Matcher, TestContext, TestMode, VarOptions}; +use azure_identity::DeveloperToolsCredential; use azure_storage_blob::{ models::{ AccessTier, AccountKind, ArchiveStatus, BlobClientAcquireLeaseOptions, @@ -23,8 +27,8 @@ use azure_storage_blob::{ BlobClientSetImmutabilityPolicyOptions, BlobClientSetMetadataOptions, BlobClientSetPropertiesOptions, BlobClientSetTierOptions, BlobClientStartCopyFromUrlResultHeaders, BlobTags, BlockBlobClientCommitBlockListOptions, - BlockBlobClientUploadOptions, CopyStatus, ImmutabilityPolicyMode, KeyInfo, LeaseState, - RehydratePriority, StorageErrorCode, + BlockBlobClientUploadOptions, CopyStatus, ImmutabilityPolicyMode, KeyInfo, + LayoutAwareRouting, LeaseState, RehydratePriority, StorageErrorCode, }, BlobClient, BlobClientOptions, BlobContainerClient, BlobContainerClientOptions, StorageError, }; @@ -44,7 +48,7 @@ use std::{ num::NonZero, sync::{ atomic::{AtomicUsize, Ordering}, - mpsc, Arc, + mpsc, Arc, Mutex, }, time::Duration, }; @@ -1528,6 +1532,403 @@ async fn test_managed_download_into_etag_lock(ctx: TestContext) -> Result<(), Bo Ok(()) } +fn parse_bytes_range(range: &str) -> (usize, usize) { + let spec = range.strip_prefix("bytes=").expect("bytes= prefix"); + let (start, end) = spec.split_once('-').expect("start-end"); + ( + start.parse().expect("range start"), + end.parse().expect("range end"), + ) +} + +fn live_layout_blob_url() -> Option { + let account = std::env::var("AZURE_STORAGE_ACCOUNT_NAME").ok()?; + let account = account.trim().trim_matches('"').trim(); + if account.is_empty() { + return None; + } + Url::parse(&format!( + "https://{account}.blob.core.windows.net/layout-routing-test/routing-test-blob" + )) + .ok() +} + +#[tokio::test] +async fn test_download_layout_aware_routing_routes_chunks() -> Result<(), Box> { + const DATA: [u8; 12] = [10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21]; + const LAYOUT: &[u8] = br#" + + + + + + + + + + + +"#; + + #[derive(Clone)] + struct Seen { + is_layout: bool, + url_host: Option, + host_header: Option, + range: Option, + } + + let seen = Arc::new(Mutex::new(Vec::::new())); + let seen_capture = seen.clone(); + let mock_client = Arc::new(MockHttpClient::new(move |request| { + let is_layout = request + .url() + .query() + .is_some_and(|query| query.contains("comp=layout")); + let url_host = request.url().host_str().map(str::to_owned); + let host_header = request + .headers() + .get_optional_str(&"host".into()) + .map(str::to_owned); + let range = request + .headers() + .get_optional_str(&"range".into()) + .map(str::to_owned); + seen_capture.lock().unwrap().push(Seen { + is_layout, + url_host, + host_header, + range: range.clone(), + }); + + let response = if is_layout { + let mut headers = Headers::new(); + headers.insert("etag", "\"routing-etag\""); + AsyncRawResponse::from_bytes(StatusCode::Ok, headers, Bytes::from_static(LAYOUT)) + } else { + let range = range.expect("data request must carry a range header"); + let (start, end) = parse_bytes_range(&range); + let slice = DATA[start..=end].to_vec(); + let mut headers = Headers::new(); + headers.insert( + "content-range", + format!("bytes {start}-{end}/{}", DATA.len()), + ); + headers.insert("content-length", slice.len().to_string()); + headers.insert("etag", "\"routing-etag\""); + headers.insert("x-ms-download-hint", "layout"); + AsyncRawResponse::from_bytes(StatusCode::PartialContent, headers, Bytes::from(slice)) + }; + async move { Ok(response) }.boxed() + })); + + let blob_client = BlobClient::new( + Url::parse("https://acct.blob.core.windows.net/container/blob")?, + None, + Some(BlobClientOptions { + client_options: ClientOptions { + transport: Some(Transport::new(mock_client)), + ..Default::default() + }, + ..Default::default() + }), + )?; + + let body = blob_client + .download(Some(BlobClientDownloadOptions { + layout_aware_routing: LayoutAwareRouting::Enabled, + partition_size: Some(NonZero::new(4).unwrap()), + parallel: Some(NonZero::new(2).unwrap()), + ..Default::default() + })) + .await? + .body + .collect() + .await?; + + assert_eq!(&body[..], &DATA[..]); + let seen = seen.lock().unwrap(); + let layout_requests: Vec<&Seen> = seen.iter().filter(|request| request.is_layout).collect(); + let data_requests: Vec<&Seen> = seen.iter().filter(|request| !request.is_layout).collect(); + assert_eq!(layout_requests.len(), 1); + assert_eq!( + layout_requests[0].url_host.as_deref(), + Some("acct.blob.core.windows.net") + ); + + assert_eq!(data_requests.len(), 3); + for request in &data_requests { + let range = request.range.as_deref().expect("range header"); + match range { + "bytes=0-3" => assert_eq!( + request.url_host.as_deref(), + Some("acct.blob.core.windows.net"), + "the initial chunk should not be routed" + ), + "bytes=4-7" | "bytes=8-11" => { + let expected = if range == "bytes=4-7" { + "epb.blob.core.windows.net" + } else { + "epc.blob.core.windows.net" + }; + assert_eq!(request.url_host.as_deref(), Some(expected)); + assert_eq!( + request.host_header.as_deref(), + Some("acct.blob.core.windows.net") + ); + } + other => panic!("unexpected data range: {other}"), + } + } + Ok(()) +} + +#[tokio::test] +async fn test_download_layout_aware_routing_skips_without_hint() -> Result<(), Box> { + const DATA: [u8; 8] = [10, 11, 12, 13, 14, 15, 16, 17]; + + let layout_requests = Arc::new(AtomicUsize::new(0)); + let hosts = Arc::new(Mutex::new(Vec::>::new())); + let layout_capture = layout_requests.clone(); + let hosts_capture = hosts.clone(); + let mock_client = Arc::new(MockHttpClient::new(move |request| { + let is_layout = request + .url() + .query() + .is_some_and(|query| query.contains("comp=layout")); + hosts_capture + .lock() + .unwrap() + .push(request.url().host_str().map(str::to_owned)); + let range = request + .headers() + .get_optional_str(&"range".into()) + .map(str::to_owned); + + let response = if is_layout { + layout_capture.fetch_add(1, Ordering::SeqCst); + AsyncRawResponse::from_bytes(StatusCode::NoContent, Headers::new(), Bytes::new()) + } else { + let range = range.expect("data request must carry a range header"); + let (start, end) = parse_bytes_range(&range); + let slice = DATA[start..=end].to_vec(); + let mut headers = Headers::new(); + headers.insert( + "content-range", + format!("bytes {start}-{end}/{}", DATA.len()), + ); + headers.insert("content-length", slice.len().to_string()); + headers.insert("etag", "\"no-hint-etag\""); + AsyncRawResponse::from_bytes(StatusCode::PartialContent, headers, Bytes::from(slice)) + }; + async move { Ok(response) }.boxed() + })); + + let blob_client = BlobClient::new( + Url::parse("https://acct.blob.core.windows.net/container/blob")?, + None, + Some(BlobClientOptions { + client_options: ClientOptions { + transport: Some(Transport::new(mock_client)), + ..Default::default() + }, + ..Default::default() + }), + )?; + + let body = blob_client + .download(Some(BlobClientDownloadOptions { + layout_aware_routing: LayoutAwareRouting::Enabled, + partition_size: Some(NonZero::new(4).unwrap()), + parallel: Some(NonZero::new(2).unwrap()), + ..Default::default() + })) + .await? + .body + .collect() + .await?; + + assert_eq!(&body[..], &DATA[..]); + assert_eq!(layout_requests.load(Ordering::SeqCst), 0); + for host in hosts.lock().unwrap().iter() { + assert_eq!(host.as_deref(), Some("acct.blob.core.windows.net")); + } + Ok(()) +} + +#[derive(Debug)] +struct RoutingObserver { + seen: Arc>>, +} + +#[derive(Debug)] +struct RoutedRequest { + is_layout: bool, + range: Option, + sent_to: Option, + original_host: Option, + download_hint: Option, + status: StatusCode, +} + +#[async_trait] +impl Policy for RoutingObserver { + async fn send( + &self, + ctx: &Context, + request: &mut Request, + next: &[Arc], + ) -> PolicyResult { + let is_layout = request + .url() + .query() + .is_some_and(|query| query.contains("comp=layout")); + let range = request + .headers() + .get_optional_str(&"range".into()) + .map(str::to_owned); + let sent_to = request.url().host_str().map(str::to_owned); + let original_host = request + .headers() + .get_optional_str(&"host".into()) + .map(str::to_owned); + let response = next[0].send(ctx, request, &next[1..]).await?; + let download_hint = response + .headers() + .get_optional_str(&"x-ms-download-hint".into()) + .map(str::to_owned); + self.seen.lock().unwrap().push(RoutedRequest { + is_layout, + range, + sent_to, + original_host, + download_hint, + status: response.status(), + }); + Ok(response) + } +} + +#[recorded::test(live)] +async fn test_download_layout_aware_routing() -> Result<(), Box> { + let Some(url) = live_layout_blob_url() else { + eprintln!( + "skipping test_download_layout_aware_routing: set AZURE_STORAGE_ACCOUNT_NAME to run it" + ); + return Ok(()); + }; + let credential: Arc = DeveloperToolsCredential::new(None)?; + let account_host = url.host_str().unwrap_or_default().to_owned(); + let observed = Arc::new(Mutex::new(Vec::::new())); + let observer: Arc = Arc::new(RoutingObserver { + seen: observed.clone(), + }); + + let mut container_url = url.clone(); + container_url + .path_segments_mut() + .expect("blob URL must be a base") + .pop(); + let container_client = BlobContainerClient::new(container_url, Some(credential.clone()), None)?; + if let Err(error) = container_client.create(None).await { + if error.http_status() != Some(StatusCode::Conflict) { + return Err(error.into()); + } + } + + let blob_client = BlobClient::new( + url, + Some(credential), + Some(BlobClientOptions { + client_options: ClientOptions { + per_try_policies: vec![observer], + ..Default::default() + }, + ..Default::default() + }), + )?; + + const BLOB_SIZE: usize = 64 * 1024 * 1024; + let data: Vec = (0..BLOB_SIZE).map(|index| (index % 251) as u8).collect(); + blob_client + .upload(RequestContent::from(data.clone()), None) + .await?; + observed.lock().unwrap().clear(); + + let body = blob_client + .download(Some(BlobClientDownloadOptions { + layout_aware_routing: LayoutAwareRouting::Enabled, + partition_size: Some(NonZero::new(16 * 1024 * 1024).unwrap()), + parallel: Some(NonZero::new(4).unwrap()), + ..Default::default() + })) + .await? + .body + .collect() + .await?; + assert_eq!(&body[..], &data[..]); + + let mut buffer = vec![0u8; data.len()]; + let result = blob_client + .download_into( + &mut buffer, + Some(BlobClientDownloadOptions { + layout_aware_routing: LayoutAwareRouting::Enabled, + partition_size: Some(NonZero::new(16 * 1024 * 1024).unwrap()), + parallel: Some(NonZero::new(4).unwrap()), + ..Default::default() + }), + ) + .await?; + assert_eq!(result.len, data.len()); + assert_eq!(&buffer[..], &data[..]); + + let observed = observed.lock().unwrap(); + let layout_calls = observed.iter().filter(|request| request.is_layout).count(); + let data_chunks = observed.iter().filter(|request| !request.is_layout).count(); + let routed: Vec<&RoutedRequest> = observed + .iter() + .filter(|request| !request.is_layout && request.original_host.is_some()) + .collect(); + + eprintln!(); + eprintln!("=== layout-aware routing ======================================"); + eprintln!(" account host : {account_host}"); + eprintln!(" Get Blob Layout calls : {layout_calls}"); + eprintln!( + " data chunks : {data_chunks} ({} routed)", + routed.len() + ); + for request in observed.iter() { + let kind = if request.is_layout { "layout" } else { "data" }; + let range = request.range.as_deref().unwrap_or("(whole blob)"); + let sent_to = request.sent_to.as_deref().unwrap_or("?"); + let hint = request.download_hint.as_deref().unwrap_or("-"); + let routed = if request.original_host.is_some() { + "yes" + } else { + "no" + }; + eprintln!(" {kind:<6} {range:<24} {routed:<7} {hint:<7} {sent_to}"); + } + + if routed.is_empty() { + eprintln!( + "NOTE: no chunk was rewritten - the account returned no layout hint for this blob" + ); + } else { + for request in &routed { + assert_eq!( + request.original_host.as_deref(), + Some(account_host.as_str()), + "a rewritten chunk did not preserve the account Host header" + ); + } + } + + Ok(()) +} + #[recorded::test] async fn test_upload_blob_content_headers(ctx: TestContext) -> Result<(), Box> { // Recording Setup From 0cee10eecee89c5bcba1e3ff239f62e6d6eca4bf Mon Sep 17 00:00:00 2001 From: vincenttran-msft <101599632+vincenttran-msft@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:24:23 -0700 Subject: [PATCH 2/6] analyze step --- sdk/storage/azure_storage_blob/tests/blob_client.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/sdk/storage/azure_storage_blob/tests/blob_client.rs b/sdk/storage/azure_storage_blob/tests/blob_client.rs index 263fd30f36e..8573ef69f85 100644 --- a/sdk/storage/azure_storage_blob/tests/blob_client.rs +++ b/sdk/storage/azure_storage_blob/tests/blob_client.rs @@ -1909,7 +1909,8 @@ async fn test_download_layout_aware_routing() -> Result<(), Box> { } else { "no" }; - eprintln!(" {kind:<6} {range:<24} {routed:<7} {hint:<7} {sent_to}"); + let status = request.status; + eprintln!(" {kind:<6} {status:<3} {range:<24} {routed:<7} {hint:<7} {sent_to}"); } if routed.is_empty() { From 1a6ae8e1e4003b1ca86113e093adf4b0c1457eb1 Mon Sep 17 00:00:00 2001 From: vincenttran-msft <101599632+vincenttran-msft@users.noreply.github.com> Date: Thu, 10 Sep 2026 12:11:46 -0700 Subject: [PATCH 3/6] nit: Auto =~ Enabled --- .../src/clients/blob_client.rs | 2 +- .../src/models/method_options.rs | 2 +- .../azure_storage_blob/tests/blob_client.rs | 63 ++++++++++++++++++- 3 files changed, 64 insertions(+), 3 deletions(-) diff --git a/sdk/storage/azure_storage_blob/src/clients/blob_client.rs b/sdk/storage/azure_storage_blob/src/clients/blob_client.rs index a8cd54b4239..07d79f2ce61 100644 --- a/sdk/storage/azure_storage_blob/src/clients/blob_client.rs +++ b/sdk/storage/azure_storage_blob/src/clients/blob_client.rs @@ -476,7 +476,7 @@ fn clone_blob_client(client: &GeneratedBlobClient) -> GeneratedBlobClient { fn layout_request_from_options( options: &BlobClientDownloadOptions<'_>, ) -> Option> { - matches!(options.layout_aware_routing, LayoutAwareRouting::Enabled) + (!matches!(options.layout_aware_routing, LayoutAwareRouting::Disabled)) .then(|| layout_options_from_download(options)) } diff --git a/sdk/storage/azure_storage_blob/src/models/method_options.rs b/sdk/storage/azure_storage_blob/src/models/method_options.rs index 19162346b75..855730b6315 100644 --- a/sdk/storage/azure_storage_blob/src/models/method_options.rs +++ b/sdk/storage/azure_storage_blob/src/models/method_options.rs @@ -28,7 +28,7 @@ use crate::models::{ pub enum LayoutAwareRouting { /// The locality-aware routing behavior is determined by the client library and /// may change in future releases. Currently equivalent to - /// [`LayoutAwareRouting::Disabled`]. + /// [`LayoutAwareRouting::Enabled`]. #[default] Auto, diff --git a/sdk/storage/azure_storage_blob/tests/blob_client.rs b/sdk/storage/azure_storage_blob/tests/blob_client.rs index 8573ef69f85..622a37b232e 100644 --- a/sdk/storage/azure_storage_blob/tests/blob_client.rs +++ b/sdk/storage/azure_storage_blob/tests/blob_client.rs @@ -1636,7 +1636,7 @@ async fn test_download_layout_aware_routing_routes_chunks() -> Result<(), Box Result<(), Box Result<(), Box> { + const DATA: [u8; 8] = [10, 11, 12, 13, 14, 15, 16, 17]; + + let request_count = Arc::new(AtomicUsize::new(0)); + let count_capture = request_count.clone(); + let mock_client = Arc::new(MockHttpClient::new(move |request| { + assert_eq!(0, count_capture.fetch_add(1, Ordering::SeqCst)); + assert!(!request + .url() + .query() + .is_some_and(|query| query.contains("comp=layout"))); + assert_eq!( + Some("bytes=0-7"), + request.headers().get_optional_str(&"range".into()) + ); + + let mut headers = Headers::new(); + headers.insert("content-range", "bytes 0-7/8"); + headers.insert("content-length", DATA.len().to_string()); + headers.insert("etag", "\"complete-range-etag\""); + headers.insert("x-ms-download-hint", "layout"); + async move { + Ok(AsyncRawResponse::from_bytes( + StatusCode::PartialContent, + headers, + Bytes::from_static(&DATA), + )) + } + .boxed() + })); + + let blob_client = BlobClient::new( + Url::parse("https://acct.blob.core.windows.net/container/blob")?, + None, + Some(BlobClientOptions { + client_options: ClientOptions { + transport: Some(Transport::new(mock_client)), + ..Default::default() + }, + ..Default::default() + }), + )?; + + let body = blob_client + .download(Some(BlobClientDownloadOptions { + layout_aware_routing: LayoutAwareRouting::Auto, + partition_size: Some(NonZero::new(DATA.len()).unwrap()), + ..Default::default() + })) + .await? + .body + .collect() + .await?; + + assert_eq!(&body[..], &DATA[..]); + assert_eq!(request_count.load(Ordering::SeqCst), 1); + Ok(()) +} + #[tokio::test] async fn test_download_layout_aware_routing_skips_without_hint() -> Result<(), Box> { const DATA: [u8; 8] = [10, 11, 12, 13, 14, 15, 16, 17]; From 44054f45e1fdec4676d1e715376cda5b3a113170 Mon Sep 17 00:00:00 2001 From: vincenttran-msft <101599632+vincenttran-msft@users.noreply.github.com> Date: Fri, 11 Sep 2026 14:37:52 -0700 Subject: [PATCH 4/6] pr comments 0.5 --- .../layout.rs => blob_layout.rs} | 450 ++++++++++++++++- .../src/clients/blob_client.rs | 452 +++--------------- sdk/storage/azure_storage_blob/src/lib.rs | 1 + .../src/models/method_options.rs | 13 +- .../src/partitioned_transfer/mod.rs | 2 - .../azure_storage_blob/tests/blob_client.rs | 136 +++++- 6 files changed, 651 insertions(+), 403 deletions(-) rename sdk/storage/azure_storage_blob/src/{partitioned_transfer/layout.rs => blob_layout.rs} (62%) diff --git a/sdk/storage/azure_storage_blob/src/partitioned_transfer/layout.rs b/sdk/storage/azure_storage_blob/src/blob_layout.rs similarity index 62% rename from sdk/storage/azure_storage_blob/src/partitioned_transfer/layout.rs rename to sdk/storage/azure_storage_blob/src/blob_layout.rs index ca3d9e6e93a..455583cc9be 100644 --- a/sdk/storage/azure_storage_blob/src/partitioned_transfer/layout.rs +++ b/sdk/storage/azure_storage_blob/src/blob_layout.rs @@ -8,10 +8,15 @@ //! flat, ascending list of [`LayoutSegment`]s and resolves the serving endpoint //! for a given byte offset via binary search. -use std::{collections::HashMap, sync::Arc}; +use std::{ + collections::HashMap, + sync::Arc, + time::{Duration, Instant}, +}; use async_trait::async_trait; use azure_core::{ + async_runtime::get_async_runtime, error::ErrorKind, http::{ policies::{Policy, PolicyResult}, @@ -19,7 +24,11 @@ use azure_core::{ }, Error, Result, }; -use futures::StreamExt as _; +use futures::{ + future::{select, Either}, + lock::Mutex, + StreamExt as _, +}; use crate::generated::{ clients::BlobClient, @@ -148,7 +157,7 @@ impl Policy for LayoutRoutingPolicy { } } -/// Rewrites `request`'s URL authority to `endpoint`, preserving the original +/// Rewrites `request`'s URL host and port to `endpoint`, preserving the original /// authority as an explicit `Host` header. /// /// The rewrite is computed on a scratch copy and only committed once fully @@ -176,6 +185,9 @@ fn apply_layout_endpoint(request: &mut Request, endpoint: &str) -> azure_core::R /// Derives the `Host` header value the client would send for `url`, i.e. the host /// plus the port when a non-default port is explicitly present. +/// +/// The value is derived rather than read back off the request: the transport adds +/// `Host` from the URL after the pipeline runs, so no such header exists here. fn original_host_header(url: &Url) -> String { match (url.host_str(), url.port()) { (Some(host), Some(port)) => format!("{host}:{port}"), @@ -274,6 +286,94 @@ fn classify_layout_error(err: Error) -> Result> { } } +const LAYOUT_TTL: Duration = Duration::from_secs(300); +const LAYOUT_REFRESH_BUFFER: Duration = Duration::from_secs(30); +const LAYOUT_REFRESH_BACKOFF: Duration = Duration::from_secs(30); +const LAYOUT_REFRESH_TIMEOUT: azure_core::time::Duration = azure_core::time::Duration::seconds(30); + +/// A blob's layout held for the lifetime of a download, refreshed in the background +/// before it expires so range requests keep routing without blocking on a fetch. +pub(crate) struct LayoutCache { + client: Arc, + layout_options: BlobClientListLayoutOptions<'static>, + state: Arc>, +} + +struct CachedLayout { + layout: Arc, + refresh_at: Instant, + expires_at: Instant, + refreshing: bool, + retry_at: Option, +} + +impl CachedLayout { + fn new(layout: Arc) -> Self { + let now = Instant::now(); + Self { + layout, + refresh_at: now + (LAYOUT_TTL - LAYOUT_REFRESH_BUFFER), + expires_at: now + LAYOUT_TTL, + refreshing: false, + retry_at: None, + } + } +} + +impl LayoutCache { + pub fn new( + client: Arc, + layout_options: BlobClientListLayoutOptions<'static>, + layout: Arc, + ) -> Self { + Self { + client, + layout_options, + state: Arc::new(Mutex::new(CachedLayout::new(layout))), + } + } + + pub async fn current(&self) -> Option> { + let mut state = self.state.lock().await; + let now = Instant::now(); + let backing_off = matches!(state.retry_at, Some(at) if now < at); + if now >= state.refresh_at && !state.refreshing && !backing_off { + state.refreshing = true; + let _refresh = get_async_runtime().spawn(Box::pin(Self::refresh( + Arc::clone(&self.client), + self.layout_options.clone(), + self.state.clone(), + ))); + } + // Routing is per-call, so a stale endpoint survives retries: drop an expired layout rather than risk an unrecoverable range request. + (now < state.expires_at).then(|| state.layout.clone()) + } + + async fn refresh( + client: Arc, + layout_options: BlobClientListLayoutOptions<'static>, + state: Arc>, + ) { + let context = Context::new(); + let fetch = Box::pin(fetch_layout(&client, &context, &layout_options)); + // Bounded so a hung fetch cannot strand `refreshing` and leak this task for the lifetime of the process; a timeout is treated as any other failed refresh. + let result = match select(fetch, get_async_runtime().sleep(LAYOUT_REFRESH_TIMEOUT)).await { + Either::Left((result, _)) => result, + Either::Right(_) => Ok(None), + }; + let mut state = state.lock().await; + match result { + Ok(Some(prefetch)) => { + *state = CachedLayout::new(Arc::new(prefetch.layout)); + } + Ok(None) | Err(_) => { + state.refreshing = false; + state.retry_at = Some(Instant::now() + LAYOUT_REFRESH_BACKOFF); + } + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -290,6 +390,7 @@ mod tests { }; use azure_core_test::http::MockHttpClient; use futures::FutureExt as _; + use std::sync::atomic::{AtomicUsize, Ordering}; fn segment(start: i64, end: i64, endpoint: Option<&str>) -> LayoutSegment { LayoutSegment { @@ -476,6 +577,11 @@ mod tests { parse_endpoint_authority("host.example.net"), Some(("host.example.net".to_owned(), None)) ); + // IPv6 endpoints keep their brackets so the value parses as a host literal. + assert_eq!( + parse_endpoint_authority("[::1]:10000"), + Some(("[::1]".to_owned(), Some(10000))) + ); // Surrounding whitespace is tolerated. assert_eq!( parse_endpoint_authority(" host.example.net:443 "), @@ -545,6 +651,48 @@ mod tests { assert_eq!(host_header(&request), None); } + #[test] + fn apply_layout_endpoint_preserves_path_style_scheme_and_account_segment() { + // Emulator shape: http, non-default port, account as the first path segment. + let mut request = request_to("http://127.0.0.1:10000/devstoreaccount1/container/blob"); + apply_layout_endpoint(&mut request, "node-b.storage.local:20000").unwrap(); + + assert_eq!(request.url().scheme(), "http"); + assert_eq!(request.url().host_str(), Some("node-b.storage.local")); + assert_eq!(request.url().port(), Some(20000)); + assert_eq!(request.url().path(), "/devstoreaccount1/container/blob"); + assert_eq!(host_header(&request).as_deref(), Some("127.0.0.1:10000")); + } + + #[test] + fn apply_layout_endpoint_custom_domain_preserves_host() { + let mut request = request_to("https://blobs.contoso.com/container/blob"); + apply_layout_endpoint(&mut request, "ep.contoso.net:443").unwrap(); + + assert_eq!(request.url().host_str(), Some("ep.contoso.net")); + assert_eq!(host_header(&request).as_deref(), Some("blobs.contoso.com")); + } + + #[test] + fn apply_layout_endpoint_keeps_ipv6_brackets_in_original_host() { + let mut request = request_to("http://[::1]:10000/devstoreaccount1/container/blob"); + apply_layout_endpoint(&mut request, "node-b.storage.local:20000").unwrap(); + + assert_eq!(request.url().host_str(), Some("node-b.storage.local")); + assert_eq!(host_header(&request).as_deref(), Some("[::1]:10000")); + } + + #[test] + fn apply_layout_endpoint_keeps_the_request_scheme() { + // The endpoint's scheme is ignored, so an http client is never silently upgraded. + let mut request = request_to("http://127.0.0.1:10000/devstoreaccount1/container/blob"); + apply_layout_endpoint(&mut request, "https://node-b.storage.local:20000").unwrap(); + + assert_eq!(request.url().scheme(), "http"); + assert_eq!(request.url().host_str(), Some("node-b.storage.local")); + assert_eq!(request.url().port(), Some(20000)); + } + const LAYOUT_SINGLE_PAGE: &[u8] = br#" @@ -822,4 +970,300 @@ mod tests { .await; assert!(result.is_err()); } + + const LAYOUT_V1: &[u8] = br#" + + + + + + + +"#; + + const LAYOUT_V2: &[u8] = br#" + + + + + + + +"#; + + async fn seed_layout( + client: &BlobClient, + layout_options: &BlobClientListLayoutOptions<'static>, + ) -> Layout { + fetch_layout(client, &Context::new(), layout_options) + .await + .unwrap() + .unwrap() + .layout + } + + #[tokio::test] + async fn refresh_success_updates_layout_and_resets_deadlines() { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let mock: Arc = Arc::new(MockHttpClient::new(move |_req| { + let count = counter.fetch_add(1, Ordering::SeqCst); + async move { + let body = if count == 0 { LAYOUT_V1 } else { LAYOUT_V2 }; + Ok(AsyncRawResponse::from_bytes( + StatusCode::Ok, + headers_with_etag("etag-1"), + Bytes::from_static(body), + )) + } + .boxed() + })); + let client = Arc::new(layout_client(mock)); + let layout_options = BlobClientListLayoutOptions::default(); + let state = Arc::new(Mutex::new(CachedLayout { + layout: Arc::new(seed_layout(&client, &layout_options).await), + refresh_at: Instant::now(), + expires_at: Instant::now(), + refreshing: true, + retry_at: None, + })); + + LayoutCache::refresh(client, layout_options, state.clone()).await; + let state = state.lock().await; + assert_eq!( + state.layout.ideal_endpoint(0), + Some("epv2.blob.storage.azure.net:443") + ); + assert!(!state.refreshing); + assert!(state.retry_at.is_none()); + assert!(state.expires_at > Instant::now()); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn current_routes_while_valid_and_suspends_once_expired() { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let mock: Arc = Arc::new(MockHttpClient::new(move |_req| { + counter.fetch_add(1, Ordering::SeqCst); + async move { + Ok(AsyncRawResponse::from_bytes( + StatusCode::Ok, + headers_with_etag("etag-1"), + Bytes::from_static(LAYOUT_V1), + )) + } + .boxed() + })); + let client = Arc::new(layout_client(mock)); + let layout_options = BlobClientListLayoutOptions::default(); + let cache = LayoutCache::new( + Arc::clone(&client), + layout_options.clone(), + Arc::new(seed_layout(&client, &layout_options).await), + ); + + assert_eq!( + cache.current().await.unwrap().ideal_endpoint(0), + Some("epv1.blob.storage.azure.net:443") + ); + assert_eq!(calls.load(Ordering::SeqCst), 1); + + { + let mut state = cache.state.lock().await; + let now = Instant::now(); + state.refresh_at = now - Duration::from_secs(1); + state.expires_at = now - Duration::from_secs(1); + state.retry_at = Some(now + Duration::from_secs(300)); + } + assert!(cache.current().await.is_none()); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn current_refreshes_in_background_without_blocking() { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let mock: Arc = Arc::new(MockHttpClient::new(move |_req| { + let count = counter.fetch_add(1, Ordering::SeqCst); + async move { + let body = if count == 0 { LAYOUT_V1 } else { LAYOUT_V2 }; + Ok(AsyncRawResponse::from_bytes( + StatusCode::Ok, + headers_with_etag("etag-1"), + Bytes::from_static(body), + )) + } + .boxed() + })); + let client = Arc::new(layout_client(mock)); + let layout_options = BlobClientListLayoutOptions::default(); + let cache = LayoutCache::new( + Arc::clone(&client), + layout_options.clone(), + Arc::new(seed_layout(&client, &layout_options).await), + ); + + { + let mut state = cache.state.lock().await; + let now = Instant::now(); + state.refresh_at = now - Duration::from_secs(1); + state.expires_at = now + Duration::from_secs(100); + } + + assert_eq!( + cache.current().await.unwrap().ideal_endpoint(0), + Some("epv1.blob.storage.azure.net:443") + ); + + let mut spins = 0; + loop { + { + let state = cache.state.lock().await; + if !state.refreshing + && state.layout.ideal_endpoint(0) == Some("epv2.blob.storage.azure.net:443") + { + assert!(state.expires_at > Instant::now() + Duration::from_secs(200)); + assert!(state.retry_at.is_none()); + break; + } + } + assert!(spins < 10_000, "background refresh did not complete"); + spins += 1; + tokio::task::yield_now().await; + } + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn refresh_failure_keeps_layout_and_backs_off() { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let mock: Arc = Arc::new(MockHttpClient::new(move |_req| { + let count = counter.fetch_add(1, Ordering::SeqCst); + async move { + if count == 0 { + Ok(AsyncRawResponse::from_bytes( + StatusCode::Ok, + headers_with_etag("etag-1"), + Bytes::from_static(LAYOUT_V1), + )) + } else { + Ok(AsyncRawResponse::from_bytes( + StatusCode::InternalServerError, + Headers::new(), + Bytes::from_static(b""), + )) + } + } + .boxed() + })); + let client = Arc::new(layout_client(mock)); + let layout_options = BlobClientListLayoutOptions::default(); + let expires_at = Instant::now() + Duration::from_secs(100); + let state = Arc::new(Mutex::new(CachedLayout { + layout: Arc::new(seed_layout(&client, &layout_options).await), + refresh_at: Instant::now(), + expires_at, + refreshing: true, + retry_at: None, + })); + + LayoutCache::refresh(client, layout_options, state.clone()).await; + let state = state.lock().await; + assert_eq!( + state.layout.ideal_endpoint(0), + Some("epv1.blob.storage.azure.net:443") + ); + assert!(!state.refreshing); + assert_eq!(state.expires_at, expires_at); + match state.retry_at { + Some(retry_at) => assert!(retry_at > Instant::now()), + None => panic!("expected a backoff to be set"), + } + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn current_resumes_routing_when_a_refresh_recovers_after_expiry() { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = calls.clone(); + let mock: Arc = Arc::new(MockHttpClient::new(move |_req| { + let count = counter.fetch_add(1, Ordering::SeqCst); + async move { + // Seed, then fail one refresh, then recover. + if count == 1 { + return Ok(AsyncRawResponse::from_bytes( + StatusCode::InternalServerError, + Headers::new(), + Bytes::new(), + )); + } + let body = if count == 0 { LAYOUT_V1 } else { LAYOUT_V2 }; + Ok(AsyncRawResponse::from_bytes( + StatusCode::Ok, + headers_with_etag("etag-1"), + Bytes::from_static(body), + )) + } + .boxed() + })); + let client = Arc::new(layout_client(mock)); + let layout_options = BlobClientListLayoutOptions::default(); + let cache = LayoutCache::new( + Arc::clone(&client), + layout_options.clone(), + Arc::new(seed_layout(&client, &layout_options).await), + ); + + { + let mut state = cache.state.lock().await; + let now = Instant::now(); + state.refresh_at = now - Duration::from_secs(1); + state.expires_at = now - Duration::from_secs(1); + } + + // Expired, and the refresh it triggers fails: routing stays suspended. + assert!(cache.current().await.is_none()); + let mut spins = 0; + loop { + { + let state = cache.state.lock().await; + if !state.refreshing && state.retry_at.is_some() { + break; + } + } + assert!(spins < 10_000, "failed refresh did not settle"); + spins += 1; + tokio::task::yield_now().await; + } + assert!(cache.current().await.is_none()); + + // Once the backoff elapses the next refresh succeeds and routing resumes. + { + let mut state = cache.state.lock().await; + state.retry_at = Some(Instant::now() - Duration::from_secs(1)); + } + assert!(cache.current().await.is_none()); + spins = 0; + loop { + { + let state = cache.state.lock().await; + if !state.refreshing + && state.layout.ideal_endpoint(0) == Some("epv2.blob.storage.azure.net:443") + { + break; + } + } + assert!(spins < 10_000, "recovering refresh did not complete"); + spins += 1; + tokio::task::yield_now().await; + } + + assert_eq!( + cache.current().await.unwrap().ideal_endpoint(0), + Some("epv2.blob.storage.azure.net:443") + ); + assert_eq!(calls.load(Ordering::SeqCst), 3); + } } diff --git a/sdk/storage/azure_storage_blob/src/clients/blob_client.rs b/sdk/storage/azure_storage_blob/src/clients/blob_client.rs index 07d79f2ce61..a4e8f935bf0 100644 --- a/sdk/storage/azure_storage_blob/src/clients/blob_client.rs +++ b/sdk/storage/azure_storage_blob/src/clients/blob_client.rs @@ -4,6 +4,7 @@ pub use crate::generated::clients::{BlobClient, BlobClientOptions}; use crate::{ + blob_layout::{fetch_layout, LayoutCache, LayoutEndpoint, LayoutRoutingPolicy}, generated::{ clients::BlobClient as GeneratedBlobClient, models::{BlobClientDownloadInternalOptions, BlobClientListLayoutOptions}, @@ -13,27 +14,23 @@ use crate::{ BlobClientUploadOptions, BlobClientUploadResult, BlobDownloadProperties, HttpRange, LayoutAwareRouting, StorageErrorCode, }, - partitioned_transfer::{self, Layout, LayoutEndpoint, PartitionedDownloadBehavior}, + partitioned_transfer::{self, PartitionedDownloadBehavior}, AppendBlobClient, BlockBlobClient, PageBlobClient, }; use async_trait::async_trait; use azure_core::{ - async_runtime::get_async_runtime, credentials::TokenCredential, error::ErrorKind, http::{ headers::Headers, policies::{auth::BearerTokenAuthorizationPolicy, Policy}, - AsyncRawResponse, Context, Etag, NoFormat, Pipeline, RequestContent, StatusCode, Url, - UrlExt, + AsyncRawResponse, Etag, NoFormat, Pipeline, RequestContent, StatusCode, Url, UrlExt, }, tracing, Bytes, Result, }; -use futures::lock::Mutex; use std::{ ops::Range, sync::{Arc, OnceLock}, - time::{Duration, Instant}, }; impl BlobClient { @@ -80,7 +77,7 @@ impl BlobClient { option_env!("CARGO_PKG_NAME"), option_env!("CARGO_PKG_VERSION"), options.client_options.clone(), - vec![Arc::new(partitioned_transfer::LayoutRoutingPolicy)], + vec![Arc::new(LayoutRoutingPolicy)], per_retry_policies, None, ); @@ -209,10 +206,14 @@ impl BlobClient { version: self.version.clone(), tracer: self.tracer.clone(), }; - let layout_request = layout_request_from_options(&options); let range = options.range.clone(); - let behavior = - BlobClientDownloadBehavior::new(inner_client, options.into(), layout_request); + let layout_aware_routing = options.layout_aware_routing; + let behavior = BlobClientDownloadBehavior::new( + inner_client, + options.into(), + layout_aware_routing, + range.clone(), + ); let response = partitioned_transfer::download(range, parallel, partition_size, Arc::new(behavior)) .await?; @@ -255,10 +256,14 @@ impl BlobClient { version: self.version.clone(), tracer: self.tracer.clone(), }; - let layout_request = layout_request_from_options(&options); let range = options.range.clone(); - let behavior = - BlobClientDownloadBehavior::new(inner_client, options.into(), layout_request); + let layout_aware_routing = options.layout_aware_routing; + let behavior = BlobClientDownloadBehavior::new( + inner_client, + options.into(), + layout_aware_routing, + range.clone(), + ); let (_, headers, len) = partitioned_transfer::download_into( buffer, range, @@ -315,23 +320,46 @@ impl BlobClient { } struct BlobClientDownloadBehavior<'a> { - client: GeneratedBlobClient, + client: Arc, options: BlobClientDownloadInternalOptions<'a>, - layout_request: Option>, - cache: OnceLock>, + layout_aware_routing: LayoutAwareRouting, + /// The caller-requested range, which `options.range` does not retain because it is rewritten for each partition. + requested_range: Option, + layout_cache: OnceLock>, } impl<'a> BlobClientDownloadBehavior<'a> { fn new( client: GeneratedBlobClient, options: BlobClientDownloadInternalOptions<'a>, - layout_request: Option>, + layout_aware_routing: LayoutAwareRouting, + requested_range: Option, ) -> Self { Self { - client, + client: Arc::new(client), options, - layout_request, - cache: OnceLock::new(), + layout_aware_routing, + requested_range, + layout_cache: OnceLock::new(), + } + } + + fn layout_options(&self) -> BlobClientListLayoutOptions<'static> { + BlobClientListLayoutOptions { + encryption_algorithm: self.options.encryption_algorithm, + encryption_key: self.options.encryption_key.clone(), + encryption_key_sha256: self.options.encryption_key_sha256.clone(), + if_match: self.options.if_match.clone(), + if_modified_since: self.options.if_modified_since, + if_none_match: self.options.if_none_match.clone(), + if_tags: self.options.if_tags.clone(), + if_unmodified_since: self.options.if_unmodified_since, + lease_id: self.options.lease_id.clone(), + range: self.requested_range.clone(), + snapshot: self.options.snapshot.clone(), + timeout: self.options.timeout, + version_id: self.options.version_id.clone(), + ..Default::default() } } } @@ -344,9 +372,9 @@ impl PartitionedDownloadBehavior for BlobClientDownloadBehavior<'_> { etag_lock: Option, ) -> Result { let mut opt = self.options.clone(); - if let Some(Some(cache)) = self.cache.get() { + if let Some(Some(cache)) = self.layout_cache.get() { if let Some(range) = &range { - if let Some(layout) = cache.current(&self.client).await { + if let Some(layout) = cache.current().await { if let Some(endpoint) = layout.ideal_endpoint(range.start as i64) { opt.method_options.context = opt .method_options @@ -372,378 +400,28 @@ impl PartitionedDownloadBehavior for BlobClientDownloadBehavior<'_> { } async fn prepare(&self, initial_headers: &Headers, etag_lock: Option<&Etag>) -> Result<()> { - let Some(layout_request) = &self.layout_request else { + if matches!(self.layout_aware_routing, LayoutAwareRouting::Disabled) { return Ok(()); - }; + } if initial_headers.get_optional_str(&"x-ms-download-hint".into()) != Some("layout") { - let _ = self.cache.set(None); + let _ = self.layout_cache.set(None); return Ok(()); } - let mut request = layout_request.clone(); - if request.if_match.is_none() { - request.if_match = etag_lock.cloned(); + let mut layout_options = self.layout_options(); + if layout_options.if_match.is_none() { + layout_options.if_match = etag_lock.cloned(); } let context = self.options.method_options.context.clone(); - let cache = partitioned_transfer::fetch_layout(&self.client, &context, &request) + let cache = fetch_layout(&self.client, &context, &layout_options) .await? - .map(|prefetch| LayoutCache::new(request, Arc::new(prefetch.layout))); - let _ = self.cache.set(cache); + .map(|prefetch| { + LayoutCache::new( + Arc::clone(&self.client), + layout_options, + Arc::new(prefetch.layout), + ) + }); + let _ = self.layout_cache.set(cache); Ok(()) } } - -const LAYOUT_TTL: Duration = Duration::from_secs(300); -const LAYOUT_REFRESH_BUFFER: Duration = Duration::from_secs(30); -const LAYOUT_REFRESH_BACKOFF: Duration = Duration::from_secs(30); - -struct LayoutCache { - request: BlobClientListLayoutOptions<'static>, - state: Arc>, -} - -struct CachedLayout { - layout: Arc, - refresh_at: Instant, - expires_at: Instant, - refreshing: bool, - retry_at: Option, -} - -impl CachedLayout { - fn new(layout: Arc) -> Self { - let now = Instant::now(); - Self { - layout, - refresh_at: now + (LAYOUT_TTL - LAYOUT_REFRESH_BUFFER), - expires_at: now + LAYOUT_TTL, - refreshing: false, - retry_at: None, - } - } -} - -impl LayoutCache { - fn new(request: BlobClientListLayoutOptions<'static>, layout: Arc) -> Self { - Self { - request, - state: Arc::new(Mutex::new(CachedLayout::new(layout))), - } - } - - async fn current(&self, client: &GeneratedBlobClient) -> Option> { - let mut state = self.state.lock().await; - let now = Instant::now(); - let backing_off = matches!(state.retry_at, Some(at) if now < at); - if now >= state.refresh_at && !state.refreshing && !backing_off { - state.refreshing = true; - let _refresh = get_async_runtime().spawn(Box::pin(Self::refresh( - clone_blob_client(client), - self.request.clone(), - self.state.clone(), - ))); - } - (now < state.expires_at).then(|| state.layout.clone()) - } - - async fn refresh( - client: GeneratedBlobClient, - request: BlobClientListLayoutOptions<'static>, - state: Arc>, - ) { - let result = partitioned_transfer::fetch_layout(&client, &Context::new(), &request).await; - let mut state = state.lock().await; - match result { - Ok(Some(prefetch)) => { - *state = CachedLayout::new(Arc::new(prefetch.layout)); - } - Ok(None) | Err(_) => { - state.refreshing = false; - state.retry_at = Some(Instant::now() + LAYOUT_REFRESH_BACKOFF); - } - } - } -} - -fn clone_blob_client(client: &GeneratedBlobClient) -> GeneratedBlobClient { - GeneratedBlobClient { - endpoint: client.endpoint.clone(), - pipeline: client.pipeline.clone(), - version: client.version.clone(), - tracer: client.tracer.clone(), - } -} - -fn layout_request_from_options( - options: &BlobClientDownloadOptions<'_>, -) -> Option> { - (!matches!(options.layout_aware_routing, LayoutAwareRouting::Disabled)) - .then(|| layout_options_from_download(options)) -} - -fn layout_options_from_download( - options: &BlobClientDownloadOptions<'_>, -) -> BlobClientListLayoutOptions<'static> { - BlobClientListLayoutOptions { - encryption_algorithm: options.encryption_algorithm, - encryption_key: options.encryption_key.clone(), - encryption_key_sha256: options.encryption_key_sha256.clone(), - if_match: options.if_match.clone(), - if_modified_since: options.if_modified_since, - if_none_match: options.if_none_match.clone(), - if_tags: options.if_tags.clone(), - if_unmodified_since: options.if_unmodified_since, - lease_id: options.lease_id.clone(), - range: options.range.clone(), - snapshot: options.snapshot.clone(), - timeout: options.timeout, - version_id: options.version_id.clone(), - ..Default::default() - } -} - -#[cfg(test)] -mod tests { - use super::*; - use azure_core::http::{ClientOptions, FixedRetryOptions, HttpClient, RetryOptions, Transport}; - use azure_core_test::http::MockHttpClient; - use futures::FutureExt as _; - use std::sync::atomic::{AtomicUsize, Ordering}; - - const LAYOUT_V1: &[u8] = br#" - - - - - - - -"#; - - const LAYOUT_V2: &[u8] = br#" - - - - - - - -"#; - - fn mock_client(transport: Arc) -> GeneratedBlobClient { - BlobClient::new( - "https://acct.blob.core.windows.net/container/blob" - .parse() - .unwrap(), - None, - Some(BlobClientOptions { - client_options: ClientOptions { - transport: Some(Transport::new(transport)), - retry: RetryOptions::fixed(FixedRetryOptions { - max_retries: 0, - ..Default::default() - }), - ..Default::default() - }, - ..Default::default() - }), - ) - .unwrap() - } - - fn etag_headers() -> Headers { - let mut headers = Headers::new(); - headers.insert("etag", "etag-1".to_owned()); - headers - } - - async fn seed_layout( - client: &GeneratedBlobClient, - request: &BlobClientListLayoutOptions<'static>, - ) -> Layout { - partitioned_transfer::fetch_layout(client, &Context::new(), request) - .await - .unwrap() - .unwrap() - .layout - } - - #[tokio::test] - async fn refresh_success_updates_layout_and_resets_deadlines() { - let calls = Arc::new(AtomicUsize::new(0)); - let counter = calls.clone(); - let mock: Arc = Arc::new(MockHttpClient::new(move |_req| { - let count = counter.fetch_add(1, Ordering::SeqCst); - async move { - let body = if count == 0 { LAYOUT_V1 } else { LAYOUT_V2 }; - Ok(AsyncRawResponse::from_bytes( - StatusCode::Ok, - etag_headers(), - Bytes::from_static(body), - )) - } - .boxed() - })); - let client = mock_client(mock); - let request = BlobClientListLayoutOptions::default(); - let state = Arc::new(Mutex::new(CachedLayout { - layout: Arc::new(seed_layout(&client, &request).await), - refresh_at: Instant::now(), - expires_at: Instant::now(), - refreshing: true, - retry_at: None, - })); - - LayoutCache::refresh(client, request, state.clone()).await; - let state = state.lock().await; - assert_eq!( - state.layout.ideal_endpoint(0), - Some("epv2.blob.storage.azure.net:443") - ); - assert!(!state.refreshing); - assert!(state.retry_at.is_none()); - assert!(state.expires_at > Instant::now()); - assert_eq!(calls.load(Ordering::SeqCst), 2); - } - - #[tokio::test] - async fn current_routes_while_valid_and_suspends_once_expired() { - let calls = Arc::new(AtomicUsize::new(0)); - let counter = calls.clone(); - let mock: Arc = Arc::new(MockHttpClient::new(move |_req| { - counter.fetch_add(1, Ordering::SeqCst); - async move { - Ok(AsyncRawResponse::from_bytes( - StatusCode::Ok, - etag_headers(), - Bytes::from_static(LAYOUT_V1), - )) - } - .boxed() - })); - let client = mock_client(mock); - let request = BlobClientListLayoutOptions::default(); - let cache = LayoutCache::new( - request.clone(), - Arc::new(seed_layout(&client, &request).await), - ); - - assert_eq!( - cache.current(&client).await.unwrap().ideal_endpoint(0), - Some("epv1.blob.storage.azure.net:443") - ); - assert_eq!(calls.load(Ordering::SeqCst), 1); - - { - let mut state = cache.state.lock().await; - let now = Instant::now(); - state.refresh_at = now - Duration::from_secs(1); - state.expires_at = now - Duration::from_secs(1); - state.retry_at = Some(now + Duration::from_secs(300)); - } - assert!(cache.current(&client).await.is_none()); - assert_eq!(calls.load(Ordering::SeqCst), 1); - } - - #[tokio::test] - async fn current_refreshes_in_background_without_blocking() { - let calls = Arc::new(AtomicUsize::new(0)); - let counter = calls.clone(); - let mock: Arc = Arc::new(MockHttpClient::new(move |_req| { - let count = counter.fetch_add(1, Ordering::SeqCst); - async move { - let body = if count == 0 { LAYOUT_V1 } else { LAYOUT_V2 }; - Ok(AsyncRawResponse::from_bytes( - StatusCode::Ok, - etag_headers(), - Bytes::from_static(body), - )) - } - .boxed() - })); - let client = mock_client(mock); - let request = BlobClientListLayoutOptions::default(); - let cache = LayoutCache::new( - request.clone(), - Arc::new(seed_layout(&client, &request).await), - ); - - { - let mut state = cache.state.lock().await; - let now = Instant::now(); - state.refresh_at = now - Duration::from_secs(1); - state.expires_at = now + Duration::from_secs(100); - } - - assert_eq!( - cache.current(&client).await.unwrap().ideal_endpoint(0), - Some("epv1.blob.storage.azure.net:443") - ); - - let mut spins = 0; - loop { - { - let state = cache.state.lock().await; - if !state.refreshing - && state.layout.ideal_endpoint(0) == Some("epv2.blob.storage.azure.net:443") - { - assert!(state.expires_at > Instant::now() + Duration::from_secs(200)); - assert!(state.retry_at.is_none()); - break; - } - } - assert!(spins < 10_000, "background refresh did not complete"); - spins += 1; - tokio::task::yield_now().await; - } - assert_eq!(calls.load(Ordering::SeqCst), 2); - } - - #[tokio::test] - async fn refresh_failure_keeps_layout_and_backs_off() { - let calls = Arc::new(AtomicUsize::new(0)); - let counter = calls.clone(); - let mock: Arc = Arc::new(MockHttpClient::new(move |_req| { - let count = counter.fetch_add(1, Ordering::SeqCst); - async move { - if count == 0 { - Ok(AsyncRawResponse::from_bytes( - StatusCode::Ok, - etag_headers(), - Bytes::from_static(LAYOUT_V1), - )) - } else { - Ok(AsyncRawResponse::from_bytes( - StatusCode::InternalServerError, - Headers::new(), - Bytes::from_static(b""), - )) - } - } - .boxed() - })); - let client = mock_client(mock); - let request = BlobClientListLayoutOptions::default(); - let expires_at = Instant::now() + Duration::from_secs(100); - let state = Arc::new(Mutex::new(CachedLayout { - layout: Arc::new(seed_layout(&client, &request).await), - refresh_at: Instant::now(), - expires_at, - refreshing: true, - retry_at: None, - })); - - LayoutCache::refresh(client, request, state.clone()).await; - let state = state.lock().await; - assert_eq!( - state.layout.ideal_endpoint(0), - Some("epv1.blob.storage.azure.net:443") - ); - assert!(!state.refreshing); - assert_eq!(state.expires_at, expires_at); - match state.retry_at { - Some(retry_at) => assert!(retry_at > Instant::now()), - None => panic!("expected a backoff to be set"), - } - assert_eq!(calls.load(Ordering::SeqCst), 2); - } -} diff --git a/sdk/storage/azure_storage_blob/src/lib.rs b/sdk/storage/azure_storage_blob/src/lib.rs index 1b65251e6af..adcd7f1b7d0 100644 --- a/sdk/storage/azure_storage_blob/src/lib.rs +++ b/sdk/storage/azure_storage_blob/src/lib.rs @@ -9,6 +9,7 @@ #[cfg(feature = "arrow")] mod arrow; +mod blob_layout; pub(crate) mod buffers; pub mod clients; #[allow(unused_imports)] diff --git a/sdk/storage/azure_storage_blob/src/models/method_options.rs b/sdk/storage/azure_storage_blob/src/models/method_options.rs index 855730b6315..8733c6c37dd 100644 --- a/sdk/storage/azure_storage_blob/src/models/method_options.rs +++ b/sdk/storage/azure_storage_blob/src/models/method_options.rs @@ -20,24 +20,19 @@ use crate::models::{ /// requests issued by a download. /// /// This is a performance optimization only - the bytes returned are identical -/// regardless of the mode used. When routing is enabled, the blob's layout is +/// regardless of the mode used. Routing is enabled by default: the blob's layout is /// fetched and each range request is sent to the endpoint that serves it, falling /// back to the client's configured endpoint when no layout is available. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] #[non_exhaustive] pub enum LayoutAwareRouting { - /// The locality-aware routing behavior is determined by the client library and - /// may change in future releases. Currently equivalent to - /// [`LayoutAwareRouting::Enabled`]. - #[default] - Auto, - /// Never route range requests based on the blob's layout. All requests are sent /// to the client's configured endpoint. Disabled, - /// Opt in to locality-aware routing. The blob's layout is fetched and each range + /// Use locality-aware routing. The blob's layout is fetched and each range /// download is routed to the endpoint that serves it. + #[default] Enabled, } @@ -75,7 +70,7 @@ pub struct BlobClientDownloadOptions<'a> { /// requests issued by this download. This is a performance optimization only: /// the bytes returned are identical regardless of the mode. /// - /// Defaults to [`LayoutAwareRouting::Auto`]. + /// Defaults to [`LayoutAwareRouting::Enabled`]. pub layout_aware_routing: LayoutAwareRouting, /// Allows customization of the method call. diff --git a/sdk/storage/azure_storage_blob/src/partitioned_transfer/mod.rs b/sdk/storage/azure_storage_blob/src/partitioned_transfer/mod.rs index f7ee6ce14eb..716bc4ce9ea 100644 --- a/sdk/storage/azure_storage_blob/src/partitioned_transfer/mod.rs +++ b/sdk/storage/azure_storage_blob/src/partitioned_transfer/mod.rs @@ -3,11 +3,9 @@ pub(crate) mod defaults; mod download; -mod layout; mod upload; pub(crate) use download::*; -pub(crate) use layout::{fetch_layout, Layout, LayoutEndpoint, LayoutRoutingPolicy}; pub(crate) use upload::*; use std::{cmp::max, future::Future, num::NonZero, pin::Pin}; diff --git a/sdk/storage/azure_storage_blob/tests/blob_client.rs b/sdk/storage/azure_storage_blob/tests/blob_client.rs index 622a37b232e..b7ab5af48b0 100644 --- a/sdk/storage/azure_storage_blob/tests/blob_client.rs +++ b/sdk/storage/azure_storage_blob/tests/blob_client.rs @@ -1636,7 +1636,7 @@ async fn test_download_layout_aware_routing_routes_chunks() -> Result<(), Box Result<(), Box Result<(), Box> { + const DATA: [u8; 12] = [10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21]; + const LAYOUT: &[u8] = br#" + + + + + + + + + + + +"#; + + struct Seen { + is_layout: bool, + url: Url, + host_header: Option, + range: Option, + } + + let seen = Arc::new(Mutex::new(Vec::::new())); + let seen_capture = seen.clone(); + let mock_client = Arc::new(MockHttpClient::new(move |request| { + let is_layout = request + .url() + .query() + .is_some_and(|query| query.contains("comp=layout")); + let host_header = request + .headers() + .get_optional_str(&"host".into()) + .map(str::to_owned); + let range = request + .headers() + .get_optional_str(&"range".into()) + .map(str::to_owned); + seen_capture.lock().unwrap().push(Seen { + is_layout, + url: request.url().clone(), + host_header, + range: range.clone(), + }); + + let response = if is_layout { + let mut headers = Headers::new(); + headers.insert("etag", "\"path-style-etag\""); + AsyncRawResponse::from_bytes(StatusCode::Ok, headers, Bytes::from_static(LAYOUT)) + } else { + let range = range.expect("data request must carry a range header"); + let (start, end) = parse_bytes_range(&range); + let slice = DATA[start..=end].to_vec(); + let mut headers = Headers::new(); + headers.insert( + "content-range", + format!("bytes {start}-{end}/{}", DATA.len()), + ); + headers.insert("content-length", slice.len().to_string()); + headers.insert("etag", "\"path-style-etag\""); + headers.insert("x-ms-download-hint", "layout"); + AsyncRawResponse::from_bytes(StatusCode::PartialContent, headers, Bytes::from(slice)) + }; + async move { Ok(response) }.boxed() + })); + + let blob_client = BlobClient::new( + Url::parse("http://127.0.0.1:10000/devstoreaccount1/container/blob")?, + None, + Some(BlobClientOptions { + client_options: ClientOptions { + transport: Some(Transport::new(mock_client)), + ..Default::default() + }, + ..Default::default() + }), + )?; + + let body = blob_client + .download(Some(BlobClientDownloadOptions { + layout_aware_routing: LayoutAwareRouting::Enabled, + partition_size: Some(NonZero::new(4).unwrap()), + parallel: Some(NonZero::new(2).unwrap()), + ..Default::default() + })) + .await? + .body + .collect() + .await?; + + assert_eq!(&body[..], &DATA[..]); + let seen = seen.lock().unwrap(); + let layout_requests: Vec<&Seen> = seen.iter().filter(|request| request.is_layout).collect(); + let data_requests: Vec<&Seen> = seen.iter().filter(|request| !request.is_layout).collect(); + + assert_eq!(layout_requests.len(), 1); + assert_eq!(layout_requests[0].url.host_str(), Some("127.0.0.1")); + + assert_eq!(data_requests.len(), 3); + for request in &data_requests { + let range = request.range.as_deref().expect("range header"); + // Rerouting must not disturb the emulator's scheme or its account path segment. + assert_eq!(request.url.scheme(), "http"); + assert_eq!(request.url.path(), "/devstoreaccount1/container/blob"); + match range { + "bytes=0-3" => { + assert_eq!( + request.url.host_str(), + Some("127.0.0.1"), + "the initial chunk should not be routed" + ); + assert_eq!(request.url.port(), Some(10000)); + assert_eq!(request.host_header, None); + } + "bytes=4-7" | "bytes=8-11" => { + let expected = if range == "bytes=4-7" { + "node-b.storage.local" + } else { + "node-c.storage.local" + }; + assert_eq!(request.url.host_str(), Some(expected)); + assert_eq!(request.url.port(), Some(20000)); + assert_eq!(request.host_header.as_deref(), Some("127.0.0.1:10000")); + } + other => panic!("unexpected data range: {other}"), + } + } + Ok(()) +} + #[tokio::test] async fn test_download_layout_aware_routing_skips_layout_for_complete_initial_range( ) -> Result<(), Box> { @@ -1730,7 +1862,7 @@ async fn test_download_layout_aware_routing_skips_layout_for_complete_initial_ra let body = blob_client .download(Some(BlobClientDownloadOptions { - layout_aware_routing: LayoutAwareRouting::Auto, + layout_aware_routing: LayoutAwareRouting::Enabled, partition_size: Some(NonZero::new(DATA.len()).unwrap()), ..Default::default() })) From dc58969013b00852d78b3c5a9b39fba037c19b3e Mon Sep 17 00:00:00 2001 From: vincenttran-msft <101599632+vincenttran-msft@users.noreply.github.com> Date: Mon, 14 Sep 2026 13:36:27 -0700 Subject: [PATCH 5/6] rename to get_layout --- .../azure_storage_blob/src/blob_layout.rs | 40 +- .../src/clients/blob_client.rs | 6 +- .../src/generated/clients/blob_client.rs | 420 +++++++++--------- .../src/generated/models/header_traits.rs | 2 +- .../src/generated/models/method_options.rs | 108 ++--- .../azure_storage_blob/tsp-location.yaml | 2 +- 6 files changed, 290 insertions(+), 288 deletions(-) diff --git a/sdk/storage/azure_storage_blob/src/blob_layout.rs b/sdk/storage/azure_storage_blob/src/blob_layout.rs index 455583cc9be..5160676d4b7 100644 --- a/sdk/storage/azure_storage_blob/src/blob_layout.rs +++ b/sdk/storage/azure_storage_blob/src/blob_layout.rs @@ -32,7 +32,7 @@ use futures::{ use crate::generated::{ clients::BlobClient, - models::{BlobClientListLayoutOptions, BlobLayout}, + models::{BlobClientGetLayoutOptions, BlobLayout}, }; /// A contiguous byte range of a blob and the endpoint that serves it. @@ -242,13 +242,13 @@ pub(crate) struct LayoutPrefetch { pub(crate) async fn fetch_layout( client: &BlobClient, context: &Context<'_>, - options: &BlobClientListLayoutOptions<'_>, + options: &BlobClientGetLayoutOptions<'_>, ) -> Result> { let mut layout = Layout::default(); let mut locked_etag = options.if_match.clone(); let mut options = options.clone(); options.method_options.context = context.clone().into_owned(); - let mut pages = client.list_layout(Some(options))?; + let mut pages = client.get_layout(Some(options))?; while let Some(response) = pages.next().await { let response = match response { @@ -295,7 +295,7 @@ const LAYOUT_REFRESH_TIMEOUT: azure_core::time::Duration = azure_core::time::Dur /// before it expires so range requests keep routing without blocking on a fetch. pub(crate) struct LayoutCache { client: Arc, - layout_options: BlobClientListLayoutOptions<'static>, + layout_options: BlobClientGetLayoutOptions<'static>, state: Arc>, } @@ -323,7 +323,7 @@ impl CachedLayout { impl LayoutCache { pub fn new( client: Arc, - layout_options: BlobClientListLayoutOptions<'static>, + layout_options: BlobClientGetLayoutOptions<'static>, layout: Arc, ) -> Self { Self { @@ -351,7 +351,7 @@ impl LayoutCache { async fn refresh( client: Arc, - layout_options: BlobClientListLayoutOptions<'static>, + layout_options: BlobClientGetLayoutOptions<'static>, state: Arc>, ) { let context = Context::new(); @@ -778,7 +778,7 @@ mod tests { let prefetch = fetch_layout( &client, &Context::new(), - &BlobClientListLayoutOptions::default(), + &BlobClientGetLayoutOptions::default(), ) .await .unwrap() @@ -824,7 +824,7 @@ mod tests { let prefetch = fetch_layout( &client, &Context::new(), - &BlobClientListLayoutOptions { + &BlobClientGetLayoutOptions { if_match: Some(Etag::from("etag-1")), ..Default::default() }, @@ -849,7 +849,7 @@ mod tests { // Once the emitter is fixed, this test should be changed to assert the corrected empty or // optional result generated for a successful 204 response. #[tokio::test] - async fn generated_list_layout_fails_to_deserialize_valid_no_content_response() { + async fn generated_get_layout_fails_to_deserialize_valid_no_content_response() { let mock: Arc = Arc::new(MockHttpClient::new(|_req| { async move { Ok(AsyncRawResponse::from_bytes( @@ -862,7 +862,7 @@ mod tests { })); let client = layout_client(mock); - let mut pages = client.list_layout(None).unwrap(); + let mut pages = client.get_layout(None).unwrap(); let error = pages .next() @@ -893,7 +893,7 @@ mod tests { let result = fetch_layout( &client, &Context::new(), - &BlobClientListLayoutOptions::default(), + &BlobClientGetLayoutOptions::default(), ) .await .unwrap(); @@ -917,7 +917,7 @@ mod tests { let result = fetch_layout( &client, &Context::new(), - &BlobClientListLayoutOptions::default(), + &BlobClientGetLayoutOptions::default(), ) .await .unwrap(); @@ -941,7 +941,7 @@ mod tests { let result = fetch_layout( &client, &Context::new(), - &BlobClientListLayoutOptions::default(), + &BlobClientGetLayoutOptions::default(), ) .await .unwrap(); @@ -965,7 +965,7 @@ mod tests { let result = fetch_layout( &client, &Context::new(), - &BlobClientListLayoutOptions::default(), + &BlobClientGetLayoutOptions::default(), ) .await; assert!(result.is_err()); @@ -993,7 +993,7 @@ mod tests { async fn seed_layout( client: &BlobClient, - layout_options: &BlobClientListLayoutOptions<'static>, + layout_options: &BlobClientGetLayoutOptions<'static>, ) -> Layout { fetch_layout(client, &Context::new(), layout_options) .await @@ -1019,7 +1019,7 @@ mod tests { .boxed() })); let client = Arc::new(layout_client(mock)); - let layout_options = BlobClientListLayoutOptions::default(); + let layout_options = BlobClientGetLayoutOptions::default(); let state = Arc::new(Mutex::new(CachedLayout { layout: Arc::new(seed_layout(&client, &layout_options).await), refresh_at: Instant::now(), @@ -1056,7 +1056,7 @@ mod tests { .boxed() })); let client = Arc::new(layout_client(mock)); - let layout_options = BlobClientListLayoutOptions::default(); + let layout_options = BlobClientGetLayoutOptions::default(); let cache = LayoutCache::new( Arc::clone(&client), layout_options.clone(), @@ -1097,7 +1097,7 @@ mod tests { .boxed() })); let client = Arc::new(layout_client(mock)); - let layout_options = BlobClientListLayoutOptions::default(); + let layout_options = BlobClientGetLayoutOptions::default(); let cache = LayoutCache::new( Arc::clone(&client), layout_options.clone(), @@ -1159,7 +1159,7 @@ mod tests { .boxed() })); let client = Arc::new(layout_client(mock)); - let layout_options = BlobClientListLayoutOptions::default(); + let layout_options = BlobClientGetLayoutOptions::default(); let expires_at = Instant::now() + Duration::from_secs(100); let state = Arc::new(Mutex::new(CachedLayout { layout: Arc::new(seed_layout(&client, &layout_options).await), @@ -1209,7 +1209,7 @@ mod tests { .boxed() })); let client = Arc::new(layout_client(mock)); - let layout_options = BlobClientListLayoutOptions::default(); + let layout_options = BlobClientGetLayoutOptions::default(); let cache = LayoutCache::new( Arc::clone(&client), layout_options.clone(), diff --git a/sdk/storage/azure_storage_blob/src/clients/blob_client.rs b/sdk/storage/azure_storage_blob/src/clients/blob_client.rs index a4e8f935bf0..dd6559b1cd5 100644 --- a/sdk/storage/azure_storage_blob/src/clients/blob_client.rs +++ b/sdk/storage/azure_storage_blob/src/clients/blob_client.rs @@ -7,7 +7,7 @@ use crate::{ blob_layout::{fetch_layout, LayoutCache, LayoutEndpoint, LayoutRoutingPolicy}, generated::{ clients::BlobClient as GeneratedBlobClient, - models::{BlobClientDownloadInternalOptions, BlobClientListLayoutOptions}, + models::{BlobClientDownloadInternalOptions, BlobClientGetLayoutOptions}, }, models::{ BlobClientDownloadIntoResult, BlobClientDownloadOptions, BlobClientDownloadResult, @@ -344,8 +344,8 @@ impl<'a> BlobClientDownloadBehavior<'a> { } } - fn layout_options(&self) -> BlobClientListLayoutOptions<'static> { - BlobClientListLayoutOptions { + fn layout_options(&self) -> BlobClientGetLayoutOptions<'static> { + BlobClientGetLayoutOptions { encryption_algorithm: self.options.encryption_algorithm, encryption_key: self.options.encryption_key.clone(), encryption_key_sha256: self.options.encryption_key_sha256.clone(), diff --git a/sdk/storage/azure_storage_blob/src/generated/clients/blob_client.rs b/sdk/storage/azure_storage_blob/src/generated/clients/blob_client.rs index 312daab9723..3e4e27264e9 100644 --- a/sdk/storage/azure_storage_blob/src/generated/clients/blob_client.rs +++ b/sdk/storage/azure_storage_blob/src/generated/clients/blob_client.rs @@ -3,22 +3,24 @@ // // Code generated by Microsoft (R) Rust Code Generator. DO NOT EDIT. -use crate::generated::models::{ - AccessTier, BlobClientAbortCopyOptions, BlobClientAcquireLeaseOptions, - BlobClientAcquireLeaseResult, BlobClientBreakLeaseOptions, BlobClientBreakLeaseResult, - BlobClientChangeLeaseOptions, BlobClientChangeLeaseResult, BlobClientCreateSnapshotOptions, - BlobClientCreateSnapshotResult, BlobClientDeleteImmutabilityPolicyOptions, - BlobClientDeleteOptions, BlobClientDownloadInternalOptions, BlobClientDownloadInternalResult, - BlobClientGetAccountInfoOptions, BlobClientGetAccountInfoResult, - BlobClientGetPropertiesOptions, BlobClientGetPropertiesResult, BlobClientGetTagsOptions, - BlobClientListLayoutOptions, BlobClientReleaseLeaseOptions, BlobClientReleaseLeaseResult, - BlobClientRenewLeaseOptions, BlobClientRenewLeaseResult, - BlobClientSetImmutabilityPolicyOptions, BlobClientSetLegalHoldOptions, - BlobClientSetMetadataOptions, BlobClientSetPropertiesOptions, BlobClientSetTagsOptions, - BlobClientSetTierOptions, BlobClientStartCopyFromUrlOptions, BlobClientStartCopyFromUrlResult, - BlobClientUndeleteOptions, BlobLayout, BlobTags, +use crate::{ + generated::models::{ + AccessTier, BlobClientAbortCopyOptions, BlobClientAcquireLeaseOptions, + BlobClientAcquireLeaseResult, BlobClientBreakLeaseOptions, BlobClientBreakLeaseResult, + BlobClientChangeLeaseOptions, BlobClientChangeLeaseResult, BlobClientCreateSnapshotOptions, + BlobClientCreateSnapshotResult, BlobClientDeleteImmutabilityPolicyOptions, + BlobClientDeleteOptions, BlobClientDownloadInternalOptions, + BlobClientDownloadInternalResult, BlobClientGetAccountInfoOptions, + BlobClientGetAccountInfoResult, BlobClientGetLayoutOptions, BlobClientGetPropertiesOptions, + BlobClientGetPropertiesResult, BlobClientGetTagsOptions, BlobClientReleaseLeaseOptions, + BlobClientReleaseLeaseResult, BlobClientRenewLeaseOptions, BlobClientRenewLeaseResult, + BlobClientSetImmutabilityPolicyOptions, BlobClientSetLegalHoldOptions, + BlobClientSetMetadataOptions, BlobClientSetPropertiesOptions, BlobClientSetTagsOptions, + BlobClientSetTierOptions, BlobClientStartCopyFromUrlOptions, + BlobClientStartCopyFromUrlResult, BlobClientUndeleteOptions, BlobLayout, BlobTags, + }, + // SessionOptions, }; -// use crate::SessionOptions; use azure_core::{ base64, error::CheckSuccessOptions, @@ -47,7 +49,7 @@ pub struct BlobClient { pub struct BlobClientOptions { /// Allows customization of the client. pub client_options: ClientOptions, - // /// Options for session token authentication. + /// Options for session token authentication. // pub session_options: Option, /// Specifies the version of the operation to use for this request. pub version: String, @@ -828,6 +830,199 @@ impl BlobClient { Ok(rsp.into()) } + /// The Get Blob Layout operation returns all user-defined metadata, standard HTTP properties, and system properties for the + /// blob. In addition, it may optionally return the layout of the blob. + /// + /// # Arguments + /// + /// * `options` - Optional parameters for the request. + /// + /// ## Response Headers + /// + /// The returned [`Response`](azure_core::http::Response) implements the [`BlobLayoutHeaders`] trait, which provides + /// access to response headers. For example: + /// + /// ```no_run + /// use azure_core::{Result, http::{Response, XmlFormat}}; + /// use azure_storage_blob::models::{BlobLayout, BlobLayoutHeaders}; + /// async fn example() -> Result<()> { + /// let response: Response = unimplemented!(); + /// // Access response headers + /// if let Some(cache_control) = response.cache_control()? { + /// println!("cache-control: {:?}", cache_control); + /// } + /// if let Some(content_disposition) = response.content_disposition()? { + /// println!("content-disposition: {:?}", content_disposition); + /// } + /// if let Some(content_encoding) = response.content_encoding()? { + /// println!("content-encoding: {:?}", content_encoding); + /// } + /// Ok(()) + /// } + /// ``` + /// + /// ### Available headers + /// * [`cache_control`()](crate::generated::models::BlobLayoutHeaders::cache_control) - cache-control + /// * [`content_disposition`()](crate::generated::models::BlobLayoutHeaders::content_disposition) - content-disposition + /// * [`content_encoding`()](crate::generated::models::BlobLayoutHeaders::content_encoding) - content-encoding + /// * [`content_language`()](crate::generated::models::BlobLayoutHeaders::content_language) - content-language + /// * [`content_length`()](crate::generated::models::BlobLayoutHeaders::content_length) - content-length + /// * [`content_md5`()](crate::generated::models::BlobLayoutHeaders::content_md5) - content-md5 + /// * [`etag`()](crate::generated::models::BlobLayoutHeaders::etag) - etag + /// * [`last_modified`()](crate::generated::models::BlobLayoutHeaders::last_modified) - last-modified + /// * [`access_tier`()](crate::generated::models::BlobLayoutHeaders::access_tier) - x-ms-access-tier + /// * [`access_tier_change_time`()](crate::generated::models::BlobLayoutHeaders::access_tier_change_time) - x-ms-access-tier-change-time + /// * [`access_tier_inferred`()](crate::generated::models::BlobLayoutHeaders::access_tier_inferred) - x-ms-access-tier-inferred + /// * [`archive_status`()](crate::generated::models::BlobLayoutHeaders::archive_status) - x-ms-archive-status + /// * [`blob_committed_block_count`()](crate::generated::models::BlobLayoutHeaders::blob_committed_block_count) - x-ms-blob-committed-block-count + /// * [`blob_content_encoding`()](crate::generated::models::BlobLayoutHeaders::blob_content_encoding) - x-ms-blob-content-encoding + /// * [`blob_content_length`()](crate::generated::models::BlobLayoutHeaders::blob_content_length) - x-ms-blob-content-length + /// * [`blob_content_md5`()](crate::generated::models::BlobLayoutHeaders::blob_content_md5) - x-ms-blob-content-md5 + /// * [`blob_content_type`()](crate::generated::models::BlobLayoutHeaders::blob_content_type) - x-ms-blob-content-type + /// * [`blob_creation_time`()](crate::generated::models::BlobLayoutHeaders::blob_creation_time) - x-ms-blob-creation-time + /// * [`is_sealed`()](crate::generated::models::BlobLayoutHeaders::is_sealed) - x-ms-blob-sealed + /// * [`blob_sequence_number`()](crate::generated::models::BlobLayoutHeaders::blob_sequence_number) - x-ms-blob-sequence-number + /// * [`blob_type`()](crate::generated::models::BlobLayoutHeaders::blob_type) - x-ms-blob-type + /// * [`copy_completion_time`()](crate::generated::models::BlobLayoutHeaders::copy_completion_time) - x-ms-copy-completion-time + /// * [`destination_snapshot`()](crate::generated::models::BlobLayoutHeaders::destination_snapshot) - x-ms-copy-destination-snapshot + /// * [`copy_id`()](crate::generated::models::BlobLayoutHeaders::copy_id) - x-ms-copy-id + /// * [`copy_progress`()](crate::generated::models::BlobLayoutHeaders::copy_progress) - x-ms-copy-progress + /// * [`copy_source`()](crate::generated::models::BlobLayoutHeaders::copy_source) - x-ms-copy-source + /// * [`copy_status`()](crate::generated::models::BlobLayoutHeaders::copy_status) - x-ms-copy-status + /// * [`copy_status_description`()](crate::generated::models::BlobLayoutHeaders::copy_status_description) - x-ms-copy-status-description + /// * [`creation_time`()](crate::generated::models::BlobLayoutHeaders::creation_time) - x-ms-creation-time + /// * [`encryption_key_sha256`()](crate::generated::models::BlobLayoutHeaders::encryption_key_sha256) - x-ms-encryption-key-sha256 + /// * [`encryption_scope`()](crate::generated::models::BlobLayoutHeaders::encryption_scope) - x-ms-encryption-scope + /// * [`expires_on`()](crate::generated::models::BlobLayoutHeaders::expires_on) - x-ms-expiry-time + /// * [`immutability_policy_mode`()](crate::generated::models::BlobLayoutHeaders::immutability_policy_mode) - x-ms-immutability-policy-mode + /// * [`immutability_policy_expires_on`()](crate::generated::models::BlobLayoutHeaders::immutability_policy_expires_on) - x-ms-immutability-policy-until-date + /// * [`is_incremental_copy`()](crate::generated::models::BlobLayoutHeaders::is_incremental_copy) - x-ms-incremental-copy + /// * [`is_current_version`()](crate::generated::models::BlobLayoutHeaders::is_current_version) - x-ms-is-current-version + /// * [`last_accessed`()](crate::generated::models::BlobLayoutHeaders::last_accessed) - x-ms-last-access-time + /// * [`duration`()](crate::generated::models::BlobLayoutHeaders::duration) - x-ms-lease-duration + /// * [`lease_state`()](crate::generated::models::BlobLayoutHeaders::lease_state) - x-ms-lease-state + /// * [`lease_status`()](crate::generated::models::BlobLayoutHeaders::lease_status) - x-ms-lease-status + /// * [`legal_hold`()](crate::generated::models::BlobLayoutHeaders::legal_hold) - x-ms-legal-hold + /// * [`metadata`()](crate::generated::models::BlobLayoutHeaders::metadata) - x-ms-meta + /// * [`object_replication_rules`()](crate::generated::models::BlobLayoutHeaders::object_replication_rules) - x-ms-or + /// * [`object_replication_policy_id`()](crate::generated::models::BlobLayoutHeaders::object_replication_policy_id) - x-ms-or-policy-id + /// * [`rehydrate_priority`()](crate::generated::models::BlobLayoutHeaders::rehydrate_priority) - x-ms-rehydrate-priority + /// * [`is_server_encrypted`()](crate::generated::models::BlobLayoutHeaders::is_server_encrypted) - x-ms-server-encrypted + /// * [`smart_access_tier`()](crate::generated::models::BlobLayoutHeaders::smart_access_tier) - x-ms-smart-access-tier + /// * [`tag_count`()](crate::generated::models::BlobLayoutHeaders::tag_count) - x-ms-tag-count + /// * [`version_id`()](crate::generated::models::BlobLayoutHeaders::version_id) - x-ms-version-id + /// + /// [`BlobLayoutHeaders`]: crate::generated::models::BlobLayoutHeaders + #[tracing::function("Storage.Blob.BlobClient.getLayout")] + pub fn get_layout( + &self, + options: Option>, + ) -> Result>> { + let options = options.unwrap_or_default().into_owned(); + let pipeline = self.pipeline.clone(); + let mut first_url = self.endpoint.clone(); + let mut query_builder = first_url.query_builder(); + query_builder.append_pair("comp", "layout"); + if let Some(marker) = options.marker.as_ref() { + query_builder.set_pair("marker", marker); + } + if let Some(maxresults) = options.maxresults { + query_builder.set_pair("maxresults", maxresults.to_string()); + } + if let Some(snapshot) = options.snapshot.as_ref() { + query_builder.set_pair("snapshot", snapshot); + } + if let Some(timeout) = options.timeout { + query_builder.set_pair("timeout", timeout.to_string()); + } + if let Some(version_id) = options.version_id.as_ref() { + query_builder.set_pair("versionid", version_id); + } + query_builder.build(); + #[derive(serde::Deserialize)] + struct BlobClientGetLayoutPage { + #[serde(rename = "NextMarker")] + next_marker: Option, + } + + let version = self.version.clone(); + Ok(PageIterator::new( + move |marker: PagerState, pager_options| { + let mut url = first_url.clone(); + if let PagerState::More(marker) = marker { + let mut query_builder = url.query_builder(); + query_builder.set_pair("marker", marker.as_ref()); + query_builder.build(); + } + let mut request = Request::new(url, Method::Get); + request.insert_header("accept", "application/xml"); + if let Some(if_match) = options.if_match.as_ref() { + request.insert_header("if-match", if_match.to_string()); + } + if let Some(if_modified_since) = options.if_modified_since { + request.insert_header("if-modified-since", to_rfc7231(&if_modified_since)); + } + if let Some(if_none_match) = options.if_none_match.as_ref() { + request.insert_header("if-none-match", if_none_match.to_string()); + } + if let Some(if_unmodified_since) = options.if_unmodified_since { + request.insert_header("if-unmodified-since", to_rfc7231(&if_unmodified_since)); + } + if let Some(range) = options.range.as_ref() { + request.insert_header("range", range.to_string()); + } + if let Some(encryption_algorithm) = options.encryption_algorithm.as_ref() { + request.insert_header( + "x-ms-encryption-algorithm", + encryption_algorithm.to_string(), + ); + } + if let Some(encryption_key) = options.encryption_key.as_ref() { + request.insert_header("x-ms-encryption-key", encryption_key); + } + if let Some(encryption_key_sha256) = options.encryption_key_sha256.as_ref() { + request.insert_header("x-ms-encryption-key-sha256", encryption_key_sha256); + } + if let Some(if_tags) = options.if_tags.as_ref() { + request.insert_header("x-ms-if-tags", if_tags); + } + if let Some(lease_id) = options.lease_id.as_ref() { + request.insert_header("x-ms-lease-id", lease_id); + } + request.insert_header("x-ms-version", &version); + let pipeline = pipeline.clone(); + Box::pin(async move { + let rsp = pipeline + .send( + &pager_options.context, + &mut request, + Some(PipelineSendOptions { + check_success: CheckSuccessOptions { + success_codes: &[200, 204], + }, + ..Default::default() + }), + ) + .await?; + let (status, headers, body) = rsp.deconstruct(); + // TODO: [Emitter Fix Needed] This pageable operation accepts HTTP 204 as + // success, but the generated code still attempts to deserialize its empty + // body as XML, so a valid 204 fails before returning `PagerResult::Done`. + let res: BlobClientGetLayoutPage = xml::from_xml(&body)?; + let rsp = RawResponse::from_bytes(status, headers, body).into(); + Ok(match res.next_marker { + Some(next_marker) if !next_marker.is_empty() => PagerResult::More { + response: rsp, + continuation: PagerContinuation::Token(next_marker), + }, + _ => PagerResult::Done { response: rsp }, + }) + }) + }, + Some(options.method_options), + )) + } + /// Returns all user-defined metadata, standard HTTP properties, and system properties for the specified blob. It does not /// return the content of the blob. /// @@ -1042,199 +1237,6 @@ impl BlobClient { Ok(rsp.into()) } - /// The Get Blob Layout operation returns all user-defined metadata, standard HTTP properties, and system properties for the - /// blob. In addition, it may optionally return the layout of the blob. - /// - /// # Arguments - /// - /// * `options` - Optional parameters for the request. - /// - /// ## Response Headers - /// - /// The returned [`Response`](azure_core::http::Response) implements the [`BlobLayoutHeaders`] trait, which provides - /// access to response headers. For example: - /// - /// ```no_run - /// use azure_core::{Result, http::{Response, XmlFormat}}; - /// use azure_storage_blob::models::{BlobLayout, BlobLayoutHeaders}; - /// async fn example() -> Result<()> { - /// let response: Response = unimplemented!(); - /// // Access response headers - /// if let Some(cache_control) = response.cache_control()? { - /// println!("cache-control: {:?}", cache_control); - /// } - /// if let Some(content_disposition) = response.content_disposition()? { - /// println!("content-disposition: {:?}", content_disposition); - /// } - /// if let Some(content_encoding) = response.content_encoding()? { - /// println!("content-encoding: {:?}", content_encoding); - /// } - /// Ok(()) - /// } - /// ``` - /// - /// ### Available headers - /// * [`cache_control`()](crate::generated::models::BlobLayoutHeaders::cache_control) - cache-control - /// * [`content_disposition`()](crate::generated::models::BlobLayoutHeaders::content_disposition) - content-disposition - /// * [`content_encoding`()](crate::generated::models::BlobLayoutHeaders::content_encoding) - content-encoding - /// * [`content_language`()](crate::generated::models::BlobLayoutHeaders::content_language) - content-language - /// * [`content_length`()](crate::generated::models::BlobLayoutHeaders::content_length) - content-length - /// * [`content_md5`()](crate::generated::models::BlobLayoutHeaders::content_md5) - content-md5 - /// * [`etag`()](crate::generated::models::BlobLayoutHeaders::etag) - etag - /// * [`last_modified`()](crate::generated::models::BlobLayoutHeaders::last_modified) - last-modified - /// * [`access_tier`()](crate::generated::models::BlobLayoutHeaders::access_tier) - x-ms-access-tier - /// * [`access_tier_change_time`()](crate::generated::models::BlobLayoutHeaders::access_tier_change_time) - x-ms-access-tier-change-time - /// * [`access_tier_inferred`()](crate::generated::models::BlobLayoutHeaders::access_tier_inferred) - x-ms-access-tier-inferred - /// * [`archive_status`()](crate::generated::models::BlobLayoutHeaders::archive_status) - x-ms-archive-status - /// * [`blob_committed_block_count`()](crate::generated::models::BlobLayoutHeaders::blob_committed_block_count) - x-ms-blob-committed-block-count - /// * [`blob_content_encoding`()](crate::generated::models::BlobLayoutHeaders::blob_content_encoding) - x-ms-blob-content-encoding - /// * [`blob_content_length`()](crate::generated::models::BlobLayoutHeaders::blob_content_length) - x-ms-blob-content-length - /// * [`blob_content_md5`()](crate::generated::models::BlobLayoutHeaders::blob_content_md5) - x-ms-blob-content-md5 - /// * [`blob_content_type`()](crate::generated::models::BlobLayoutHeaders::blob_content_type) - x-ms-blob-content-type - /// * [`blob_creation_time`()](crate::generated::models::BlobLayoutHeaders::blob_creation_time) - x-ms-blob-creation-time - /// * [`is_sealed`()](crate::generated::models::BlobLayoutHeaders::is_sealed) - x-ms-blob-sealed - /// * [`blob_sequence_number`()](crate::generated::models::BlobLayoutHeaders::blob_sequence_number) - x-ms-blob-sequence-number - /// * [`blob_type`()](crate::generated::models::BlobLayoutHeaders::blob_type) - x-ms-blob-type - /// * [`copy_completion_time`()](crate::generated::models::BlobLayoutHeaders::copy_completion_time) - x-ms-copy-completion-time - /// * [`destination_snapshot`()](crate::generated::models::BlobLayoutHeaders::destination_snapshot) - x-ms-copy-destination-snapshot - /// * [`copy_id`()](crate::generated::models::BlobLayoutHeaders::copy_id) - x-ms-copy-id - /// * [`copy_progress`()](crate::generated::models::BlobLayoutHeaders::copy_progress) - x-ms-copy-progress - /// * [`copy_source`()](crate::generated::models::BlobLayoutHeaders::copy_source) - x-ms-copy-source - /// * [`copy_status`()](crate::generated::models::BlobLayoutHeaders::copy_status) - x-ms-copy-status - /// * [`copy_status_description`()](crate::generated::models::BlobLayoutHeaders::copy_status_description) - x-ms-copy-status-description - /// * [`creation_time`()](crate::generated::models::BlobLayoutHeaders::creation_time) - x-ms-creation-time - /// * [`encryption_key_sha256`()](crate::generated::models::BlobLayoutHeaders::encryption_key_sha256) - x-ms-encryption-key-sha256 - /// * [`encryption_scope`()](crate::generated::models::BlobLayoutHeaders::encryption_scope) - x-ms-encryption-scope - /// * [`expires_on`()](crate::generated::models::BlobLayoutHeaders::expires_on) - x-ms-expiry-time - /// * [`immutability_policy_mode`()](crate::generated::models::BlobLayoutHeaders::immutability_policy_mode) - x-ms-immutability-policy-mode - /// * [`immutability_policy_expires_on`()](crate::generated::models::BlobLayoutHeaders::immutability_policy_expires_on) - x-ms-immutability-policy-until-date - /// * [`is_incremental_copy`()](crate::generated::models::BlobLayoutHeaders::is_incremental_copy) - x-ms-incremental-copy - /// * [`is_current_version`()](crate::generated::models::BlobLayoutHeaders::is_current_version) - x-ms-is-current-version - /// * [`last_accessed`()](crate::generated::models::BlobLayoutHeaders::last_accessed) - x-ms-last-access-time - /// * [`duration`()](crate::generated::models::BlobLayoutHeaders::duration) - x-ms-lease-duration - /// * [`lease_state`()](crate::generated::models::BlobLayoutHeaders::lease_state) - x-ms-lease-state - /// * [`lease_status`()](crate::generated::models::BlobLayoutHeaders::lease_status) - x-ms-lease-status - /// * [`legal_hold`()](crate::generated::models::BlobLayoutHeaders::legal_hold) - x-ms-legal-hold - /// * [`metadata`()](crate::generated::models::BlobLayoutHeaders::metadata) - x-ms-meta - /// * [`object_replication_rules`()](crate::generated::models::BlobLayoutHeaders::object_replication_rules) - x-ms-or - /// * [`object_replication_policy_id`()](crate::generated::models::BlobLayoutHeaders::object_replication_policy_id) - x-ms-or-policy-id - /// * [`rehydrate_priority`()](crate::generated::models::BlobLayoutHeaders::rehydrate_priority) - x-ms-rehydrate-priority - /// * [`is_server_encrypted`()](crate::generated::models::BlobLayoutHeaders::is_server_encrypted) - x-ms-server-encrypted - /// * [`smart_access_tier`()](crate::generated::models::BlobLayoutHeaders::smart_access_tier) - x-ms-smart-access-tier - /// * [`tag_count`()](crate::generated::models::BlobLayoutHeaders::tag_count) - x-ms-tag-count - /// * [`version_id`()](crate::generated::models::BlobLayoutHeaders::version_id) - x-ms-version-id - /// - /// [`BlobLayoutHeaders`]: crate::generated::models::BlobLayoutHeaders - #[tracing::function("Storage.Blob.BlobClient.getLayout")] - pub fn list_layout( - &self, - options: Option>, - ) -> Result>> { - let options = options.unwrap_or_default().into_owned(); - let pipeline = self.pipeline.clone(); - let mut first_url = self.endpoint.clone(); - let mut query_builder = first_url.query_builder(); - query_builder.append_pair("comp", "layout"); - if let Some(marker) = options.marker.as_ref() { - query_builder.set_pair("marker", marker); - } - if let Some(maxresults) = options.maxresults { - query_builder.set_pair("maxresults", maxresults.to_string()); - } - if let Some(snapshot) = options.snapshot.as_ref() { - query_builder.set_pair("snapshot", snapshot); - } - if let Some(timeout) = options.timeout { - query_builder.set_pair("timeout", timeout.to_string()); - } - if let Some(version_id) = options.version_id.as_ref() { - query_builder.set_pair("versionid", version_id); - } - query_builder.build(); - #[derive(serde::Deserialize)] - struct BlobClientListLayoutPage { - #[serde(rename = "NextMarker")] - next_marker: Option, - } - - let version = self.version.clone(); - Ok(PageIterator::new( - move |marker: PagerState, pager_options| { - let mut url = first_url.clone(); - if let PagerState::More(marker) = marker { - let mut query_builder = url.query_builder(); - query_builder.set_pair("marker", marker.as_ref()); - query_builder.build(); - } - let mut request = Request::new(url, Method::Get); - request.insert_header("accept", "application/xml"); - if let Some(if_match) = options.if_match.as_ref() { - request.insert_header("if-match", if_match.to_string()); - } - if let Some(if_modified_since) = options.if_modified_since { - request.insert_header("if-modified-since", to_rfc7231(&if_modified_since)); - } - if let Some(if_none_match) = options.if_none_match.as_ref() { - request.insert_header("if-none-match", if_none_match.to_string()); - } - if let Some(if_unmodified_since) = options.if_unmodified_since { - request.insert_header("if-unmodified-since", to_rfc7231(&if_unmodified_since)); - } - if let Some(range) = options.range.as_ref() { - request.insert_header("range", range.to_string()); - } - if let Some(encryption_algorithm) = options.encryption_algorithm.as_ref() { - request.insert_header( - "x-ms-encryption-algorithm", - encryption_algorithm.to_string(), - ); - } - if let Some(encryption_key) = options.encryption_key.as_ref() { - request.insert_header("x-ms-encryption-key", encryption_key); - } - if let Some(encryption_key_sha256) = options.encryption_key_sha256.as_ref() { - request.insert_header("x-ms-encryption-key-sha256", encryption_key_sha256); - } - if let Some(if_tags) = options.if_tags.as_ref() { - request.insert_header("x-ms-if-tags", if_tags); - } - if let Some(lease_id) = options.lease_id.as_ref() { - request.insert_header("x-ms-lease-id", lease_id); - } - request.insert_header("x-ms-version", &version); - let pipeline = pipeline.clone(); - Box::pin(async move { - let rsp = pipeline - .send( - &pager_options.context, - &mut request, - Some(PipelineSendOptions { - check_success: CheckSuccessOptions { - success_codes: &[200, 204], - }, - ..Default::default() - }), - ) - .await?; - let (status, headers, body) = rsp.deconstruct(); - // TODO: [Emitter Fix Needed] This pageable operation accepts HTTP 204 as - // success, but the generated code still attempts to deserialize its empty - // body as XML, so a valid 204 fails before returning `PagerResult::Done`. - let res: BlobClientListLayoutPage = xml::from_xml(&body)?; - let rsp = RawResponse::from_bytes(status, headers, body).into(); - Ok(match res.next_marker { - Some(next_marker) if !next_marker.is_empty() => PagerResult::More { - response: rsp, - continuation: PagerContinuation::Token(next_marker), - }, - _ => PagerResult::Done { response: rsp }, - }) - }) - }, - Some(options.method_options), - )) - } - /// Frees the lease if it's no longer needed, so that another client can immediately acquire a lease against the blob. /// /// # Arguments diff --git a/sdk/storage/azure_storage_blob/src/generated/models/header_traits.rs b/sdk/storage/azure_storage_blob/src/generated/models/header_traits.rs index dbfe8fa65be..a28ea3802af 100644 --- a/sdk/storage/azure_storage_blob/src/generated/models/header_traits.rs +++ b/sdk/storage/azure_storage_blob/src/generated/models/header_traits.rs @@ -1887,7 +1887,7 @@ impl BlobContainerClientRenewLeaseResultHeaders } } -/// Provides access to typed response headers for `BlobClient::list_layout()` +/// Provides access to typed response headers for `BlobClient::get_layout()` /// /// # Examples /// diff --git a/sdk/storage/azure_storage_blob/src/generated/models/method_options.rs b/sdk/storage/azure_storage_blob/src/generated/models/method_options.rs index bce1658c99a..09b49de8581 100644 --- a/sdk/storage/azure_storage_blob/src/generated/models/method_options.rs +++ b/sdk/storage/azure_storage_blob/src/generated/models/method_options.rs @@ -529,9 +529,9 @@ pub struct BlobClientGetAccountInfoOptions<'a> { pub timeout: Option, } -/// Options to be passed to `BlobClient::get_properties()` +/// Options to be passed to `BlobClient::get_layout()` #[derive(Clone, Default, SafeDebug)] -pub struct BlobClientGetPropertiesOptions<'a> { +pub struct BlobClientGetLayoutOptions<'a> { /// The algorithm used to produce the encryption key hash. Must be provided if the encryption key is provided. pub encryption_algorithm: Option, @@ -559,8 +559,18 @@ pub struct BlobClientGetPropertiesOptions<'a> { /// If specified, the operation only succeeds if the resource's lease is active and matches this ID. pub lease_id: Option, + /// An opaque string value that identifies the portion of the result set to return with this operation. + pub marker: Option, + + /// Specifies the maximum number of resources to return. If the request does not specify maxresults, or specifies a value + /// greater than 5000, the server will return up to 5000 items. + pub maxresults: Option, + /// Allows customization of the method call. - pub method_options: ClientMethodOptions<'a>, + pub method_options: PagerOptions<'a>, + + /// Specifies the range of the blob to operate on. + pub range: Option, /// Specifies the snapshot of the blob. pub snapshot: Option, @@ -572,9 +582,45 @@ pub struct BlobClientGetPropertiesOptions<'a> { pub version_id: Option, } -/// Options to be passed to `BlobClient::get_tags()` +impl BlobClientGetLayoutOptions<'_> { + /// Transforms this [`BlobClientGetLayoutOptions`] into a new `BlobClientGetLayoutOptions` that owns the underlying data, cloning it if necessary. + pub fn into_owned(self) -> BlobClientGetLayoutOptions<'static> { + BlobClientGetLayoutOptions { + encryption_algorithm: self.encryption_algorithm, + encryption_key: self.encryption_key, + encryption_key_sha256: self.encryption_key_sha256, + if_match: self.if_match, + if_modified_since: self.if_modified_since, + if_none_match: self.if_none_match, + if_tags: self.if_tags, + if_unmodified_since: self.if_unmodified_since, + lease_id: self.lease_id, + marker: self.marker, + maxresults: self.maxresults, + method_options: PagerOptions { + context: self.method_options.context.into_owned(), + ..self.method_options + }, + range: self.range, + snapshot: self.snapshot, + timeout: self.timeout, + version_id: self.version_id, + } + } +} + +/// Options to be passed to `BlobClient::get_properties()` #[derive(Clone, Default, SafeDebug)] -pub struct BlobClientGetTagsOptions<'a> { +pub struct BlobClientGetPropertiesOptions<'a> { + /// The algorithm used to produce the encryption key hash. Must be provided if the encryption key is provided. + pub encryption_algorithm: Option, + + /// Specifies the encryption key to use to encrypt the data provided in the request. + pub encryption_key: Option, + + /// The SHA-256 hash of the provided encryption key. Must be provided if the encryption key is provided. + pub encryption_key_sha256: Option, + /// Specify this value to operate only on a blob with a matching Etag value. pub if_match: Option, @@ -606,18 +652,9 @@ pub struct BlobClientGetTagsOptions<'a> { pub version_id: Option, } -/// Options to be passed to `BlobClient::list_layout()` +/// Options to be passed to `BlobClient::get_tags()` #[derive(Clone, Default, SafeDebug)] -pub struct BlobClientListLayoutOptions<'a> { - /// The algorithm used to produce the encryption key hash. Must be provided if the encryption key is provided. - pub encryption_algorithm: Option, - - /// Specifies the encryption key to use to encrypt the data provided in the request. - pub encryption_key: Option, - - /// The SHA-256 hash of the provided encryption key. Must be provided if the encryption key is provided. - pub encryption_key_sha256: Option, - +pub struct BlobClientGetTagsOptions<'a> { /// Specify this value to operate only on a blob with a matching Etag value. pub if_match: Option, @@ -636,18 +673,8 @@ pub struct BlobClientListLayoutOptions<'a> { /// If specified, the operation only succeeds if the resource's lease is active and matches this ID. pub lease_id: Option, - /// An opaque string value that identifies the portion of the result set to return with this operation. - pub marker: Option, - - /// Specifies the maximum number of resources to return. If the request does not specify maxresults, or specifies a value - /// greater than 5000, the server will return up to 5000 items. - pub maxresults: Option, - /// Allows customization of the method call. - pub method_options: PagerOptions<'a>, - - /// Specifies the range of the blob to operate on. - pub range: Option, + pub method_options: ClientMethodOptions<'a>, /// Specifies the snapshot of the blob. pub snapshot: Option, @@ -659,33 +686,6 @@ pub struct BlobClientListLayoutOptions<'a> { pub version_id: Option, } -impl BlobClientListLayoutOptions<'_> { - /// Transforms this [`BlobClientListLayoutOptions`] into a new `BlobClientListLayoutOptions` that owns the underlying data, cloning it if necessary. - pub fn into_owned(self) -> BlobClientListLayoutOptions<'static> { - BlobClientListLayoutOptions { - encryption_algorithm: self.encryption_algorithm, - encryption_key: self.encryption_key, - encryption_key_sha256: self.encryption_key_sha256, - if_match: self.if_match, - if_modified_since: self.if_modified_since, - if_none_match: self.if_none_match, - if_tags: self.if_tags, - if_unmodified_since: self.if_unmodified_since, - lease_id: self.lease_id, - marker: self.marker, - maxresults: self.maxresults, - method_options: PagerOptions { - context: self.method_options.context.into_owned(), - ..self.method_options - }, - range: self.range, - snapshot: self.snapshot, - timeout: self.timeout, - version_id: self.version_id, - } - } -} - /// Options to be passed to `BlobClient::release_lease()` #[derive(Clone, Default, SafeDebug)] pub struct BlobClientReleaseLeaseOptions<'a> { diff --git a/sdk/storage/azure_storage_blob/tsp-location.yaml b/sdk/storage/azure_storage_blob/tsp-location.yaml index 3f59bbd952f..e3b4229c30c 100644 --- a/sdk/storage/azure_storage_blob/tsp-location.yaml +++ b/sdk/storage/azure_storage_blob/tsp-location.yaml @@ -1,4 +1,4 @@ directory: specification/storage/data-plane/BlobStorage -commit: 51c2be144398b7ab9ff7e03c8afa217227cec28f +commit: f25cb2f1a14e113dca1d7df37a164e7a2f84f314 repo: Azure/azure-rest-api-specs additionalDirectories: From b38f2a7989208cfb4afaf421791dff6b17c65fef Mon Sep 17 00:00:00 2001 From: vincenttran-msft <101599632+vincenttran-msft@users.noreply.github.com> Date: Tue, 15 Sep 2026 14:54:17 -0700 Subject: [PATCH 6/6] accessor-derived clients now also inherit layout pipeline policy, expose layout_endpoint to manually override, add testcases, unify to singular test file --- .../src/clients/blob_client.rs | 51 +- .../azure_storage_blob/src/clients/mod.rs | 10 +- .../src/models/method_options.rs | 23 +- .../azure_storage_blob/tests/blob_client.rs | 605 +-------------- .../tests/blob_client_data_locality.rs | 734 ++++++++++++++++++ 5 files changed, 806 insertions(+), 617 deletions(-) create mode 100644 sdk/storage/azure_storage_blob/tests/blob_client_data_locality.rs diff --git a/sdk/storage/azure_storage_blob/src/clients/blob_client.rs b/sdk/storage/azure_storage_blob/src/clients/blob_client.rs index dd6559b1cd5..b0238c38614 100644 --- a/sdk/storage/azure_storage_blob/src/clients/blob_client.rs +++ b/sdk/storage/azure_storage_blob/src/clients/blob_client.rs @@ -4,7 +4,7 @@ pub use crate::generated::clients::{BlobClient, BlobClientOptions}; use crate::{ - blob_layout::{fetch_layout, LayoutCache, LayoutEndpoint, LayoutRoutingPolicy}, + blob_layout::{fetch_layout, LayoutCache, LayoutEndpoint}, generated::{ clients::BlobClient as GeneratedBlobClient, models::{BlobClientDownloadInternalOptions, BlobClientGetLayoutOptions}, @@ -77,7 +77,7 @@ impl BlobClient { option_env!("CARGO_PKG_NAME"), option_env!("CARGO_PKG_VERSION"), options.client_options.clone(), - vec![Arc::new(LayoutRoutingPolicy)], + Vec::default(), per_retry_policies, None, ); @@ -208,10 +208,12 @@ impl BlobClient { }; let range = options.range.clone(); let layout_aware_routing = options.layout_aware_routing; + let layout_endpoint = options.layout_endpoint.clone(); let behavior = BlobClientDownloadBehavior::new( inner_client, options.into(), layout_aware_routing, + layout_endpoint, range.clone(), ); let response = @@ -258,10 +260,12 @@ impl BlobClient { }; let range = options.range.clone(); let layout_aware_routing = options.layout_aware_routing; + let layout_endpoint = options.layout_endpoint.clone(); let behavior = BlobClientDownloadBehavior::new( inner_client, options.into(), layout_aware_routing, + layout_endpoint, range.clone(), ); let (_, headers, len) = partitioned_transfer::download_into( @@ -323,6 +327,8 @@ struct BlobClientDownloadBehavior<'a> { client: Arc, options: BlobClientDownloadInternalOptions<'a>, layout_aware_routing: LayoutAwareRouting, + /// Caller-supplied endpoint that pins every request this download issues. + layout_endpoint: Option, /// The caller-requested range, which `options.range` does not retain because it is rewritten for each partition. requested_range: Option, layout_cache: OnceLock>, @@ -333,17 +339,36 @@ impl<'a> BlobClientDownloadBehavior<'a> { client: GeneratedBlobClient, options: BlobClientDownloadInternalOptions<'a>, layout_aware_routing: LayoutAwareRouting, + layout_endpoint: Option, requested_range: Option, ) -> Self { Self { client: Arc::new(client), options, layout_aware_routing, + layout_endpoint, requested_range, layout_cache: OnceLock::new(), } } + /// Resolves the endpoint a request covering `range` should be sent to, preferring a + /// caller-supplied endpoint over the fetched layout. + /// + /// A caller-supplied endpoint applies to every request, including the initial one, + /// because it pins the whole call rather than following the blob's layout. + async fn resolve_endpoint(&self, range: Option<&Range>) -> Option { + if let Some(endpoint) = self.layout_endpoint.as_ref() { + return Some(endpoint.clone()); + } + if let (Some(Some(cache)), Some(range)) = (self.layout_cache.get(), range) { + if let Some(layout) = cache.current().await { + return layout.ideal_endpoint(range.start as i64).map(str::to_owned); + } + } + None + } + fn layout_options(&self) -> BlobClientGetLayoutOptions<'static> { BlobClientGetLayoutOptions { encryption_algorithm: self.options.encryption_algorithm, @@ -372,18 +397,12 @@ impl PartitionedDownloadBehavior for BlobClientDownloadBehavior<'_> { etag_lock: Option, ) -> Result { let mut opt = self.options.clone(); - if let Some(Some(cache)) = self.layout_cache.get() { - if let Some(range) = &range { - if let Some(layout) = cache.current().await { - if let Some(endpoint) = layout.ideal_endpoint(range.start as i64) { - opt.method_options.context = opt - .method_options - .context - .clone() - .with_value(LayoutEndpoint(endpoint.to_owned())); - } - } - } + if let Some(endpoint) = self.resolve_endpoint(range.as_ref()).await { + opt.method_options.context = opt + .method_options + .context + .clone() + .with_value(LayoutEndpoint(endpoint)); } opt.range = range.map(HttpRange::from); if let Some(etag) = etag_lock { @@ -400,7 +419,9 @@ impl PartitionedDownloadBehavior for BlobClientDownloadBehavior<'_> { } async fn prepare(&self, initial_headers: &Headers, etag_lock: Option<&Etag>) -> Result<()> { - if matches!(self.layout_aware_routing, LayoutAwareRouting::Disabled) { + if self.layout_endpoint.is_some() + || matches!(self.layout_aware_routing, LayoutAwareRouting::Disabled) + { return Ok(()); } if initial_headers.get_optional_str(&"x-ms-download-hint".into()) != Some("layout") { diff --git a/sdk/storage/azure_storage_blob/src/clients/mod.rs b/sdk/storage/azure_storage_blob/src/clients/mod.rs index 2ea0b70b3c7..2b5fbf484b8 100644 --- a/sdk/storage/azure_storage_blob/src/clients/mod.rs +++ b/sdk/storage/azure_storage_blob/src/clients/mod.rs @@ -4,8 +4,9 @@ //! Clients used to communicate with Azure Blob Storage. use azure_core::http::{new_http_client, ClientOptions, HttpClientOptions, Transport}; +use std::sync::Arc; -use crate::logging::apply_storage_logging_defaults; +use crate::{blob_layout::LayoutRoutingPolicy, logging::apply_storage_logging_defaults}; mod append_blob_client; mod blob_client; @@ -22,6 +23,10 @@ pub use block_blob_client::{BlockBlobClient, BlockBlobClientOptions}; pub use page_blob_client::{PageBlobClient, PageBlobClientOptions}; #[allow(clippy::needless_update)] +/// Applies defaults shared by every client. +/// +/// Derived clients reuse their parent's pipeline rather than rebuilding it, so +/// anything added here must be valid for all client types. fn apply_client_defaults(options: &mut ClientOptions) { if options.transport.is_none() { options.transport = Some(Transport::new(new_http_client(Some(HttpClientOptions { @@ -29,5 +34,8 @@ fn apply_client_defaults(options: &mut ClientOptions) { ..Default::default() })))) } + options + .per_call_policies + .push(Arc::new(LayoutRoutingPolicy)); apply_storage_logging_defaults(options); } diff --git a/sdk/storage/azure_storage_blob/src/models/method_options.rs b/sdk/storage/azure_storage_blob/src/models/method_options.rs index 8733c6c37dd..5c28e4b7965 100644 --- a/sdk/storage/azure_storage_blob/src/models/method_options.rs +++ b/sdk/storage/azure_storage_blob/src/models/method_options.rs @@ -70,9 +70,30 @@ pub struct BlobClientDownloadOptions<'a> { /// requests issued by this download. This is a performance optimization only: /// the bytes returned are identical regardless of the mode. /// - /// Defaults to [`LayoutAwareRouting::Enabled`]. + /// Defaults to [`LayoutAwareRouting::Enabled`]. Ignored when + /// [`layout_endpoint`](Self::layout_endpoint) is set. pub layout_aware_routing: LayoutAwareRouting, + /// Optional. The layout endpoint (`host:port`) that every request issued by this + /// download is sent to, with the client's configured authority preserved as the + /// `Host` header. + /// + /// Setting this disables automatic layout lookup: the blob's layout is not fetched + /// and [`layout_aware_routing`](Self::layout_aware_routing) is ignored. To choose an + /// endpoint, enumerate the pages returned by + /// [`BlobClient::get_layout()`](crate::BlobClient::get_layout) and select the endpoint + /// whose layout range covers the offset of the requested [`range`](Self::range). + /// + /// Because the endpoint applies to every request the download issues, set a + /// [`range`](Self::range) that falls within a single layout range; a range spanning + /// several layout ranges still returns the correct bytes, but sends them all to one + /// endpoint. A malformed or unreachable endpoint fails the download; the request is + /// not retried against the client's configured endpoint. + /// + /// When `None` (the default), requests are sent to the client's configured endpoint + /// unless automatic routing applies. + pub layout_endpoint: Option, + /// Allows customization of the method call. pub method_options: ClientMethodOptions<'a>, diff --git a/sdk/storage/azure_storage_blob/tests/blob_client.rs b/sdk/storage/azure_storage_blob/tests/blob_client.rs index b7ab5af48b0..59295cd1561 100644 --- a/sdk/storage/azure_storage_blob/tests/blob_client.rs +++ b/sdk/storage/azure_storage_blob/tests/blob_client.rs @@ -3,21 +3,17 @@ mod common; -use async_trait::async_trait; use azure_core::{ - credentials::TokenCredential, error::ErrorKind, http::{ headers::{HeaderName, Headers, CONTENT_TYPE}, - policies::{Policy, PolicyResult}, - AsyncRawResponse, ClientOptions, Context, Method, Request, RequestContent, StatusCode, - Transport, Url, XmlFormat, + AsyncRawResponse, ClientOptions, Method, RequestContent, StatusCode, Transport, Url, + XmlFormat, }, time::{parse_rfc3339, to_rfc3339, OffsetDateTime}, Bytes, }; use azure_core_test::{http::MockHttpClient, recorded, Matcher, TestContext, TestMode, VarOptions}; -use azure_identity::DeveloperToolsCredential; use azure_storage_blob::{ models::{ AccessTier, AccountKind, ArchiveStatus, BlobClientAcquireLeaseOptions, @@ -27,8 +23,8 @@ use azure_storage_blob::{ BlobClientSetImmutabilityPolicyOptions, BlobClientSetMetadataOptions, BlobClientSetPropertiesOptions, BlobClientSetTierOptions, BlobClientStartCopyFromUrlResultHeaders, BlobTags, BlockBlobClientCommitBlockListOptions, - BlockBlobClientUploadOptions, CopyStatus, ImmutabilityPolicyMode, KeyInfo, - LayoutAwareRouting, LeaseState, RehydratePriority, StorageErrorCode, + BlockBlobClientUploadOptions, CopyStatus, ImmutabilityPolicyMode, KeyInfo, LeaseState, + RehydratePriority, StorageErrorCode, }, BlobClient, BlobClientOptions, BlobContainerClient, BlobContainerClientOptions, StorageError, }; @@ -48,7 +44,7 @@ use std::{ num::NonZero, sync::{ atomic::{AtomicUsize, Ordering}, - mpsc, Arc, Mutex, + mpsc, Arc, }, time::Duration, }; @@ -1532,597 +1528,6 @@ async fn test_managed_download_into_etag_lock(ctx: TestContext) -> Result<(), Bo Ok(()) } -fn parse_bytes_range(range: &str) -> (usize, usize) { - let spec = range.strip_prefix("bytes=").expect("bytes= prefix"); - let (start, end) = spec.split_once('-').expect("start-end"); - ( - start.parse().expect("range start"), - end.parse().expect("range end"), - ) -} - -fn live_layout_blob_url() -> Option { - let account = std::env::var("AZURE_STORAGE_ACCOUNT_NAME").ok()?; - let account = account.trim().trim_matches('"').trim(); - if account.is_empty() { - return None; - } - Url::parse(&format!( - "https://{account}.blob.core.windows.net/layout-routing-test/routing-test-blob" - )) - .ok() -} - -#[tokio::test] -async fn test_download_layout_aware_routing_routes_chunks() -> Result<(), Box> { - const DATA: [u8; 12] = [10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21]; - const LAYOUT: &[u8] = br#" - - - - - - - - - - - -"#; - - #[derive(Clone)] - struct Seen { - is_layout: bool, - url_host: Option, - host_header: Option, - range: Option, - } - - let seen = Arc::new(Mutex::new(Vec::::new())); - let seen_capture = seen.clone(); - let mock_client = Arc::new(MockHttpClient::new(move |request| { - let is_layout = request - .url() - .query() - .is_some_and(|query| query.contains("comp=layout")); - let url_host = request.url().host_str().map(str::to_owned); - let host_header = request - .headers() - .get_optional_str(&"host".into()) - .map(str::to_owned); - let range = request - .headers() - .get_optional_str(&"range".into()) - .map(str::to_owned); - seen_capture.lock().unwrap().push(Seen { - is_layout, - url_host, - host_header, - range: range.clone(), - }); - - let response = if is_layout { - let mut headers = Headers::new(); - headers.insert("etag", "\"routing-etag\""); - AsyncRawResponse::from_bytes(StatusCode::Ok, headers, Bytes::from_static(LAYOUT)) - } else { - let range = range.expect("data request must carry a range header"); - let (start, end) = parse_bytes_range(&range); - let slice = DATA[start..=end].to_vec(); - let mut headers = Headers::new(); - headers.insert( - "content-range", - format!("bytes {start}-{end}/{}", DATA.len()), - ); - headers.insert("content-length", slice.len().to_string()); - headers.insert("etag", "\"routing-etag\""); - headers.insert("x-ms-download-hint", "layout"); - AsyncRawResponse::from_bytes(StatusCode::PartialContent, headers, Bytes::from(slice)) - }; - async move { Ok(response) }.boxed() - })); - - let blob_client = BlobClient::new( - Url::parse("https://acct.blob.core.windows.net/container/blob")?, - None, - Some(BlobClientOptions { - client_options: ClientOptions { - transport: Some(Transport::new(mock_client)), - ..Default::default() - }, - ..Default::default() - }), - )?; - - let body = blob_client - .download(Some(BlobClientDownloadOptions { - layout_aware_routing: LayoutAwareRouting::Enabled, - partition_size: Some(NonZero::new(4).unwrap()), - parallel: Some(NonZero::new(2).unwrap()), - ..Default::default() - })) - .await? - .body - .collect() - .await?; - - assert_eq!(&body[..], &DATA[..]); - let seen = seen.lock().unwrap(); - let layout_requests: Vec<&Seen> = seen.iter().filter(|request| request.is_layout).collect(); - let data_requests: Vec<&Seen> = seen.iter().filter(|request| !request.is_layout).collect(); - assert_eq!(layout_requests.len(), 1); - assert_eq!( - layout_requests[0].url_host.as_deref(), - Some("acct.blob.core.windows.net") - ); - - assert_eq!(data_requests.len(), 3); - for request in &data_requests { - let range = request.range.as_deref().expect("range header"); - match range { - "bytes=0-3" => assert_eq!( - request.url_host.as_deref(), - Some("acct.blob.core.windows.net"), - "the initial chunk should not be routed" - ), - "bytes=4-7" | "bytes=8-11" => { - let expected = if range == "bytes=4-7" { - "epb.blob.core.windows.net" - } else { - "epc.blob.core.windows.net" - }; - assert_eq!(request.url_host.as_deref(), Some(expected)); - assert_eq!( - request.host_header.as_deref(), - Some("acct.blob.core.windows.net") - ); - } - other => panic!("unexpected data range: {other}"), - } - } - Ok(()) -} - -#[tokio::test] -async fn test_download_layout_aware_routing_routes_chunks_for_path_style_endpoint( -) -> Result<(), Box> { - const DATA: [u8; 12] = [10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21]; - const LAYOUT: &[u8] = br#" - - - - - - - - - - - -"#; - - struct Seen { - is_layout: bool, - url: Url, - host_header: Option, - range: Option, - } - - let seen = Arc::new(Mutex::new(Vec::::new())); - let seen_capture = seen.clone(); - let mock_client = Arc::new(MockHttpClient::new(move |request| { - let is_layout = request - .url() - .query() - .is_some_and(|query| query.contains("comp=layout")); - let host_header = request - .headers() - .get_optional_str(&"host".into()) - .map(str::to_owned); - let range = request - .headers() - .get_optional_str(&"range".into()) - .map(str::to_owned); - seen_capture.lock().unwrap().push(Seen { - is_layout, - url: request.url().clone(), - host_header, - range: range.clone(), - }); - - let response = if is_layout { - let mut headers = Headers::new(); - headers.insert("etag", "\"path-style-etag\""); - AsyncRawResponse::from_bytes(StatusCode::Ok, headers, Bytes::from_static(LAYOUT)) - } else { - let range = range.expect("data request must carry a range header"); - let (start, end) = parse_bytes_range(&range); - let slice = DATA[start..=end].to_vec(); - let mut headers = Headers::new(); - headers.insert( - "content-range", - format!("bytes {start}-{end}/{}", DATA.len()), - ); - headers.insert("content-length", slice.len().to_string()); - headers.insert("etag", "\"path-style-etag\""); - headers.insert("x-ms-download-hint", "layout"); - AsyncRawResponse::from_bytes(StatusCode::PartialContent, headers, Bytes::from(slice)) - }; - async move { Ok(response) }.boxed() - })); - - let blob_client = BlobClient::new( - Url::parse("http://127.0.0.1:10000/devstoreaccount1/container/blob")?, - None, - Some(BlobClientOptions { - client_options: ClientOptions { - transport: Some(Transport::new(mock_client)), - ..Default::default() - }, - ..Default::default() - }), - )?; - - let body = blob_client - .download(Some(BlobClientDownloadOptions { - layout_aware_routing: LayoutAwareRouting::Enabled, - partition_size: Some(NonZero::new(4).unwrap()), - parallel: Some(NonZero::new(2).unwrap()), - ..Default::default() - })) - .await? - .body - .collect() - .await?; - - assert_eq!(&body[..], &DATA[..]); - let seen = seen.lock().unwrap(); - let layout_requests: Vec<&Seen> = seen.iter().filter(|request| request.is_layout).collect(); - let data_requests: Vec<&Seen> = seen.iter().filter(|request| !request.is_layout).collect(); - - assert_eq!(layout_requests.len(), 1); - assert_eq!(layout_requests[0].url.host_str(), Some("127.0.0.1")); - - assert_eq!(data_requests.len(), 3); - for request in &data_requests { - let range = request.range.as_deref().expect("range header"); - // Rerouting must not disturb the emulator's scheme or its account path segment. - assert_eq!(request.url.scheme(), "http"); - assert_eq!(request.url.path(), "/devstoreaccount1/container/blob"); - match range { - "bytes=0-3" => { - assert_eq!( - request.url.host_str(), - Some("127.0.0.1"), - "the initial chunk should not be routed" - ); - assert_eq!(request.url.port(), Some(10000)); - assert_eq!(request.host_header, None); - } - "bytes=4-7" | "bytes=8-11" => { - let expected = if range == "bytes=4-7" { - "node-b.storage.local" - } else { - "node-c.storage.local" - }; - assert_eq!(request.url.host_str(), Some(expected)); - assert_eq!(request.url.port(), Some(20000)); - assert_eq!(request.host_header.as_deref(), Some("127.0.0.1:10000")); - } - other => panic!("unexpected data range: {other}"), - } - } - Ok(()) -} - -#[tokio::test] -async fn test_download_layout_aware_routing_skips_layout_for_complete_initial_range( -) -> Result<(), Box> { - const DATA: [u8; 8] = [10, 11, 12, 13, 14, 15, 16, 17]; - - let request_count = Arc::new(AtomicUsize::new(0)); - let count_capture = request_count.clone(); - let mock_client = Arc::new(MockHttpClient::new(move |request| { - assert_eq!(0, count_capture.fetch_add(1, Ordering::SeqCst)); - assert!(!request - .url() - .query() - .is_some_and(|query| query.contains("comp=layout"))); - assert_eq!( - Some("bytes=0-7"), - request.headers().get_optional_str(&"range".into()) - ); - - let mut headers = Headers::new(); - headers.insert("content-range", "bytes 0-7/8"); - headers.insert("content-length", DATA.len().to_string()); - headers.insert("etag", "\"complete-range-etag\""); - headers.insert("x-ms-download-hint", "layout"); - async move { - Ok(AsyncRawResponse::from_bytes( - StatusCode::PartialContent, - headers, - Bytes::from_static(&DATA), - )) - } - .boxed() - })); - - let blob_client = BlobClient::new( - Url::parse("https://acct.blob.core.windows.net/container/blob")?, - None, - Some(BlobClientOptions { - client_options: ClientOptions { - transport: Some(Transport::new(mock_client)), - ..Default::default() - }, - ..Default::default() - }), - )?; - - let body = blob_client - .download(Some(BlobClientDownloadOptions { - layout_aware_routing: LayoutAwareRouting::Enabled, - partition_size: Some(NonZero::new(DATA.len()).unwrap()), - ..Default::default() - })) - .await? - .body - .collect() - .await?; - - assert_eq!(&body[..], &DATA[..]); - assert_eq!(request_count.load(Ordering::SeqCst), 1); - Ok(()) -} - -#[tokio::test] -async fn test_download_layout_aware_routing_skips_without_hint() -> Result<(), Box> { - const DATA: [u8; 8] = [10, 11, 12, 13, 14, 15, 16, 17]; - - let layout_requests = Arc::new(AtomicUsize::new(0)); - let hosts = Arc::new(Mutex::new(Vec::>::new())); - let layout_capture = layout_requests.clone(); - let hosts_capture = hosts.clone(); - let mock_client = Arc::new(MockHttpClient::new(move |request| { - let is_layout = request - .url() - .query() - .is_some_and(|query| query.contains("comp=layout")); - hosts_capture - .lock() - .unwrap() - .push(request.url().host_str().map(str::to_owned)); - let range = request - .headers() - .get_optional_str(&"range".into()) - .map(str::to_owned); - - let response = if is_layout { - layout_capture.fetch_add(1, Ordering::SeqCst); - AsyncRawResponse::from_bytes(StatusCode::NoContent, Headers::new(), Bytes::new()) - } else { - let range = range.expect("data request must carry a range header"); - let (start, end) = parse_bytes_range(&range); - let slice = DATA[start..=end].to_vec(); - let mut headers = Headers::new(); - headers.insert( - "content-range", - format!("bytes {start}-{end}/{}", DATA.len()), - ); - headers.insert("content-length", slice.len().to_string()); - headers.insert("etag", "\"no-hint-etag\""); - AsyncRawResponse::from_bytes(StatusCode::PartialContent, headers, Bytes::from(slice)) - }; - async move { Ok(response) }.boxed() - })); - - let blob_client = BlobClient::new( - Url::parse("https://acct.blob.core.windows.net/container/blob")?, - None, - Some(BlobClientOptions { - client_options: ClientOptions { - transport: Some(Transport::new(mock_client)), - ..Default::default() - }, - ..Default::default() - }), - )?; - - let body = blob_client - .download(Some(BlobClientDownloadOptions { - layout_aware_routing: LayoutAwareRouting::Enabled, - partition_size: Some(NonZero::new(4).unwrap()), - parallel: Some(NonZero::new(2).unwrap()), - ..Default::default() - })) - .await? - .body - .collect() - .await?; - - assert_eq!(&body[..], &DATA[..]); - assert_eq!(layout_requests.load(Ordering::SeqCst), 0); - for host in hosts.lock().unwrap().iter() { - assert_eq!(host.as_deref(), Some("acct.blob.core.windows.net")); - } - Ok(()) -} - -#[derive(Debug)] -struct RoutingObserver { - seen: Arc>>, -} - -#[derive(Debug)] -struct RoutedRequest { - is_layout: bool, - range: Option, - sent_to: Option, - original_host: Option, - download_hint: Option, - status: StatusCode, -} - -#[async_trait] -impl Policy for RoutingObserver { - async fn send( - &self, - ctx: &Context, - request: &mut Request, - next: &[Arc], - ) -> PolicyResult { - let is_layout = request - .url() - .query() - .is_some_and(|query| query.contains("comp=layout")); - let range = request - .headers() - .get_optional_str(&"range".into()) - .map(str::to_owned); - let sent_to = request.url().host_str().map(str::to_owned); - let original_host = request - .headers() - .get_optional_str(&"host".into()) - .map(str::to_owned); - let response = next[0].send(ctx, request, &next[1..]).await?; - let download_hint = response - .headers() - .get_optional_str(&"x-ms-download-hint".into()) - .map(str::to_owned); - self.seen.lock().unwrap().push(RoutedRequest { - is_layout, - range, - sent_to, - original_host, - download_hint, - status: response.status(), - }); - Ok(response) - } -} - -#[recorded::test(live)] -async fn test_download_layout_aware_routing() -> Result<(), Box> { - let Some(url) = live_layout_blob_url() else { - eprintln!( - "skipping test_download_layout_aware_routing: set AZURE_STORAGE_ACCOUNT_NAME to run it" - ); - return Ok(()); - }; - let credential: Arc = DeveloperToolsCredential::new(None)?; - let account_host = url.host_str().unwrap_or_default().to_owned(); - let observed = Arc::new(Mutex::new(Vec::::new())); - let observer: Arc = Arc::new(RoutingObserver { - seen: observed.clone(), - }); - - let mut container_url = url.clone(); - container_url - .path_segments_mut() - .expect("blob URL must be a base") - .pop(); - let container_client = BlobContainerClient::new(container_url, Some(credential.clone()), None)?; - if let Err(error) = container_client.create(None).await { - if error.http_status() != Some(StatusCode::Conflict) { - return Err(error.into()); - } - } - - let blob_client = BlobClient::new( - url, - Some(credential), - Some(BlobClientOptions { - client_options: ClientOptions { - per_try_policies: vec![observer], - ..Default::default() - }, - ..Default::default() - }), - )?; - - const BLOB_SIZE: usize = 64 * 1024 * 1024; - let data: Vec = (0..BLOB_SIZE).map(|index| (index % 251) as u8).collect(); - blob_client - .upload(RequestContent::from(data.clone()), None) - .await?; - observed.lock().unwrap().clear(); - - let body = blob_client - .download(Some(BlobClientDownloadOptions { - layout_aware_routing: LayoutAwareRouting::Enabled, - partition_size: Some(NonZero::new(16 * 1024 * 1024).unwrap()), - parallel: Some(NonZero::new(4).unwrap()), - ..Default::default() - })) - .await? - .body - .collect() - .await?; - assert_eq!(&body[..], &data[..]); - - let mut buffer = vec![0u8; data.len()]; - let result = blob_client - .download_into( - &mut buffer, - Some(BlobClientDownloadOptions { - layout_aware_routing: LayoutAwareRouting::Enabled, - partition_size: Some(NonZero::new(16 * 1024 * 1024).unwrap()), - parallel: Some(NonZero::new(4).unwrap()), - ..Default::default() - }), - ) - .await?; - assert_eq!(result.len, data.len()); - assert_eq!(&buffer[..], &data[..]); - - let observed = observed.lock().unwrap(); - let layout_calls = observed.iter().filter(|request| request.is_layout).count(); - let data_chunks = observed.iter().filter(|request| !request.is_layout).count(); - let routed: Vec<&RoutedRequest> = observed - .iter() - .filter(|request| !request.is_layout && request.original_host.is_some()) - .collect(); - - eprintln!(); - eprintln!("=== layout-aware routing ======================================"); - eprintln!(" account host : {account_host}"); - eprintln!(" Get Blob Layout calls : {layout_calls}"); - eprintln!( - " data chunks : {data_chunks} ({} routed)", - routed.len() - ); - for request in observed.iter() { - let kind = if request.is_layout { "layout" } else { "data" }; - let range = request.range.as_deref().unwrap_or("(whole blob)"); - let sent_to = request.sent_to.as_deref().unwrap_or("?"); - let hint = request.download_hint.as_deref().unwrap_or("-"); - let routed = if request.original_host.is_some() { - "yes" - } else { - "no" - }; - let status = request.status; - eprintln!(" {kind:<6} {status:<3} {range:<24} {routed:<7} {hint:<7} {sent_to}"); - } - - if routed.is_empty() { - eprintln!( - "NOTE: no chunk was rewritten - the account returned no layout hint for this blob" - ); - } else { - for request in &routed { - assert_eq!( - request.original_host.as_deref(), - Some(account_host.as_str()), - "a rewritten chunk did not preserve the account Host header" - ); - } - } - - Ok(()) -} - #[recorded::test] async fn test_upload_blob_content_headers(ctx: TestContext) -> Result<(), Box> { // Recording Setup diff --git a/sdk/storage/azure_storage_blob/tests/blob_client_data_locality.rs b/sdk/storage/azure_storage_blob/tests/blob_client_data_locality.rs new file mode 100644 index 00000000000..6187c87826e --- /dev/null +++ b/sdk/storage/azure_storage_blob/tests/blob_client_data_locality.rs @@ -0,0 +1,734 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Data-locality tests: the SDK's automatic layout-aware routing of managed downloads, +//! and the caller-controlled `layout_endpoint` override. + +use async_trait::async_trait; +use azure_core::{ + credentials::TokenCredential, + http::{ + headers::Headers, + policies::{Policy, PolicyResult}, + AsyncRawResponse, ClientOptions, Context, Request, RequestContent, StatusCode, Transport, + Url, + }, + Bytes, +}; +use azure_core_test::{http::MockHttpClient, recorded}; +use azure_identity::DeveloperToolsCredential; +use azure_storage_blob::{ + models::{BlobClientDownloadOptions, BlobLayout, LayoutAwareRouting}, + BlobClient, BlobClientOptions, BlobContainerClient, BlobContainerClientOptions, +}; +use futures::{FutureExt as _, TryStreamExt}; +use std::{ + collections::HashMap, + error::Error, + num::NonZero, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, Mutex, + }, +}; + +const ACCOUNT_HOST: &str = "acct.blob.core.windows.net"; +const ETAG: &str = "\"layout-test-etag\""; + +const BLOB_DATA: [u8; 12] = [10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21]; + +/// Three endpoints addressed by hostname, each serving a four-byte range of `BLOB_DATA`. +const LAYOUT: &[u8] = br#" + + + + + + + + + + + +"#; + +/// The same shape as [`LAYOUT`], but with the non-default port an emulator would report. +const EMULATOR_LAYOUT: &[u8] = br#" + + + + + + + + + + + +"#; + +/// A request as the transport saw it, after any rewrite the routing policy applied. +#[derive(Clone)] +struct ObservedRequest { + is_layout: bool, + url: Url, + original_host: Option, + range: Option, +} + +/// Whether data responses carry `x-ms-download-hint`, which is what gates the SDK's +/// automatic layout lookup. +#[derive(Clone, Copy, PartialEq)] +enum LayoutHint { + Advertised, + Absent, +} + +fn parse_bytes_range(range: &str) -> (usize, usize) { + let spec = range.strip_prefix("bytes=").expect("bytes= prefix"); + let (start, end) = spec.split_once('-').expect("start-end"); + ( + start.parse().expect("range start"), + end.parse().expect("range end"), + ) +} + +/// Builds a transport that serves `layout` for Get Blob Layout and slices `BLOB_DATA` +/// for range requests, recording where each request was sent. +fn layout_mock_transport( + seen: Arc>>, + layout: &'static [u8], + hint: LayoutHint, +) -> Transport { + Transport::new(Arc::new(MockHttpClient::new(move |request| { + let is_layout = request + .url() + .query() + .is_some_and(|query| query.contains("comp=layout")); + let range = request + .headers() + .get_optional_str(&"range".into()) + .map(str::to_owned); + seen.lock().unwrap().push(ObservedRequest { + is_layout, + url: request.url().clone(), + original_host: request + .headers() + .get_optional_str(&"host".into()) + .map(str::to_owned), + range: range.clone(), + }); + + let response = if is_layout { + let mut headers = Headers::new(); + headers.insert("etag", ETAG); + AsyncRawResponse::from_bytes(StatusCode::Ok, headers, Bytes::from_static(layout)) + } else { + let range = range.expect("data request must carry a range header"); + let (start, end) = parse_bytes_range(&range); + let slice = BLOB_DATA[start..=end].to_vec(); + let mut headers = Headers::new(); + headers.insert( + "content-range", + format!("bytes {start}-{end}/{}", BLOB_DATA.len()), + ); + headers.insert("content-length", slice.len().to_string()); + headers.insert("etag", ETAG); + if hint == LayoutHint::Advertised { + headers.insert("x-ms-download-hint", "layout"); + } + AsyncRawResponse::from_bytes(StatusCode::PartialContent, headers, Bytes::from(slice)) + }; + async move { Ok(response) }.boxed() + }))) +} + +fn blob_client_with(transport: Transport, url: &str) -> Result> { + Ok(BlobClient::new( + Url::parse(url)?, + None, + Some(BlobClientOptions { + client_options: ClientOptions { + transport: Some(transport), + ..Default::default() + }, + ..Default::default() + }), + )?) +} + +/// Mirrors the work a caller does with the pages returned by `get_layout()`: find the +/// endpoint whose layout range covers `offset`. +fn endpoint_for_offset(layout: &BlobLayout, offset: i64) -> Option { + let endpoints = layout.endpoints.as_ref()?.endpoint.as_ref()?; + let index = layout + .ranges + .range + .iter() + .find(|range| { + range.start.unwrap_or_default() <= offset && offset <= range.end.unwrap_or_default() + })? + .endpoint_index?; + endpoints + .iter() + .find(|endpoint| endpoint.index == Some(index)) + .and_then(|endpoint| endpoint.value.clone()) +} + +#[tokio::test] +async fn test_download_layout_aware_routing_routes_chunks() -> Result<(), Box> { + let seen = Arc::new(Mutex::new(Vec::::new())); + let transport = layout_mock_transport(seen.clone(), LAYOUT, LayoutHint::Advertised); + let blob_client = blob_client_with( + transport, + "https://acct.blob.core.windows.net/container/blob", + )?; + + let body = blob_client + .download(Some(BlobClientDownloadOptions { + layout_aware_routing: LayoutAwareRouting::Enabled, + partition_size: Some(NonZero::new(4).unwrap()), + parallel: Some(NonZero::new(2).unwrap()), + ..Default::default() + })) + .await? + .body + .collect() + .await?; + + assert_eq!(&body[..], &BLOB_DATA[..]); + let seen = seen.lock().unwrap(); + let layout_requests: Vec<&ObservedRequest> = + seen.iter().filter(|request| request.is_layout).collect(); + let data_requests: Vec<&ObservedRequest> = + seen.iter().filter(|request| !request.is_layout).collect(); + + assert_eq!(layout_requests.len(), 1); + assert_eq!(layout_requests[0].url.host_str(), Some(ACCOUNT_HOST)); + + assert_eq!(data_requests.len(), 3); + for request in &data_requests { + let range = request.range.as_deref().expect("range header"); + match range { + "bytes=0-3" => assert_eq!( + request.url.host_str(), + Some(ACCOUNT_HOST), + "the initial chunk should not be routed" + ), + "bytes=4-7" | "bytes=8-11" => { + let expected = if range == "bytes=4-7" { + "epb.blob.core.windows.net" + } else { + "epc.blob.core.windows.net" + }; + assert_eq!(request.url.host_str(), Some(expected)); + assert_eq!(request.original_host.as_deref(), Some(ACCOUNT_HOST)); + } + other => panic!("unexpected data range: {other}"), + } + } + Ok(()) +} + +#[tokio::test] +async fn test_download_layout_aware_routing_routes_chunks_for_path_style_endpoint( +) -> Result<(), Box> { + let seen = Arc::new(Mutex::new(Vec::::new())); + let transport = layout_mock_transport(seen.clone(), EMULATOR_LAYOUT, LayoutHint::Advertised); + let blob_client = blob_client_with( + transport, + "http://127.0.0.1:10000/devstoreaccount1/container/blob", + )?; + + let body = blob_client + .download(Some(BlobClientDownloadOptions { + layout_aware_routing: LayoutAwareRouting::Enabled, + partition_size: Some(NonZero::new(4).unwrap()), + parallel: Some(NonZero::new(2).unwrap()), + ..Default::default() + })) + .await? + .body + .collect() + .await?; + + assert_eq!(&body[..], &BLOB_DATA[..]); + let seen = seen.lock().unwrap(); + let layout_requests: Vec<&ObservedRequest> = + seen.iter().filter(|request| request.is_layout).collect(); + let data_requests: Vec<&ObservedRequest> = + seen.iter().filter(|request| !request.is_layout).collect(); + + assert_eq!(layout_requests.len(), 1); + assert_eq!(layout_requests[0].url.host_str(), Some("127.0.0.1")); + + assert_eq!(data_requests.len(), 3); + for request in &data_requests { + let range = request.range.as_deref().expect("range header"); + // Rerouting must not disturb the emulator's scheme or its account path segment. + assert_eq!(request.url.scheme(), "http"); + assert_eq!(request.url.path(), "/devstoreaccount1/container/blob"); + match range { + "bytes=0-3" => { + assert_eq!( + request.url.host_str(), + Some("127.0.0.1"), + "the initial chunk should not be routed" + ); + assert_eq!(request.url.port(), Some(10000)); + assert_eq!(request.original_host, None); + } + "bytes=4-7" | "bytes=8-11" => { + let expected = if range == "bytes=4-7" { + "node-b.storage.local" + } else { + "node-c.storage.local" + }; + assert_eq!(request.url.host_str(), Some(expected)); + assert_eq!(request.url.port(), Some(20000)); + assert_eq!(request.original_host.as_deref(), Some("127.0.0.1:10000")); + } + other => panic!("unexpected data range: {other}"), + } + } + Ok(()) +} + +/// Without a layout hint the SDK must not look the layout up, and must leave every +/// request pointed at the client's configured endpoint. +#[tokio::test] +async fn test_download_layout_aware_routing_skips_without_hint() -> Result<(), Box> { + let seen = Arc::new(Mutex::new(Vec::::new())); + let transport = layout_mock_transport(seen.clone(), LAYOUT, LayoutHint::Absent); + let blob_client = blob_client_with( + transport, + "https://acct.blob.core.windows.net/container/blob", + )?; + + let body = blob_client + .download(Some(BlobClientDownloadOptions { + layout_aware_routing: LayoutAwareRouting::Enabled, + partition_size: Some(NonZero::new(4).unwrap()), + parallel: Some(NonZero::new(2).unwrap()), + ..Default::default() + })) + .await? + .body + .collect() + .await?; + + assert_eq!(&body[..], &BLOB_DATA[..]); + let seen = seen.lock().unwrap(); + assert!( + !seen.iter().any(|request| request.is_layout), + "no layout should be fetched without a hint" + ); + for request in seen.iter() { + assert_eq!(request.url.host_str(), Some(ACCOUNT_HOST)); + } + Ok(()) +} + +/// A download that fits in one partition is served by the initial request, which is +/// never routed, so the layout is not worth fetching. +#[tokio::test] +async fn test_download_layout_aware_routing_skips_layout_for_complete_initial_range( +) -> Result<(), Box> { + const DATA: [u8; 8] = [10, 11, 12, 13, 14, 15, 16, 17]; + + let request_count = Arc::new(AtomicUsize::new(0)); + let count_capture = request_count.clone(); + let mock_client = Arc::new(MockHttpClient::new(move |request| { + assert_eq!(0, count_capture.fetch_add(1, Ordering::SeqCst)); + assert!(!request + .url() + .query() + .is_some_and(|query| query.contains("comp=layout"))); + assert_eq!( + Some("bytes=0-7"), + request.headers().get_optional_str(&"range".into()) + ); + + let mut headers = Headers::new(); + headers.insert("content-range", "bytes 0-7/8"); + headers.insert("content-length", DATA.len().to_string()); + headers.insert("etag", ETAG); + headers.insert("x-ms-download-hint", "layout"); + async move { + Ok(AsyncRawResponse::from_bytes( + StatusCode::PartialContent, + headers, + Bytes::from_static(&DATA), + )) + } + .boxed() + })); + + let blob_client = blob_client_with( + Transport::new(mock_client), + "https://acct.blob.core.windows.net/container/blob", + )?; + + let body = blob_client + .download(Some(BlobClientDownloadOptions { + layout_aware_routing: LayoutAwareRouting::Enabled, + partition_size: Some(NonZero::new(DATA.len()).unwrap()), + ..Default::default() + })) + .await? + .body + .collect() + .await?; + + assert_eq!(&body[..], &DATA[..]); + assert_eq!(request_count.load(Ordering::SeqCst), 1); + Ok(()) +} + +/// A `BlobClient` reached through a container client shares that client's pipeline, so +/// this covers the routing policy being installed for every client rather than only the +/// ones built by `BlobClient::new`. +#[tokio::test] +async fn test_download_layout_aware_routing_from_container_client() -> Result<(), Box> { + let seen = Arc::new(Mutex::new(Vec::::new())); + let transport = layout_mock_transport(seen.clone(), LAYOUT, LayoutHint::Advertised); + + let container_client = BlobContainerClient::new( + Url::parse("https://acct.blob.core.windows.net/container")?, + None, + Some(BlobContainerClientOptions { + client_options: ClientOptions { + transport: Some(transport), + ..Default::default() + }, + ..Default::default() + }), + )?; + let blob_client = container_client.blob_client("blob"); + + let body = blob_client + .download(Some(BlobClientDownloadOptions { + layout_aware_routing: LayoutAwareRouting::Enabled, + partition_size: Some(NonZero::new(4).unwrap()), + parallel: Some(NonZero::new(2).unwrap()), + ..Default::default() + })) + .await? + .body + .collect() + .await?; + + assert_eq!(&body[..], &BLOB_DATA[..]); + let seen = seen.lock().unwrap(); + let routed: HashMap = seen + .iter() + .filter(|request| !request.is_layout) + .map(|request| { + ( + request.range.clone().expect("range header"), + request.url.host_str().unwrap_or_default().to_owned(), + ) + }) + .collect(); + assert_eq!( + routed.get("bytes=4-7").map(String::as_str), + Some("epb.blob.core.windows.net"), + "a blob client from a container client should still route" + ); + assert_eq!( + routed.get("bytes=8-11").map(String::as_str), + Some("epc.blob.core.windows.net") + ); + Ok(()) +} + +/// A caller can fetch the layout themselves and pin a download to the endpoint serving +/// the range they want, without the SDK looking up the layout on their behalf. +#[tokio::test] +async fn test_download_layout_endpoint_selected_by_caller() -> Result<(), Box> { + let seen = Arc::new(Mutex::new(Vec::::new())); + // No layout hint, so nothing but the caller's endpoint can cause a rewrite. + let transport = layout_mock_transport(seen.clone(), LAYOUT, LayoutHint::Absent); + let blob_client = blob_client_with( + transport, + "https://acct.blob.core.windows.net/container/blob", + )?; + + let mut pages = blob_client.get_layout(None)?; + let layout = pages + .try_next() + .await? + .expect("a layout page") + .into_model()?; + let endpoint = endpoint_for_offset(&layout, 4).expect("an endpoint covering offset 4"); + assert_eq!(endpoint, "epb.blob.core.windows.net:443"); + + let body = blob_client + .download(Some(BlobClientDownloadOptions { + layout_endpoint: Some(endpoint), + range: Some((4u64..8).into()), + partition_size: Some(NonZero::new(4).unwrap()), + ..Default::default() + })) + .await? + .body + .collect() + .await?; + + assert_eq!(&body[..], &BLOB_DATA[4..8]); + let seen = seen.lock().unwrap(); + let layout_requests: Vec<&ObservedRequest> = + seen.iter().filter(|request| request.is_layout).collect(); + let data_requests: Vec<&ObservedRequest> = + seen.iter().filter(|request| !request.is_layout).collect(); + + assert_eq!( + layout_requests.len(), + 1, + "only the caller's own get_layout call should be issued" + ); + assert_eq!(layout_requests[0].url.host_str(), Some(ACCOUNT_HOST)); + + assert_eq!(data_requests.len(), 1); + assert_eq!(data_requests[0].range.as_deref(), Some("bytes=4-7")); + assert_eq!( + data_requests[0].url.host_str(), + Some("epb.blob.core.windows.net"), + "the download should be sent to the caller's endpoint" + ); + assert_eq!( + data_requests[0].original_host.as_deref(), + Some(ACCOUNT_HOST), + "the account authority should be preserved as the Host header" + ); + Ok(()) +} + +/// A caller-supplied endpoint takes precedence over automatic routing: the layout is +/// never fetched and every request the download issues is pinned, including the first. +#[tokio::test] +async fn test_download_layout_endpoint_overrides_layout_aware_routing() -> Result<(), Box> +{ + let seen = Arc::new(Mutex::new(Vec::::new())); + let transport = layout_mock_transport(seen.clone(), LAYOUT, LayoutHint::Advertised); + let blob_client = blob_client_with( + transport, + "https://acct.blob.core.windows.net/container/blob", + )?; + + let body = blob_client + .download(Some(BlobClientDownloadOptions { + layout_endpoint: Some("epa.blob.core.windows.net:443".into()), + // Left at its default of `Enabled` to show the explicit endpoint wins. + partition_size: Some(NonZero::new(4).unwrap()), + parallel: Some(NonZero::new(2).unwrap()), + ..Default::default() + })) + .await? + .body + .collect() + .await?; + + assert_eq!(&body[..], &BLOB_DATA[..]); + let seen = seen.lock().unwrap(); + assert!( + !seen.iter().any(|request| request.is_layout), + "an explicit endpoint should suppress the automatic layout lookup" + ); + assert_eq!(seen.len(), 3); + for request in seen.iter() { + assert_eq!( + request.url.host_str(), + Some("epa.blob.core.windows.net"), + "every request, including the first, should be pinned" + ); + assert_eq!(request.original_host.as_deref(), Some(ACCOUNT_HOST)); + } + Ok(()) +} + +fn live_layout_blob_url() -> Option { + let account = std::env::var("AZURE_STORAGE_ACCOUNT_NAME").ok()?; + let account = account.trim().trim_matches('"').trim(); + if account.is_empty() { + return None; + } + Url::parse(&format!( + "https://{account}.blob.core.windows.net/layout-routing-test/routing-test-blob" + )) + .ok() +} + +#[derive(Debug)] +struct RoutingObserver { + seen: Arc>>, +} + +#[derive(Debug)] +struct RoutedRequest { + is_layout: bool, + range: Option, + sent_to: Option, + original_host: Option, + download_hint: Option, + status: StatusCode, +} + +#[async_trait] +impl Policy for RoutingObserver { + async fn send( + &self, + ctx: &Context, + request: &mut Request, + next: &[Arc], + ) -> PolicyResult { + let is_layout = request + .url() + .query() + .is_some_and(|query| query.contains("comp=layout")); + let range = request + .headers() + .get_optional_str(&"range".into()) + .map(str::to_owned); + let sent_to = request.url().host_str().map(str::to_owned); + let original_host = request + .headers() + .get_optional_str(&"host".into()) + .map(str::to_owned); + let response = next[0].send(ctx, request, &next[1..]).await?; + let download_hint = response + .headers() + .get_optional_str(&"x-ms-download-hint".into()) + .map(str::to_owned); + self.seen.lock().unwrap().push(RoutedRequest { + is_layout, + range, + sent_to, + original_host, + download_hint, + status: response.status(), + }); + Ok(response) + } +} + +#[recorded::test(live)] +async fn test_download_layout_aware_routing() -> Result<(), Box> { + let Some(url) = live_layout_blob_url() else { + eprintln!( + "skipping test_download_layout_aware_routing: set AZURE_STORAGE_ACCOUNT_NAME to run it" + ); + return Ok(()); + }; + let credential: Arc = DeveloperToolsCredential::new(None)?; + let account_host = url.host_str().unwrap_or_default().to_owned(); + let observed = Arc::new(Mutex::new(Vec::::new())); + let observer: Arc = Arc::new(RoutingObserver { + seen: observed.clone(), + }); + + let mut container_url = url.clone(); + container_url + .path_segments_mut() + .expect("blob URL must be a base") + .pop(); + let container_client = BlobContainerClient::new(container_url, Some(credential.clone()), None)?; + if let Err(error) = container_client.create(None).await { + if error.http_status() != Some(StatusCode::Conflict) { + return Err(error.into()); + } + } + + let blob_client = BlobClient::new( + url, + Some(credential), + Some(BlobClientOptions { + client_options: ClientOptions { + per_try_policies: vec![observer], + ..Default::default() + }, + ..Default::default() + }), + )?; + + const BLOB_SIZE: usize = 64 * 1024 * 1024; + let data: Vec = (0..BLOB_SIZE).map(|index| (index % 251) as u8).collect(); + blob_client + .upload(RequestContent::from(data.clone()), None) + .await?; + observed.lock().unwrap().clear(); + + let body = blob_client + .download(Some(BlobClientDownloadOptions { + layout_aware_routing: LayoutAwareRouting::Enabled, + partition_size: Some(NonZero::new(16 * 1024 * 1024).unwrap()), + parallel: Some(NonZero::new(4).unwrap()), + ..Default::default() + })) + .await? + .body + .collect() + .await?; + assert_eq!(&body[..], &data[..]); + + let mut buffer = vec![0u8; data.len()]; + let result = blob_client + .download_into( + &mut buffer, + Some(BlobClientDownloadOptions { + layout_aware_routing: LayoutAwareRouting::Enabled, + partition_size: Some(NonZero::new(16 * 1024 * 1024).unwrap()), + parallel: Some(NonZero::new(4).unwrap()), + ..Default::default() + }), + ) + .await?; + assert_eq!(result.len, data.len()); + assert_eq!(&buffer[..], &data[..]); + + let observed = observed.lock().unwrap(); + let layout_calls = observed.iter().filter(|request| request.is_layout).count(); + let data_chunks = observed.iter().filter(|request| !request.is_layout).count(); + let routed: Vec<&RoutedRequest> = observed + .iter() + .filter(|request| !request.is_layout && request.original_host.is_some()) + .collect(); + + eprintln!(); + eprintln!("=== layout-aware routing ======================================"); + eprintln!(" account host : {account_host}"); + eprintln!(" Get Blob Layout calls : {layout_calls}"); + eprintln!( + " data chunks : {data_chunks} ({} routed)", + routed.len() + ); + for request in observed.iter() { + let kind = if request.is_layout { "layout" } else { "data" }; + let range = request.range.as_deref().unwrap_or("(whole blob)"); + let sent_to = request.sent_to.as_deref().unwrap_or("?"); + let hint = request.download_hint.as_deref().unwrap_or("-"); + let routed = if request.original_host.is_some() { + "yes" + } else { + "no" + }; + let status = request.status; + eprintln!(" {kind:<6} {status:<3} {range:<24} {routed:<7} {hint:<7} {sent_to}"); + } + + if routed.is_empty() { + eprintln!( + "NOTE: no chunk was rewritten - the account returned no layout hint for this blob" + ); + } else { + for request in &routed { + assert_eq!( + request.original_host.as_deref(), + Some(account_host.as_str()), + "a rewritten chunk did not preserve the account Host header" + ); + } + } + + Ok(()) +}