diff --git a/sdk/storage/azure_storage_blob/src/blob_layout.rs b/sdk/storage/azure_storage_blob/src/blob_layout.rs new file mode 100644 index 00000000000..5160676d4b7 --- /dev/null +++ b/sdk/storage/azure_storage_blob/src/blob_layout.rs @@ -0,0 +1,1269 @@ +// 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, + time::{Duration, Instant}, +}; + +use async_trait::async_trait; +use azure_core::{ + async_runtime::get_async_runtime, + error::ErrorKind, + http::{ + policies::{Policy, PolicyResult}, + Context, Etag, Request, StatusCode, Url, + }, + Error, Result, +}; +use futures::{ + future::{select, Either}, + lock::Mutex, + StreamExt as _, +}; + +use crate::generated::{ + clients::BlobClient, + models::{BlobClientGetLayoutOptions, 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 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 +/// 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. +/// +/// 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}"), + (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: &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.get_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), + } +} + +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: BlobClientGetLayoutOptions<'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: BlobClientGetLayoutOptions<'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: BlobClientGetLayoutOptions<'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::*; + 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 _; + use std::sync::atomic::{AtomicUsize, Ordering}; + + 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)) + ); + // 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 "), + 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); + } + + #[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#" + + + + + + + + + +"#; + + 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(), + &BlobClientGetLayoutOptions::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(), + &BlobClientGetLayoutOptions { + 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_get_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.get_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(), + &BlobClientGetLayoutOptions::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(), + &BlobClientGetLayoutOptions::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(), + &BlobClientGetLayoutOptions::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(), + &BlobClientGetLayoutOptions::default(), + ) + .await; + assert!(result.is_err()); + } + + const LAYOUT_V1: &[u8] = br#" + + + + + + + +"#; + + const LAYOUT_V2: &[u8] = br#" + + + + + + + +"#; + + async fn seed_layout( + client: &BlobClient, + layout_options: &BlobClientGetLayoutOptions<'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 = BlobClientGetLayoutOptions::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 = BlobClientGetLayoutOptions::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 = BlobClientGetLayoutOptions::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 = 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), + 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 = BlobClientGetLayoutOptions::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 5bf60b8db1e..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,13 +4,15 @@ pub use crate::generated::clients::{BlobClient, BlobClientOptions}; use crate::{ + blob_layout::{fetch_layout, LayoutCache, LayoutEndpoint}, generated::{ - clients::BlobClient as GeneratedBlobClient, models::BlobClientDownloadInternalOptions, + clients::BlobClient as GeneratedBlobClient, + models::{BlobClientDownloadInternalOptions, BlobClientGetLayoutOptions}, }, models::{ BlobClientDownloadIntoResult, BlobClientDownloadOptions, BlobClientDownloadResult, BlobClientUploadOptions, BlobClientUploadResult, BlobDownloadProperties, HttpRange, - StorageErrorCode, + LayoutAwareRouting, StorageErrorCode, }, partitioned_transfer::{self, PartitionedDownloadBehavior}, AppendBlobClient, BlockBlobClient, PageBlobClient, @@ -20,12 +22,16 @@ use azure_core::{ credentials::TokenCredential, error::ErrorKind, http::{ + headers::Headers, policies::{auth::BearerTokenAuthorizationPolicy, Policy}, AsyncRawResponse, Etag, NoFormat, Pipeline, RequestContent, StatusCode, Url, UrlExt, }, tracing, Bytes, Result, }; -use std::{ops::Range, sync::Arc}; +use std::{ + ops::Range, + sync::{Arc, OnceLock}, +}; impl BlobClient { /// Creates a new BlobClient from a blob URL. @@ -194,14 +200,22 @@ 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 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 = partitioned_transfer::download(range, parallel, partition_size, Arc::new(behavior)) .await?; @@ -238,14 +252,22 @@ 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 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( buffer, range, @@ -302,13 +324,68 @@ impl BlobClient { } struct BlobClientDownloadBehavior<'a> { - client: GeneratedBlobClient, + 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>, } impl<'a> BlobClientDownloadBehavior<'a> { - fn new(client: GeneratedBlobClient, options: BlobClientDownloadInternalOptions<'a>) -> Self { - Self { client, options } + fn new( + 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, + 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() + } } } @@ -320,6 +397,13 @@ impl PartitionedDownloadBehavior for BlobClientDownloadBehavior<'_> { etag_lock: Option, ) -> Result { let mut opt = self.options.clone(); + 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 { opt.if_match = Some(etag); @@ -333,4 +417,32 @@ impl PartitionedDownloadBehavior for BlobClientDownloadBehavior<'_> { .await .map(AsyncRawResponse::from) } + + async fn prepare(&self, initial_headers: &Headers, etag_lock: Option<&Etag>) -> Result<()> { + 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") { + let _ = self.layout_cache.set(None); + return Ok(()); + } + 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 = fetch_layout(&self.client, &context, &layout_options) + .await? + .map(|prefetch| { + LayoutCache::new( + Arc::clone(&self.client), + layout_options, + Arc::new(prefetch.layout), + ) + }); + let _ = self.layout_cache.set(cache); + Ok(()) + } } 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/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..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 @@ -11,16 +11,15 @@ use crate::{ 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, + BlobClientGetAccountInfoResult, BlobClientGetLayoutOptions, BlobClientGetPropertiesOptions, + BlobClientGetPropertiesResult, BlobClientGetTagsOptions, BlobClientReleaseLeaseOptions, + BlobClientReleaseLeaseResult, BlobClientRenewLeaseOptions, BlobClientRenewLeaseResult, + BlobClientSetImmutabilityPolicyOptions, BlobClientSetLegalHoldOptions, + BlobClientSetMetadataOptions, BlobClientSetPropertiesOptions, BlobClientSetTagsOptions, + BlobClientSetTierOptions, BlobClientStartCopyFromUrlOptions, + BlobClientStartCopyFromUrlResult, BlobClientUndeleteOptions, BlobLayout, BlobTags, }, - SessionOptions, + // SessionOptions, }; use azure_core::{ base64, @@ -41,7 +40,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, } @@ -51,7 +50,7 @@ pub struct BlobClientOptions { /// Allows customization of the client. pub client_options: ClientOptions, /// Options for session token authentication. - pub session_options: Option, + // pub session_options: Option, /// Specifies the version of the operation to use for this request. pub version: String, } @@ -831,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. /// @@ -1045,196 +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(); - 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 @@ -1978,7 +1980,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/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/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 a43712cb80d..5c28e4b7965 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,26 @@ 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. 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 { + /// Never route range requests based on the blob's layout. All requests are sent + /// to the client's configured endpoint. + Disabled, + + /// Use locality-aware routing. The blob's layout is fetched and each range + /// download is routed to the endpoint that serves it. + #[default] + Enabled, +} + /// Options to be passed to `BlobClient::download()` #[derive(Clone, Default, SafeDebug)] pub struct BlobClientDownloadOptions<'a> { @@ -46,6 +66,34 @@ 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::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/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/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(()) +} 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: