LCOV - code coverage report
Current view: top level - libs/remote_storage/src - azure_blob.rs (source / functions) Coverage Total Hit
Test: 6df3fc19ec669bcfbbf9aba41d1338898d24eaa0.info Lines: 78.9 % 617 487
Test Date: 2025-03-12 18:28:53 Functions: 53.2 % 79 42

            Line data    Source code
       1              : //! Azure Blob Storage wrapper
       2              : 
       3              : use std::borrow::Cow;
       4              : use std::collections::HashMap;
       5              : use std::fmt::Display;
       6              : use std::num::NonZeroU32;
       7              : use std::pin::Pin;
       8              : use std::str::FromStr;
       9              : use std::sync::Arc;
      10              : use std::time::{Duration, SystemTime};
      11              : use std::{env, io};
      12              : 
      13              : use anyhow::{Context, Result};
      14              : use azure_core::request_options::{IfMatchCondition, MaxResults, Metadata, Range};
      15              : use azure_core::{Continuable, HttpClient, RetryOptions, TransportOptions};
      16              : use azure_storage::StorageCredentials;
      17              : use azure_storage_blobs::blob::CopyStatus;
      18              : use azure_storage_blobs::blob::operations::GetBlobBuilder;
      19              : use azure_storage_blobs::prelude::{ClientBuilder, ContainerClient};
      20              : use bytes::Bytes;
      21              : use futures::FutureExt;
      22              : use futures::future::Either;
      23              : use futures::stream::Stream;
      24              : use futures_util::{StreamExt, TryStreamExt};
      25              : use http_types::{StatusCode, Url};
      26              : use scopeguard::ScopeGuard;
      27              : use tokio_util::sync::CancellationToken;
      28              : use tracing::debug;
      29              : use utils::backoff;
      30              : use utils::backoff::exponential_backoff_duration_seconds;
      31              : 
      32              : use super::REMOTE_STORAGE_PREFIX_SEPARATOR;
      33              : use crate::config::AzureConfig;
      34              : use crate::error::Cancelled;
      35              : use crate::metrics::{AttemptOutcome, RequestKind, start_measuring_requests};
      36              : use crate::{
      37              :     ConcurrencyLimiter, Download, DownloadError, DownloadKind, DownloadOpts, Listing, ListingMode,
      38              :     ListingObject, RemotePath, RemoteStorage, StorageMetadata, TimeTravelError, TimeoutOrCancel,
      39              : };
      40              : 
      41              : pub struct AzureBlobStorage {
      42              :     client: ContainerClient,
      43              :     container_name: String,
      44              :     prefix_in_container: Option<String>,
      45              :     max_keys_per_list_response: Option<NonZeroU32>,
      46              :     concurrency_limiter: ConcurrencyLimiter,
      47              :     // Per-request timeout. Accessible for tests.
      48              :     pub timeout: Duration,
      49              : 
      50              :     // Alternative timeout used for metadata objects which are expected to be small
      51              :     pub small_timeout: Duration,
      52              : }
      53              : 
      54              : impl AzureBlobStorage {
      55           10 :     pub fn new(
      56           10 :         azure_config: &AzureConfig,
      57           10 :         timeout: Duration,
      58           10 :         small_timeout: Duration,
      59           10 :     ) -> Result<Self> {
      60           10 :         debug!(
      61            0 :             "Creating azure remote storage for azure container {}",
      62              :             azure_config.container_name
      63              :         );
      64              : 
      65              :         // Use the storage account from the config by default, fall back to env var if not present.
      66           10 :         let account = azure_config.storage_account.clone().unwrap_or_else(|| {
      67           10 :             env::var("AZURE_STORAGE_ACCOUNT").expect("missing AZURE_STORAGE_ACCOUNT")
      68           10 :         });
      69              : 
      70              :         // If the `AZURE_STORAGE_ACCESS_KEY` env var has an access key, use that,
      71              :         // otherwise try the token based credentials.
      72           10 :         let credentials = if let Ok(access_key) = env::var("AZURE_STORAGE_ACCESS_KEY") {
      73           10 :             StorageCredentials::access_key(account.clone(), access_key)
      74              :         } else {
      75            0 :             let token_credential = azure_identity::create_default_credential()
      76            0 :                 .context("trying to obtain Azure default credentials")?;
      77            0 :             StorageCredentials::token_credential(token_credential)
      78              :         };
      79              : 
      80           10 :         let builder = ClientBuilder::new(account, credentials)
      81           10 :             // we have an outer retry
      82           10 :             .retry(RetryOptions::none())
      83           10 :             // Customize transport to configure conneciton pooling
      84           10 :             .transport(TransportOptions::new(Self::reqwest_client(
      85           10 :                 azure_config.conn_pool_size,
      86           10 :             )));
      87           10 : 
      88           10 :         let client = builder.container_client(azure_config.container_name.to_owned());
      89              : 
      90           10 :         let max_keys_per_list_response =
      91           10 :             if let Some(limit) = azure_config.max_keys_per_list_response {
      92              :                 Some(
      93            4 :                     NonZeroU32::new(limit as u32)
      94            4 :                         .ok_or_else(|| anyhow::anyhow!("max_keys_per_list_response can't be 0"))?,
      95              :                 )
      96              :             } else {
      97            6 :                 None
      98              :             };
      99              : 
     100           10 :         Ok(AzureBlobStorage {
     101           10 :             client,
     102           10 :             container_name: azure_config.container_name.to_owned(),
     103           10 :             prefix_in_container: azure_config.prefix_in_container.to_owned(),
     104           10 :             max_keys_per_list_response,
     105           10 :             concurrency_limiter: ConcurrencyLimiter::new(azure_config.concurrency_limit.get()),
     106           10 :             timeout,
     107           10 :             small_timeout,
     108           10 :         })
     109           10 :     }
     110              : 
     111           10 :     fn reqwest_client(conn_pool_size: usize) -> Arc<dyn HttpClient> {
     112           10 :         let client = reqwest::ClientBuilder::new()
     113           10 :             .pool_max_idle_per_host(conn_pool_size)
     114           10 :             .build()
     115           10 :             .expect("failed to build `reqwest` client");
     116           10 :         Arc::new(client)
     117           10 :     }
     118              : 
     119          237 :     pub fn relative_path_to_name(&self, path: &RemotePath) -> String {
     120          237 :         assert_eq!(std::path::MAIN_SEPARATOR, REMOTE_STORAGE_PREFIX_SEPARATOR);
     121          237 :         let path_string = path.get_path().as_str();
     122          237 :         match &self.prefix_in_container {
     123          237 :             Some(prefix) => {
     124          237 :                 if prefix.ends_with(REMOTE_STORAGE_PREFIX_SEPARATOR) {
     125          237 :                     prefix.clone() + path_string
     126              :                 } else {
     127            0 :                     format!("{prefix}{REMOTE_STORAGE_PREFIX_SEPARATOR}{path_string}")
     128              :                 }
     129              :             }
     130            0 :             None => path_string.to_string(),
     131              :         }
     132          237 :     }
     133              : 
     134          250 :     fn name_to_relative_path(&self, key: &str) -> RemotePath {
     135          250 :         let relative_path =
     136          250 :             match key.strip_prefix(self.prefix_in_container.as_deref().unwrap_or_default()) {
     137          250 :                 Some(stripped) => stripped,
     138              :                 // we rely on Azure to return properly prefixed paths
     139              :                 // for requests with a certain prefix
     140            0 :                 None => panic!(
     141            0 :                     "Key {key} does not start with container prefix {:?}",
     142            0 :                     self.prefix_in_container
     143            0 :                 ),
     144              :             };
     145          250 :         RemotePath(
     146          250 :             relative_path
     147          250 :                 .split(REMOTE_STORAGE_PREFIX_SEPARATOR)
     148          250 :                 .collect(),
     149          250 :         )
     150          250 :     }
     151              : 
     152           11 :     async fn download_for_builder(
     153           11 :         &self,
     154           11 :         builder: GetBlobBuilder,
     155           11 :         timeout: Duration,
     156           11 :         cancel: &CancellationToken,
     157           11 :     ) -> Result<Download, DownloadError> {
     158           11 :         let kind = RequestKind::Get;
     159              : 
     160           11 :         let _permit = self.permit(kind, cancel).await?;
     161           11 :         let cancel_or_timeout = crate::support::cancel_or_timeout(self.timeout, cancel.clone());
     162           11 :         let cancel_or_timeout_ = crate::support::cancel_or_timeout(self.timeout, cancel.clone());
     163           11 : 
     164           11 :         let mut etag = None;
     165           11 :         let mut last_modified = None;
     166           11 :         let mut metadata = HashMap::new();
     167           11 : 
     168           11 :         let started_at = start_measuring_requests(kind);
     169           11 : 
     170           11 :         let download = async {
     171           11 :             let response = builder
     172           11 :                 // convert to concrete Pageable
     173           11 :                 .into_stream()
     174           11 :                 // convert to TryStream
     175           11 :                 .into_stream()
     176           11 :                 .map_err(to_download_error);
     177           11 : 
     178           11 :             // apply per request timeout
     179           11 :             let response = tokio_stream::StreamExt::timeout(response, timeout);
     180           11 : 
     181           11 :             // flatten
     182           11 :             let response = response.map(|res| match res {
     183           11 :                 Ok(res) => res,
     184            0 :                 Err(_elapsed) => Err(DownloadError::Timeout),
     185           11 :             });
     186           11 : 
     187           11 :             let mut response = Box::pin(response);
     188              : 
     189           11 :             let Some(part) = response.next().await else {
     190            0 :                 return Err(DownloadError::Other(anyhow::anyhow!(
     191            0 :                     "Azure GET response contained no response body"
     192            0 :                 )));
     193              :             };
     194           11 :             let part = part?;
     195            9 :             if etag.is_none() {
     196            9 :                 etag = Some(part.blob.properties.etag);
     197            9 :             }
     198            9 :             if last_modified.is_none() {
     199            9 :                 last_modified = Some(part.blob.properties.last_modified.into());
     200            9 :             }
     201            9 :             if let Some(blob_meta) = part.blob.metadata {
     202            0 :                 metadata.extend(blob_meta.iter().map(|(k, v)| (k.to_owned(), v.to_owned())));
     203            0 :             }
     204              : 
     205              :             // unwrap safety: if these were None, bufs would be empty and we would have returned an error already
     206            9 :             let etag = etag.unwrap();
     207            9 :             let last_modified = last_modified.unwrap();
     208            9 : 
     209            9 :             let tail_stream = response
     210            9 :                 .map(|part| match part {
     211            0 :                     Ok(part) => Either::Left(part.data.map(|r| r.map_err(io::Error::other))),
     212            0 :                     Err(e) => {
     213            0 :                         Either::Right(futures::stream::once(async { Err(io::Error::other(e)) }))
     214              :                     }
     215            9 :                 })
     216            9 :                 .flatten();
     217            9 :             let stream = part
     218            9 :                 .data
     219            9 :                 .map(|r| r.map_err(io::Error::other))
     220            9 :                 .chain(sync_wrapper::SyncStream::new(tail_stream));
     221            9 :             //.chain(SyncStream::from_pin(Box::pin(tail_stream)));
     222            9 : 
     223            9 :             let download_stream = crate::support::DownloadStream::new(cancel_or_timeout_, stream);
     224            9 : 
     225            9 :             Ok(Download {
     226            9 :                 download_stream: Box::pin(download_stream),
     227            9 :                 etag,
     228            9 :                 last_modified,
     229            9 :                 metadata: Some(StorageMetadata(metadata)),
     230            9 :             })
     231            0 :         };
     232              : 
     233           11 :         let download = tokio::select! {
     234           11 :             bufs = download => bufs,
     235           11 :             cancel_or_timeout = cancel_or_timeout => match cancel_or_timeout {
     236            0 :                 TimeoutOrCancel::Timeout => return Err(DownloadError::Timeout),
     237            0 :                 TimeoutOrCancel::Cancel => return Err(DownloadError::Cancelled),
     238              :             },
     239              :         };
     240           11 :         let started_at = ScopeGuard::into_inner(started_at);
     241           11 :         let outcome = match &download {
     242            9 :             Ok(_) => AttemptOutcome::Ok,
     243              :             // At this level in the stack 404 and 304 responses do not indicate an error.
     244              :             // There's expected cases when a blob may not exist or hasn't been modified since
     245              :             // the last get (e.g. probing for timeline indices and heatmap downloads).
     246              :             // Callers should handle errors if they are unexpected.
     247            2 :             Err(DownloadError::NotFound | DownloadError::Unmodified) => AttemptOutcome::Ok,
     248            0 :             Err(_) => AttemptOutcome::Err,
     249              :         };
     250           11 :         crate::metrics::BUCKET_METRICS
     251           11 :             .req_seconds
     252           11 :             .observe_elapsed(kind, outcome, started_at);
     253           11 :         download
     254            0 :     }
     255              : 
     256          227 :     async fn permit(
     257          227 :         &self,
     258          227 :         kind: RequestKind,
     259          227 :         cancel: &CancellationToken,
     260          227 :     ) -> Result<tokio::sync::SemaphorePermit<'_>, Cancelled> {
     261          227 :         let acquire = self.concurrency_limiter.acquire(kind);
     262          227 : 
     263          227 :         tokio::select! {
     264          227 :             permit = acquire => Ok(permit.expect("never closed")),
     265          227 :             _ = cancel.cancelled() => Err(Cancelled),
     266              :         }
     267            0 :     }
     268              : 
     269            0 :     pub fn container_name(&self) -> &str {
     270            0 :         &self.container_name
     271            0 :     }
     272              : }
     273              : 
     274            0 : fn to_azure_metadata(metadata: StorageMetadata) -> Metadata {
     275            0 :     let mut res = Metadata::new();
     276            0 :     for (k, v) in metadata.0.into_iter() {
     277            0 :         res.insert(k, v);
     278            0 :     }
     279            0 :     res
     280            0 : }
     281              : 
     282            3 : fn to_download_error(error: azure_core::Error) -> DownloadError {
     283            3 :     if let Some(http_err) = error.as_http_error() {
     284            3 :         match http_err.status() {
     285            1 :             StatusCode::NotFound => DownloadError::NotFound,
     286            2 :             StatusCode::NotModified => DownloadError::Unmodified,
     287            0 :             StatusCode::BadRequest => DownloadError::BadInput(anyhow::Error::new(error)),
     288            0 :             _ => DownloadError::Other(anyhow::Error::new(error)),
     289              :         }
     290              :     } else {
     291            0 :         DownloadError::Other(error.into())
     292              :     }
     293            3 : }
     294              : 
     295              : impl RemoteStorage for AzureBlobStorage {
     296           26 :     fn list_streaming(
     297           26 :         &self,
     298           26 :         prefix: Option<&RemotePath>,
     299           26 :         mode: ListingMode,
     300           26 :         max_keys: Option<NonZeroU32>,
     301           26 :         cancel: &CancellationToken,
     302           26 :     ) -> impl Stream<Item = Result<Listing, DownloadError>> {
     303           26 :         // get the passed prefix or if it is not set use prefix_in_bucket value
     304           26 :         let list_prefix = prefix.map(|p| self.relative_path_to_name(p)).or_else(|| {
     305           10 :             self.prefix_in_container.clone().map(|mut s| {
     306           10 :                 if !s.ends_with(REMOTE_STORAGE_PREFIX_SEPARATOR) {
     307            0 :                     s.push(REMOTE_STORAGE_PREFIX_SEPARATOR);
     308           10 :                 }
     309           10 :                 s
     310           10 :             })
     311           26 :         });
     312           26 : 
     313           26 :         async_stream::stream! {
     314           26 :             let _permit = self.permit(RequestKind::List, cancel).await?;
     315           26 : 
     316           26 :             let mut builder = self.client.list_blobs();
     317           26 : 
     318           26 :             if let ListingMode::WithDelimiter = mode {
     319           26 :                 builder = builder.delimiter(REMOTE_STORAGE_PREFIX_SEPARATOR.to_string());
     320           26 :             }
     321           26 : 
     322           26 :             if let Some(prefix) = list_prefix {
     323           26 :                 builder = builder.prefix(Cow::from(prefix.to_owned()));
     324           26 :             }
     325           26 : 
     326           26 :             if let Some(limit) = self.max_keys_per_list_response {
     327           26 :                 builder = builder.max_results(MaxResults::new(limit));
     328           26 :             }
     329           26 : 
     330           26 :             let mut next_marker = None;
     331           26 : 
     332           26 :             let mut timeout_try_cnt = 1;
     333           26 : 
     334           26 :             'outer: loop {
     335           26 :                 let mut builder = builder.clone();
     336           26 :                 if let Some(marker) = next_marker.clone() {
     337           26 :                     builder = builder.marker(marker);
     338           26 :                 }
     339           26 :                 // Azure Blob Rust SDK does not expose the list blob API directly. Users have to use
     340           26 :                 // their pageable iterator wrapper that returns all keys as a stream. We want to have
     341           26 :                 // full control of paging, and therefore we only take the first item from the stream.
     342           26 :                 let mut response_stream = builder.into_stream();
     343           26 :                 let response = response_stream.next();
     344           26 :                 // Timeout mechanism: Azure client will sometimes stuck on a request, but retrying that request
     345           26 :                 // would immediately succeed. Therefore, we use exponential backoff timeout to retry the request.
     346           26 :                 // (Usually, exponential backoff is used to determine the sleep time between two retries.) We
     347           26 :                 // start with 10.0 second timeout, and double the timeout for each failure, up to 5 failures.
     348           26 :                 // timeout = min(5 * (1.0+1.0)^n, self.timeout).
     349           26 :                 let this_timeout = (5.0 * exponential_backoff_duration_seconds(timeout_try_cnt, 1.0, self.timeout.as_secs_f64())).min(self.timeout.as_secs_f64());
     350           26 :                 let response = tokio::time::timeout(Duration::from_secs_f64(this_timeout), response);
     351           26 :                 let response = response.map(|res| {
     352           45 :                     match res {
     353           45 :                         Ok(Some(Ok(res))) => Ok(Some(res)),
     354           26 :                         Ok(Some(Err(e)))  => Err(to_download_error(e)),
     355           26 :                         Ok(None) => Ok(None),
     356           26 :                         Err(_elasped) => Err(DownloadError::Timeout),
     357           26 :                     }
     358           26 :                 });
     359           26 :                 let mut max_keys = max_keys.map(|mk| mk.get());
     360           26 :                 let next_item = tokio::select! {
     361           26 :                     op = response => op,
     362           26 :                     _ = cancel.cancelled() => Err(DownloadError::Cancelled),
     363           26 :                 };
     364           26 : 
     365           26 :                 if let Err(DownloadError::Timeout) = &next_item {
     366           26 :                     timeout_try_cnt += 1;
     367           26 :                     if timeout_try_cnt <= 5 {
     368           26 :                         continue;
     369           26 :                     }
     370           26 :                 }
     371           26 : 
     372           26 :                 let next_item = next_item?;
     373           26 : 
     374           26 :                 // Log a warning if we saw two timeouts in a row before a successful request
     375           26 :                 if timeout_try_cnt > 2 {
     376           26 :                     tracing::warn!("Azure Blob Storage list timed out and succeeded after {} tries", timeout_try_cnt);
     377           26 :                 }
     378           26 :                 timeout_try_cnt = 1;
     379           26 : 
     380           26 :                 let Some(entry) = next_item else {
     381           26 :                     // The list is complete, so yield it.
     382           26 :                     break;
     383           26 :                 };
     384           26 : 
     385           26 :                 let mut res = Listing::default();
     386           26 :                 next_marker = entry.continuation();
     387           26 :                 let prefix_iter = entry
     388           26 :                     .blobs
     389           26 :                     .prefixes()
     390           52 :                     .map(|prefix| self.name_to_relative_path(&prefix.name));
     391           26 :                 res.prefixes.extend(prefix_iter);
     392           26 : 
     393           26 :                 let blob_iter = entry
     394           26 :                     .blobs
     395           26 :                     .blobs()
     396          198 :                     .map(|k| ListingObject{
     397          198 :                         key: self.name_to_relative_path(&k.name),
     398          198 :                         last_modified: k.properties.last_modified.into(),
     399          198 :                         size: k.properties.content_length,
     400          198 :                     }
     401           26 :                 );
     402           26 : 
     403           26 :                 for key in blob_iter {
     404           26 :                     res.keys.push(key);
     405           26 : 
     406           26 :                     if let Some(mut mk) = max_keys {
     407           26 :                         assert!(mk > 0);
     408           26 :                         mk -= 1;
     409           26 :                         if mk == 0 {
     410           26 :                             yield Ok(res); // limit reached
     411           26 :                             break 'outer;
     412           26 :                         }
     413           26 :                         max_keys = Some(mk);
     414           26 :                     }
     415           26 :                 }
     416           26 :                 yield Ok(res);
     417           26 : 
     418           26 :                 // We are done here
     419           26 :                 if next_marker.is_none() {
     420           26 :                     break;
     421           26 :                 }
     422           26 :             }
     423           26 :         }
     424           26 :     }
     425              : 
     426            3 :     async fn head_object(
     427            3 :         &self,
     428            3 :         key: &RemotePath,
     429            3 :         cancel: &CancellationToken,
     430            3 :     ) -> Result<ListingObject, DownloadError> {
     431            3 :         let kind = RequestKind::Head;
     432            3 :         let _permit = self.permit(kind, cancel).await?;
     433              : 
     434            3 :         let started_at = start_measuring_requests(kind);
     435            3 : 
     436            3 :         let blob_client = self.client.blob_client(self.relative_path_to_name(key));
     437            3 :         let properties_future = blob_client.get_properties().into_future();
     438            3 : 
     439            3 :         let properties_future = tokio::time::timeout(self.small_timeout, properties_future);
     440              : 
     441            3 :         let res = tokio::select! {
     442            3 :             res = properties_future => res,
     443            3 :             _ = cancel.cancelled() => return Err(TimeoutOrCancel::Cancel.into()),
     444              :         };
     445              : 
     446            3 :         if let Ok(inner) = &res {
     447            3 :             // do not incl. timeouts as errors in metrics but cancellations
     448            3 :             let started_at = ScopeGuard::into_inner(started_at);
     449            3 :             crate::metrics::BUCKET_METRICS
     450            3 :                 .req_seconds
     451            3 :                 .observe_elapsed(kind, inner, started_at);
     452            3 :         }
     453              : 
     454            3 :         let data = match res {
     455            2 :             Ok(Ok(data)) => Ok(data),
     456            1 :             Ok(Err(sdk)) => Err(to_download_error(sdk)),
     457            0 :             Err(_timeout) => Err(DownloadError::Timeout),
     458            1 :         }?;
     459              : 
     460            2 :         let properties = data.blob.properties;
     461            2 :         Ok(ListingObject {
     462            2 :             key: key.to_owned(),
     463            2 :             last_modified: SystemTime::from(properties.last_modified),
     464            2 :             size: properties.content_length,
     465            2 :         })
     466            0 :     }
     467              : 
     468           93 :     async fn upload(
     469           93 :         &self,
     470           93 :         from: impl Stream<Item = std::io::Result<Bytes>> + Send + Sync + 'static,
     471           93 :         data_size_bytes: usize,
     472           93 :         to: &RemotePath,
     473           93 :         metadata: Option<StorageMetadata>,
     474           93 :         cancel: &CancellationToken,
     475           93 :     ) -> anyhow::Result<()> {
     476           93 :         let kind = RequestKind::Put;
     477           93 :         let _permit = self.permit(kind, cancel).await?;
     478              : 
     479           93 :         let started_at = start_measuring_requests(kind);
     480           93 : 
     481           93 :         let op = async {
     482           93 :             let blob_client = self.client.blob_client(self.relative_path_to_name(to));
     483           93 : 
     484           93 :             let from: Pin<Box<dyn Stream<Item = std::io::Result<Bytes>> + Send + Sync + 'static>> =
     485           93 :                 Box::pin(from);
     486           93 : 
     487           93 :             let from = NonSeekableStream::new(from, data_size_bytes);
     488           93 : 
     489           93 :             let body = azure_core::Body::SeekableStream(Box::new(from));
     490           93 : 
     491           93 :             let mut builder = blob_client.put_block_blob(body);
     492              : 
     493           93 :             if let Some(metadata) = metadata {
     494            0 :                 builder = builder.metadata(to_azure_metadata(metadata));
     495            0 :             }
     496              : 
     497           93 :             let fut = builder.into_future();
     498           93 :             let fut = tokio::time::timeout(self.timeout, fut);
     499           93 : 
     500           93 :             match fut.await {
     501           93 :                 Ok(Ok(_response)) => Ok(()),
     502            0 :                 Ok(Err(azure)) => Err(azure.into()),
     503            0 :                 Err(_timeout) => Err(TimeoutOrCancel::Timeout.into()),
     504              :             }
     505            0 :         };
     506              : 
     507           93 :         let res = tokio::select! {
     508           93 :             res = op => res,
     509           93 :             _ = cancel.cancelled() => return Err(TimeoutOrCancel::Cancel.into()),
     510              :         };
     511              : 
     512           93 :         let outcome = match res {
     513           93 :             Ok(_) => AttemptOutcome::Ok,
     514            0 :             Err(_) => AttemptOutcome::Err,
     515              :         };
     516           93 :         let started_at = ScopeGuard::into_inner(started_at);
     517           93 :         crate::metrics::BUCKET_METRICS
     518           93 :             .req_seconds
     519           93 :             .observe_elapsed(kind, outcome, started_at);
     520           93 : 
     521           93 :         res
     522            0 :     }
     523              : 
     524           11 :     async fn download(
     525           11 :         &self,
     526           11 :         from: &RemotePath,
     527           11 :         opts: &DownloadOpts,
     528           11 :         cancel: &CancellationToken,
     529           11 :     ) -> Result<Download, DownloadError> {
     530           11 :         let blob_client = self.client.blob_client(self.relative_path_to_name(from));
     531           11 : 
     532           11 :         let mut builder = blob_client.get();
     533              : 
     534           11 :         if let Some(ref etag) = opts.etag {
     535            3 :             builder = builder.if_match(IfMatchCondition::NotMatch(etag.to_string()))
     536            0 :         }
     537              : 
     538           11 :         if let Some((start, end)) = opts.byte_range() {
     539            5 :             builder = builder.range(match end {
     540            3 :                 Some(end) => Range::Range(start..end),
     541            2 :                 None => Range::RangeFrom(start..),
     542              :             });
     543            0 :         }
     544              : 
     545           11 :         let timeout = match opts.kind {
     546            0 :             DownloadKind::Small => self.small_timeout,
     547           11 :             DownloadKind::Large => self.timeout,
     548              :         };
     549              : 
     550           11 :         self.download_for_builder(builder, timeout, cancel).await
     551            0 :     }
     552              : 
     553           86 :     async fn delete(&self, path: &RemotePath, cancel: &CancellationToken) -> anyhow::Result<()> {
     554           86 :         self.delete_objects(std::array::from_ref(path), cancel)
     555           86 :             .await
     556            0 :     }
     557              : 
     558           93 :     async fn delete_objects(
     559           93 :         &self,
     560           93 :         paths: &[RemotePath],
     561           93 :         cancel: &CancellationToken,
     562           93 :     ) -> anyhow::Result<()> {
     563           93 :         let kind = RequestKind::Delete;
     564           93 :         let _permit = self.permit(kind, cancel).await?;
     565           93 :         let started_at = start_measuring_requests(kind);
     566           93 : 
     567           93 :         let op = async {
     568              :             // TODO batch requests are not supported by the SDK
     569              :             // https://github.com/Azure/azure-sdk-for-rust/issues/1068
     570          205 :             for path in paths {
     571          112 :                 #[derive(Debug)]
     572          112 :                 enum AzureOrTimeout {
     573              :                     AzureError(azure_core::Error),
     574              :                     Timeout,
     575              :                     Cancel,
     576          112 :                 }
     577          112 :                 impl Display for AzureOrTimeout {
     578            0 :                     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
     579            0 :                         write!(f, "{self:?}")
     580            0 :                     }
     581              :                 }
     582          112 :                 let warn_threshold = 3;
     583          112 :                 let max_retries = 5;
     584          112 :                 backoff::retry(
     585          112 :                     || async {
     586          112 :                         let blob_client = self.client.blob_client(self.relative_path_to_name(path));
     587          112 : 
     588          112 :                         let request = blob_client.delete().into_future();
     589              : 
     590          112 :                         let res = tokio::time::timeout(self.timeout, request).await;
     591              : 
     592          112 :                         match res {
     593           90 :                             Ok(Ok(_v)) => Ok(()),
     594           22 :                             Ok(Err(azure_err)) => {
     595           22 :                                 if let Some(http_err) = azure_err.as_http_error() {
     596           22 :                                     if http_err.status() == StatusCode::NotFound {
     597           22 :                                         return Ok(());
     598            0 :                                     }
     599            0 :                                 }
     600            0 :                                 Err(AzureOrTimeout::AzureError(azure_err))
     601              :                             }
     602            0 :                             Err(_elapsed) => Err(AzureOrTimeout::Timeout),
     603              :                         }
     604          112 :                     },
     605          112 :                     |err| match err {
     606            0 :                         AzureOrTimeout::AzureError(_) | AzureOrTimeout::Timeout => false,
     607            0 :                         AzureOrTimeout::Cancel => true,
     608          112 :                     },
     609          112 :                     warn_threshold,
     610          112 :                     max_retries,
     611          112 :                     "deleting remote object",
     612          112 :                     cancel,
     613          112 :                 )
     614          112 :                 .await
     615          112 :                 .ok_or_else(|| AzureOrTimeout::Cancel)
     616          112 :                 .and_then(|x| x)
     617          112 :                 .map_err(|e| match e {
     618            0 :                     AzureOrTimeout::AzureError(err) => anyhow::Error::from(err),
     619            0 :                     AzureOrTimeout::Timeout => TimeoutOrCancel::Timeout.into(),
     620            0 :                     AzureOrTimeout::Cancel => TimeoutOrCancel::Cancel.into(),
     621          112 :                 })?;
     622              :             }
     623           93 :             Ok(())
     624            0 :         };
     625              : 
     626           93 :         let res = tokio::select! {
     627           93 :             res = op => res,
     628           93 :             _ = cancel.cancelled() => return Err(TimeoutOrCancel::Cancel.into()),
     629              :         };
     630              : 
     631           93 :         let started_at = ScopeGuard::into_inner(started_at);
     632           93 :         crate::metrics::BUCKET_METRICS
     633           93 :             .req_seconds
     634           93 :             .observe_elapsed(kind, &res, started_at);
     635           93 :         res
     636            0 :     }
     637              : 
     638            0 :     fn max_keys_per_delete(&self) -> usize {
     639            0 :         super::MAX_KEYS_PER_DELETE_AZURE
     640            0 :     }
     641              : 
     642            1 :     async fn copy(
     643            1 :         &self,
     644            1 :         from: &RemotePath,
     645            1 :         to: &RemotePath,
     646            1 :         cancel: &CancellationToken,
     647            1 :     ) -> anyhow::Result<()> {
     648            1 :         let kind = RequestKind::Copy;
     649            1 :         let _permit = self.permit(kind, cancel).await?;
     650            1 :         let started_at = start_measuring_requests(kind);
     651            1 : 
     652            1 :         let timeout = tokio::time::sleep(self.timeout);
     653            1 : 
     654            1 :         let mut copy_status = None;
     655            1 : 
     656            1 :         let op = async {
     657            1 :             let blob_client = self.client.blob_client(self.relative_path_to_name(to));
     658              : 
     659            1 :             let source_url = format!(
     660            1 :                 "{}/{}",
     661            1 :                 self.client.url()?,
     662            1 :                 self.relative_path_to_name(from)
     663              :             );
     664              : 
     665            1 :             let builder = blob_client.copy(Url::from_str(&source_url)?);
     666            1 :             let copy = builder.into_future();
     667              : 
     668            1 :             let result = copy.await?;
     669              : 
     670            1 :             copy_status = Some(result.copy_status);
     671              :             loop {
     672            1 :                 match copy_status.as_ref().expect("we always set it to Some") {
     673              :                     CopyStatus::Aborted => {
     674            0 :                         anyhow::bail!("Received abort for copy from {from} to {to}.");
     675              :                     }
     676              :                     CopyStatus::Failed => {
     677            0 :                         anyhow::bail!("Received failure response for copy from {from} to {to}.");
     678              :                     }
     679            1 :                     CopyStatus::Success => return Ok(()),
     680            0 :                     CopyStatus::Pending => (),
     681            0 :                 }
     682            0 :                 // The copy is taking longer. Waiting a second and then re-trying.
     683            0 :                 // TODO estimate time based on copy_progress and adjust time based on that
     684            0 :                 tokio::time::sleep(Duration::from_millis(1000)).await;
     685            0 :                 let properties = blob_client.get_properties().into_future().await?;
     686            0 :                 let Some(status) = properties.blob.properties.copy_status else {
     687            0 :                     tracing::warn!("copy_status for copy is None!, from={from}, to={to}");
     688            0 :                     return Ok(());
     689              :                 };
     690            0 :                 copy_status = Some(status);
     691              :             }
     692            0 :         };
     693              : 
     694            1 :         let res = tokio::select! {
     695            1 :             res = op => res,
     696            1 :             _ = cancel.cancelled() => return Err(anyhow::Error::new(TimeoutOrCancel::Cancel)),
     697            1 :             _ = timeout => {
     698            0 :                 let e = anyhow::Error::new(TimeoutOrCancel::Timeout);
     699            0 :                 let e = e.context(format!("Timeout, last status: {copy_status:?}"));
     700            0 :                 Err(e)
     701              :             },
     702              :         };
     703              : 
     704            1 :         let started_at = ScopeGuard::into_inner(started_at);
     705            1 :         crate::metrics::BUCKET_METRICS
     706            1 :             .req_seconds
     707            1 :             .observe_elapsed(kind, &res, started_at);
     708            1 :         res
     709            0 :     }
     710              : 
     711            0 :     async fn time_travel_recover(
     712            0 :         &self,
     713            0 :         _prefix: Option<&RemotePath>,
     714            0 :         _timestamp: SystemTime,
     715            0 :         _done_if_after: SystemTime,
     716            0 :         _cancel: &CancellationToken,
     717            0 :     ) -> Result<(), TimeTravelError> {
     718            0 :         // TODO use Azure point in time recovery feature for this
     719            0 :         // https://learn.microsoft.com/en-us/azure/storage/blobs/point-in-time-restore-overview
     720            0 :         Err(TimeTravelError::Unimplemented)
     721            0 :     }
     722              : }
     723              : 
     724              : pin_project_lite::pin_project! {
     725              :     /// Hack to work around not being able to stream once with azure sdk.
     726              :     ///
     727              :     /// Azure sdk clones streams around with the assumption that they are like
     728              :     /// `Arc<tokio::fs::File>` (except not supporting tokio), however our streams are not like
     729              :     /// that. For example for an `index_part.json` we just have a single chunk of [`Bytes`]
     730              :     /// representing the whole serialized vec. It could be trivially cloneable and "semi-trivially"
     731              :     /// seekable, but we can also just re-try the request easier.
     732              :     #[project = NonSeekableStreamProj]
     733              :     enum NonSeekableStream<S> {
     734              :         /// A stream wrappers initial form.
     735              :         ///
     736              :         /// Mutex exists to allow moving when cloning. If the sdk changes to do less than 1
     737              :         /// clone before first request, then this must be changed.
     738              :         Initial {
     739              :             inner: std::sync::Mutex<Option<tokio_util::compat::Compat<tokio_util::io::StreamReader<S, Bytes>>>>,
     740              :             len: usize,
     741              :         },
     742              :         /// The actually readable variant, produced by cloning the Initial variant.
     743              :         ///
     744              :         /// The sdk currently always clones once, even without retry policy.
     745              :         Actual {
     746              :             #[pin]
     747              :             inner: tokio_util::compat::Compat<tokio_util::io::StreamReader<S, Bytes>>,
     748              :             len: usize,
     749              :             read_any: bool,
     750              :         },
     751              :         /// Most likely unneeded, but left to make life easier, in case more clones are added.
     752              :         Cloned {
     753              :             len_was: usize,
     754              :         }
     755              :     }
     756              : }
     757              : 
     758              : impl<S> NonSeekableStream<S>
     759              : where
     760              :     S: Stream<Item = std::io::Result<Bytes>> + Send + Sync + 'static,
     761              : {
     762           93 :     fn new(inner: S, len: usize) -> NonSeekableStream<S> {
     763              :         use tokio_util::compat::TokioAsyncReadCompatExt;
     764              : 
     765           93 :         let inner = tokio_util::io::StreamReader::new(inner).compat();
     766           93 :         let inner = Some(inner);
     767           93 :         let inner = std::sync::Mutex::new(inner);
     768           93 :         NonSeekableStream::Initial { inner, len }
     769           93 :     }
     770              : }
     771              : 
     772              : impl<S> std::fmt::Debug for NonSeekableStream<S> {
     773            0 :     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
     774            0 :         match self {
     775            0 :             Self::Initial { len, .. } => f.debug_struct("Initial").field("len", len).finish(),
     776            0 :             Self::Actual { len, .. } => f.debug_struct("Actual").field("len", len).finish(),
     777            0 :             Self::Cloned { len_was, .. } => f.debug_struct("Cloned").field("len", len_was).finish(),
     778              :         }
     779            0 :     }
     780              : }
     781              : 
     782              : impl<S> futures::io::AsyncRead for NonSeekableStream<S>
     783              : where
     784              :     S: Stream<Item = std::io::Result<Bytes>>,
     785              : {
     786           93 :     fn poll_read(
     787           93 :         self: std::pin::Pin<&mut Self>,
     788           93 :         cx: &mut std::task::Context<'_>,
     789           93 :         buf: &mut [u8],
     790           93 :     ) -> std::task::Poll<std::io::Result<usize>> {
     791           93 :         match self.project() {
     792              :             NonSeekableStreamProj::Actual {
     793           93 :                 inner, read_any, ..
     794           93 :             } => {
     795           93 :                 *read_any = true;
     796           93 :                 inner.poll_read(cx, buf)
     797              :             }
     798              :             // NonSeekableStream::Initial does not support reading because it is just much easier
     799              :             // to have the mutex in place where one does not poll the contents, or that's how it
     800              :             // seemed originally. If there is a version upgrade which changes the cloning, then
     801              :             // that support needs to be hacked in.
     802              :             //
     803              :             // including {self:?} into the message would be useful, but unsure how to unproject.
     804            0 :             _ => std::task::Poll::Ready(Err(std::io::Error::new(
     805            0 :                 std::io::ErrorKind::Other,
     806            0 :                 "cloned or initial values cannot be read",
     807            0 :             ))),
     808              :         }
     809            0 :     }
     810              : }
     811              : 
     812              : impl<S> Clone for NonSeekableStream<S> {
     813              :     /// Weird clone implementation exists to support the sdk doing cloning before issuing the first
     814              :     /// request, see type documentation.
     815           93 :     fn clone(&self) -> Self {
     816              :         use NonSeekableStream::*;
     817              : 
     818           93 :         match self {
     819           93 :             Initial { inner, len } => {
     820           93 :                 if let Some(inner) = inner.lock().unwrap().take() {
     821           93 :                     Actual {
     822           93 :                         inner,
     823           93 :                         len: *len,
     824           93 :                         read_any: false,
     825           93 :                     }
     826              :                 } else {
     827            0 :                     Self::Cloned { len_was: *len }
     828              :                 }
     829              :             }
     830            0 :             Actual { len, .. } => Cloned { len_was: *len },
     831            0 :             Cloned { len_was } => Cloned { len_was: *len_was },
     832              :         }
     833            0 :     }
     834              : }
     835              : 
     836              : #[async_trait::async_trait]
     837              : impl<S> azure_core::SeekableStream for NonSeekableStream<S>
     838              : where
     839              :     S: Stream<Item = std::io::Result<Bytes>> + Unpin + Send + Sync + 'static,
     840              : {
     841            0 :     async fn reset(&mut self) -> azure_core::error::Result<()> {
     842              :         use NonSeekableStream::*;
     843              : 
     844            0 :         let msg = match self {
     845            0 :             Initial { inner, .. } => {
     846            0 :                 if inner.get_mut().unwrap().is_some() {
     847            0 :                     return Ok(());
     848              :                 } else {
     849            0 :                     "reset after first clone is not supported"
     850              :                 }
     851              :             }
     852            0 :             Actual { read_any, .. } if !*read_any => return Ok(()),
     853            0 :             Actual { .. } => "reset after reading is not supported",
     854            0 :             Cloned { .. } => "reset after second clone is not supported",
     855              :         };
     856            0 :         Err(azure_core::error::Error::new(
     857            0 :             azure_core::error::ErrorKind::Io,
     858            0 :             std::io::Error::new(std::io::ErrorKind::Other, msg),
     859            0 :         ))
     860            0 :     }
     861              : 
     862              :     // Note: it is not documented if this should be the total or remaining length, total passes the
     863              :     // tests.
     864           93 :     fn len(&self) -> usize {
     865              :         use NonSeekableStream::*;
     866           93 :         match self {
     867           93 :             Initial { len, .. } => *len,
     868            0 :             Actual { len, .. } => *len,
     869            0 :             Cloned { len_was, .. } => *len_was,
     870              :         }
     871            0 :     }
     872              : }
        

Generated by: LCOV version 2.1-beta