/build/source/nativelink-store/src/azure_blob_store.rs
Line | Count | Source |
1 | | // Copyright 2024 The NativeLink Authors. All rights reserved. |
2 | | // |
3 | | // Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); |
4 | | // you may not use this file except in compliance with the License. |
5 | | // You may obtain a copy of the License at |
6 | | // |
7 | | // See LICENSE file for details |
8 | | // |
9 | | // Unless required by applicable law or agreed to in writing, software |
10 | | // distributed under the License is distributed on an "AS IS" BASIS, |
11 | | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | | // See the License for the specific language governing permissions and |
13 | | // limitations under the License. |
14 | | |
15 | | use core::cmp; |
16 | | use core::pin::Pin; |
17 | | use core::time::Duration; |
18 | | use std::borrow::Cow; |
19 | | use std::sync::Arc; |
20 | | |
21 | | use async_trait::async_trait; |
22 | | use azure_core::credentials::TokenCredential; |
23 | | use azure_core::error::ErrorKind; |
24 | | use azure_core::http::{RequestContent, RetryOptions, StatusCode, Transport, Url}; |
25 | | use azure_identity::WorkloadIdentityCredential; |
26 | | use azure_storage_blob::clients::{BlobContainerClient, BlobContainerClientOptions}; |
27 | | use azure_storage_blob::models::{ |
28 | | BlobClientDownloadOptions, BlobClientGetPropertiesResultHeaders, BlockLookupList, HttpRange, |
29 | | StorageErrorCode, |
30 | | }; |
31 | | use futures::future::FusedFuture; |
32 | | use futures::stream::{FuturesUnordered, unfold}; |
33 | | use futures::{FutureExt, StreamExt, TryStreamExt}; |
34 | | use nativelink_config::stores::ExperimentalAzureSpec; |
35 | | use nativelink_error::{Code, Error, ResultExt, make_err}; |
36 | | use nativelink_metric::MetricsComponent; |
37 | | use nativelink_util::buf_channel::{ |
38 | | DropCloserReadHalf, DropCloserWriteHalf, make_buf_channel_pair, |
39 | | }; |
40 | | use nativelink_util::fs; |
41 | | use nativelink_util::health_utils::{HealthRegistryBuilder, HealthStatus, HealthStatusIndicator}; |
42 | | use nativelink_util::instant_wrapper::InstantWrapper; |
43 | | use nativelink_util::retry::{Retrier, RetryResult}; |
44 | | use nativelink_util::store_trait::{ |
45 | | RemoveCallback, StoreDriver, StoreKey, StoreOptimizations, UploadSizeInfo, |
46 | | }; |
47 | | use tokio::sync::mpsc; |
48 | | use tokio::time::sleep; |
49 | | use tracing::{Level, event}; |
50 | | |
51 | | use crate::cas_utils::is_zero_digest; |
52 | | use crate::common_s3_utils::install_default_rustls_crypto_provider; |
53 | | |
54 | | // Check the below doc for the limits specific to Azure. |
55 | | // https://learn.microsoft.com/en-us/azure/storage/blobs/scalability-targets#scale-targets-for-blob-storage |
56 | | |
57 | | // Maximum number of blocks in a block blob or append blob |
58 | | const MAX_BLOCKS: usize = 50_000; |
59 | | |
60 | | // Maximum size of a block in a block blob (4,000 MiB) |
61 | | const MAX_BLOCK_SIZE: u64 = 4_000 * 1024 * 1024; // 4,000 MiB = 4 GiB |
62 | | |
63 | | // Default block size for uploads (5 MiB) |
64 | | const DEFAULT_BLOCK_SIZE: u64 = 5 * 1024 * 1024; // 5 MiB |
65 | | |
66 | | // Default maximum retry buffer per request |
67 | | const DEFAULT_MAX_RETRY_BUFFER_PER_REQUEST: usize = 5 * 1024 * 1024; // 5 MiB |
68 | | |
69 | | // Default maximum number of concurrent uploads |
70 | | const DEFAULT_MAX_CONCURRENT_UPLOADS: usize = 10; |
71 | | |
72 | | // Default public Azure Blob Storage endpoint suffix. |
73 | | const DEFAULT_BLOB_ENDPOINT_SUFFIX: &str = "blob.core.windows.net"; |
74 | | |
75 | | #[derive(MetricsComponent)] |
76 | | pub struct AzureBlobStore<NowFn> { |
77 | | client: Arc<BlobContainerClient>, |
78 | | now_fn: NowFn, |
79 | | #[metric(help = "The container name for the Azure store")] |
80 | | container: String, |
81 | | #[metric(help = "The blob prefix for the Azure store")] |
82 | | blob_prefix: String, |
83 | | retrier: Retrier, |
84 | | #[metric(help = "The number of seconds to consider an object expired")] |
85 | | consider_expired_after_s: i64, |
86 | | #[metric(help = "The number of bytes to buffer for retrying requests")] |
87 | | max_retry_buffer_per_request: usize, |
88 | | #[metric(help = "The number of concurrent uploads allowed")] |
89 | | max_concurrent_uploads: usize, |
90 | | } |
91 | | |
92 | | impl<NowFn> core::fmt::Debug for AzureBlobStore<NowFn> { |
93 | 0 | fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { |
94 | 0 | f.debug_struct("AzureBlobStore") |
95 | 0 | .field("container", &self.container) |
96 | 0 | .field("blob_prefix", &self.blob_prefix) |
97 | 0 | .field("consider_expired_after_s", &self.consider_expired_after_s) |
98 | 0 | .finish_non_exhaustive() |
99 | 0 | } |
100 | | } |
101 | | |
102 | | impl<I, NowFn> AzureBlobStore<NowFn> |
103 | | where |
104 | | I: InstantWrapper, |
105 | | NowFn: Fn() -> I + Send + Sync + Unpin + 'static, |
106 | | { |
107 | 0 | pub fn new(spec: &ExperimentalAzureSpec, now_fn: NowFn) -> Result<Arc<Self>, Error> { |
108 | 0 | let jitter_fn = spec.common.retry.make_jitter_fn(); |
109 | 0 | let client = Self::build_container_client(spec)?; |
110 | 0 | Self::new_with_client_and_jitter(spec, client, jitter_fn, now_fn) |
111 | 0 | } |
112 | | |
113 | | /// Builds the container URL and selects the auth strategy: |
114 | | /// * `sas_url` set -> use it verbatim as the container URL with no credential. |
115 | | /// * otherwise -> `https://{account}.{endpoint}/{container}` authenticated with |
116 | | /// Entra ID via Workload Identity (keyless). |
117 | 0 | fn build_container_client(spec: &ExperimentalAzureSpec) -> Result<BlobContainerClient, Error> { |
118 | 0 | let mut options = BlobContainerClientOptions::default(); |
119 | 0 | options.client_options.retry = RetryOptions::none(); |
120 | | // Hand the SDK an HTTP client with an explicit rustls (ring) config. |
121 | 0 | options.client_options.transport = Some(Self::build_http_transport()?); |
122 | | |
123 | 0 | let (container_url, credential): (Url, Option<Arc<dyn TokenCredential>>) = |
124 | 0 | if let Some(sas_url) = spec.sas_url.as_ref() { |
125 | 0 | let url = Url::parse(sas_url) |
126 | 0 | .map_err(|e| make_err!(Code::InvalidArgument, "Invalid Azure sas_url: {e}"))?; |
127 | 0 | (url, None) |
128 | | } else { |
129 | 0 | let endpoint = spec.endpoint.clone().unwrap_or_else(|| { |
130 | 0 | format!( |
131 | | "https://{}.{DEFAULT_BLOB_ENDPOINT_SUFFIX}", |
132 | | spec.account_name |
133 | | ) |
134 | 0 | }); |
135 | 0 | let mut url = Url::parse(&endpoint) |
136 | 0 | .map_err(|e| make_err!(Code::InvalidArgument, "Invalid Azure endpoint: {e}"))?; |
137 | 0 | url.path_segments_mut() |
138 | 0 | .map_err(|()| { |
139 | 0 | make_err!( |
140 | 0 | Code::InvalidArgument, |
141 | | "Azure endpoint is not a valid base URL: {endpoint}" |
142 | | ) |
143 | 0 | })? |
144 | 0 | .pop_if_empty() |
145 | 0 | .push(&spec.container); |
146 | 0 | let credential: Arc<dyn TokenCredential> = WorkloadIdentityCredential::new(None) |
147 | 0 | .map_err(|e| { |
148 | 0 | make_err!( |
149 | 0 | Code::FailedPrecondition, |
150 | | "Failed to create Azure Workload Identity credential: {e}" |
151 | | ) |
152 | 0 | })?; |
153 | 0 | (url, Some(credential)) |
154 | | }; |
155 | | |
156 | 0 | BlobContainerClient::new(container_url, credential, Some(options)) |
157 | 0 | .map_err(|e| make_err!(Code::Unavailable, "Failed to create Azure client: {e}")) |
158 | 0 | } |
159 | | |
160 | | /// Builds an HTTP transport for the Azure SDK backed by a reqwest client with |
161 | | /// an explicit rustls config using `NativeLink`'s ring crypto provider, so the |
162 | | /// SDK never falls back to guessing a provider (which breaks HTTPS here). |
163 | 0 | fn build_http_transport() -> Result<Transport, Error> { |
164 | 0 | install_default_rustls_crypto_provider(); |
165 | | |
166 | 0 | let mut roots = rustls::RootCertStore::empty(); |
167 | 0 | roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned()); |
168 | 0 | let tls_config = rustls::ClientConfig::builder() |
169 | 0 | .with_root_certificates(roots) |
170 | 0 | .with_no_client_auth(); |
171 | | |
172 | 0 | let client = reqwest::Client::builder() |
173 | 0 | .use_preconfigured_tls(tls_config) |
174 | 0 | .build() |
175 | 0 | .map_err(|e| make_err!(Code::Unavailable, "Failed to build Azure HTTP client: {e}"))?; |
176 | | |
177 | 0 | Ok(Transport::new(Arc::new(client))) |
178 | 0 | } |
179 | | |
180 | 14 | pub fn new_with_client_and_jitter( |
181 | 14 | spec: &ExperimentalAzureSpec, |
182 | 14 | client: BlobContainerClient, |
183 | 14 | jitter_fn: Arc<dyn Fn(Duration) -> Duration + Send + Sync>, |
184 | 14 | now_fn: NowFn, |
185 | 14 | ) -> Result<Arc<Self>, Error> { |
186 | 14 | Ok(Arc::new(Self { |
187 | 14 | client: Arc::new(client), |
188 | 14 | now_fn, |
189 | 14 | container: spec.container.clone(), |
190 | 14 | blob_prefix: spec |
191 | 14 | .common |
192 | 14 | .key_prefix |
193 | 14 | .as_ref() |
194 | 14 | .unwrap_or(&String::new()) |
195 | 14 | .clone(), |
196 | 14 | retrier: Retrier::new( |
197 | 14 | Arc::new(|duration| Box::pin5 (sleep5 (duration5 ))), |
198 | 14 | jitter_fn, |
199 | 14 | spec.common.retry.clone(), |
200 | | ), |
201 | 14 | consider_expired_after_s: i64::from(spec.common.consider_expired_after_s), |
202 | 14 | max_retry_buffer_per_request: spec |
203 | 14 | .common |
204 | 14 | .max_retry_buffer_per_request |
205 | 14 | .unwrap_or(DEFAULT_MAX_RETRY_BUFFER_PER_REQUEST), |
206 | 14 | max_concurrent_uploads: spec |
207 | 14 | .common |
208 | 14 | .multipart_max_concurrent_uploads |
209 | 14 | .unwrap_or(DEFAULT_MAX_CONCURRENT_UPLOADS), |
210 | | })) |
211 | 14 | } |
212 | | |
213 | 13 | fn make_blob_path(&self, key: &StoreKey<'_>) -> String { |
214 | 13 | format!("{}{}", self.blob_prefix, key.as_str()) |
215 | 13 | } |
216 | | |
217 | 5 | async fn has(self: Pin<&Self>, digest: &StoreKey<'_>) -> Result<Option<u64>, Error> { |
218 | 5 | let blob_path = self.make_blob_path(digest); |
219 | | |
220 | 5 | self.retrier |
221 | 6 | .retry5 (unfold5 (()5 , move |state| { |
222 | 6 | let blob_path = blob_path.clone(); |
223 | 6 | let client = Arc::clone(&self.client); |
224 | 6 | async move { |
225 | 6 | let _permit = match fs::get_permit().await { |
226 | 6 | Ok(permit) => permit, |
227 | 0 | Err(e) => { |
228 | 0 | return Some(( |
229 | 0 | RetryResult::Retry(make_err!( |
230 | 0 | Code::Unavailable, |
231 | 0 | "Failed to acquire permit: {e}" |
232 | 0 | )), |
233 | 0 | state, |
234 | 0 | )); |
235 | | } |
236 | | }; |
237 | | |
238 | 6 | let result = client.blob_client(&blob_path).get_properties(None).await; |
239 | | |
240 | 6 | match result { |
241 | 4 | Ok(props) => { |
242 | 4 | if self.consider_expired_after_s > 0 |
243 | 2 | && let Some(last_modified) = props.last_modified().ok().flatten() |
244 | | { |
245 | 2 | let now = (self.now_fn)() |
246 | 2 | .unix_timestamp() |
247 | 2 | .try_into() |
248 | 2 | .unwrap_or(i64::MAX); |
249 | 2 | if last_modified.unix_timestamp() + self.consider_expired_after_s |
250 | 2 | <= now |
251 | | { |
252 | 1 | return Some((RetryResult::Ok(None), state)); |
253 | 1 | } |
254 | 2 | } |
255 | 3 | let blob_size = props.content_length().ok().flatten().unwrap_or(0); |
256 | 3 | Some((RetryResult::Ok(Some(blob_size)), state)) |
257 | | } |
258 | 2 | Err(err) => { |
259 | 2 | if err.http_status() == Some(StatusCode::NotFound) { |
260 | | // Distinguish a missing container (a config error) from a |
261 | | // missing blob (a normal cache miss). |
262 | | if let ErrorKind::HttpResponse { |
263 | 0 | error_code: Some(error_code), |
264 | | .. |
265 | 1 | } = err.kind() |
266 | 0 | && error_code == StorageErrorCode::ContainerNotFound.as_ref() |
267 | | { |
268 | 0 | return Some(( |
269 | 0 | RetryResult::Err(make_err!( |
270 | 0 | Code::InvalidArgument, |
271 | 0 | "Container not found: {err}" |
272 | 0 | )), |
273 | 0 | state, |
274 | 0 | )); |
275 | 1 | } |
276 | 1 | Some((RetryResult::Ok(None), state)) |
277 | | } else { |
278 | 1 | Some(( |
279 | 1 | RetryResult::Retry(make_err!( |
280 | 1 | Code::Unavailable, |
281 | 1 | "Failed to get blob properties: {err:?}" |
282 | 1 | )), |
283 | 1 | state, |
284 | 1 | )) |
285 | | } |
286 | | } |
287 | | } |
288 | 6 | } |
289 | 6 | })) |
290 | 5 | .await |
291 | 5 | } |
292 | | } |
293 | | |
294 | | #[async_trait] |
295 | | impl<I, NowFn> StoreDriver for AzureBlobStore<NowFn> |
296 | | where |
297 | | I: InstantWrapper, |
298 | | NowFn: Fn() -> I + Send + Sync + Unpin + 'static, |
299 | | { |
300 | 0 | async fn post_init(self: Arc<Self>) -> Result<(), Error> { |
301 | | Ok(()) |
302 | 0 | } |
303 | | |
304 | | async fn has_with_results( |
305 | | self: Pin<&Self>, |
306 | | keys: &[StoreKey<'_>], |
307 | | results: &mut [Option<u64>], |
308 | 6 | ) -> Result<(), Error> { |
309 | | keys.iter() |
310 | | .zip(results.iter_mut()) |
311 | 6 | .map(|(key, result)| async move { |
312 | 6 | if is_zero_digest(key.borrow()) { |
313 | 1 | *result = Some(0); |
314 | 1 | return Ok::<_, Error>(()); |
315 | 5 | } |
316 | 5 | *result = self.has(key).await?0 ; |
317 | 5 | Ok::<_, Error>(()) |
318 | 12 | }) |
319 | | .collect::<FuturesUnordered<_>>() |
320 | | .try_collect() |
321 | | .await |
322 | 6 | } |
323 | | |
324 | 0 | fn optimized_for(&self, optimization: StoreOptimizations) -> bool { |
325 | 0 | matches!(optimization, StoreOptimizations::LazyExistenceOnSync) |
326 | 0 | } |
327 | | |
328 | | async fn update( |
329 | | self: Pin<&Self>, |
330 | | digest: StoreKey<'_>, |
331 | | mut reader: DropCloserReadHalf, |
332 | | upload_size: UploadSizeInfo, |
333 | 4 | ) -> Result<u64, Error> { |
334 | | let blob_path = self.make_blob_path(&digest); |
335 | | // Handling zero-sized content check |
336 | | if upload_size == UploadSizeInfo::ExactSize(0) { |
337 | | return Ok(0); |
338 | | } |
339 | | |
340 | | let max_size = match upload_size { |
341 | | UploadSizeInfo::ExactSize(sz) | UploadSizeInfo::MaxSize(sz) => sz, |
342 | | }; |
343 | | |
344 | | // For small files of a known size we buffer to `Bytes` and upload in a single request. |
345 | | if max_size < DEFAULT_BLOCK_SIZE && matches!(upload_size, UploadSizeInfo::ExactSize(_)) { |
346 | | let UploadSizeInfo::ExactSize(sz) = upload_size else { |
347 | | unreachable!("upload_size must be UploadSizeInfo::ExactSize here"); |
348 | | }; |
349 | | |
350 | | reader.set_max_recent_data_size( |
351 | | u64::try_from(self.max_retry_buffer_per_request) |
352 | | .err_tip(|| "Could not convert max_retry_buffer_per_request to u64")?, |
353 | | ); |
354 | | |
355 | | return self |
356 | | .retrier |
357 | 3 | .retry(unfold(reader, move |mut reader| { |
358 | 3 | let client = Arc::clone(&self.client); |
359 | 3 | let blob_path = blob_path.clone(); |
360 | 3 | async move { |
361 | 3 | let _permit = match fs::get_permit().await { |
362 | 3 | Ok(permit) => permit, |
363 | 0 | Err(e) => { |
364 | 0 | return Some(( |
365 | 0 | RetryResult::Retry(make_err!( |
366 | 0 | Code::Unavailable, |
367 | 0 | "Failed to acquire permit: {e}" |
368 | 0 | )), |
369 | 0 | reader, |
370 | 0 | )); |
371 | | } |
372 | | }; |
373 | | |
374 | 3 | let (mut tx, mut rx) = make_buf_channel_pair(); |
375 | | |
376 | 3 | let result = { |
377 | 3 | let reader_ref = &mut reader; |
378 | 3 | let (upload_res, bind_res) = tokio::join!( |
379 | 3 | async { |
380 | 3 | let mut buffer = Vec::with_capacity( |
381 | 3 | usize::try_from(sz).expect( |
382 | 3 | "size must be non-negative and fit in usize", |
383 | | ), |
384 | | ); |
385 | 9 | while let Ok(Some(chunk6 )) = rx.try_next().await { |
386 | 6 | buffer.extend_from_slice(&chunk); |
387 | 6 | } |
388 | | |
389 | 3 | client |
390 | 3 | .blob_client(&blob_path) |
391 | 3 | .block_blob_client() |
392 | 3 | .upload(RequestContent::from(buffer), None) |
393 | 3 | .await |
394 | 3 | .map(|_| ()) |
395 | 3 | .map_err(|e| make_err!1 (Code::Aborted1 , "{e:?}")) |
396 | 3 | }, |
397 | 3 | async { tx.bind_buffered(reader_ref).await } |
398 | | ); |
399 | | |
400 | 3 | match (upload_res, bind_res) { |
401 | 2 | (Ok(()), Ok(())) => Ok(()), |
402 | 1 | (Err(e), _) | (_, Err(e0 )) => Err(e), |
403 | | } |
404 | 3 | .err_tip(|| "Failed to upload blob in single chunk") |
405 | | }; |
406 | | |
407 | 3 | match result { |
408 | | Ok(()) => { |
409 | 2 | Some((RetryResult::Ok(reader.get_bytes_received()), reader)) |
410 | | } |
411 | 1 | Err(mut err) => { |
412 | 1 | err.code = Code::Aborted; |
413 | 1 | let bytes_received = reader.get_bytes_received(); |
414 | | |
415 | 1 | if let Err(try_reset_err0 ) = reader.try_reset_stream() { |
416 | 0 | event!( |
417 | 0 | Level::ERROR, |
418 | | ?bytes_received, |
419 | | err = ?try_reset_err, |
420 | | "Unable to reset stream after failed upload in AzureStore::update" |
421 | | ); |
422 | 0 | Some(( |
423 | 0 | RetryResult::Err(err.merge(try_reset_err).append(format!( |
424 | 0 | "Failed to retry upload with {bytes_received} bytes received in AzureStore::update" |
425 | 0 | ))), |
426 | 0 | reader, |
427 | 0 | )) |
428 | | } else { |
429 | 1 | let err = err.append(format!( |
430 | | "Retry on upload happened with {bytes_received} bytes received in AzureStore::update" |
431 | | )); |
432 | 1 | event!( |
433 | 1 | Level::INFO, |
434 | | ?err, |
435 | | ?bytes_received, |
436 | | "Retryable Azure error" |
437 | | ); |
438 | 1 | Some((RetryResult::Retry(err), reader)) |
439 | | } |
440 | | } |
441 | | } |
442 | 3 | } |
443 | 3 | })) |
444 | | .await; |
445 | | } |
446 | | |
447 | | // For larger files we stream the content as staged blocks and commit a block list. |
448 | | let block_size = |
449 | | cmp::min(max_size / (MAX_BLOCKS as u64 - 1), MAX_BLOCK_SIZE).max(DEFAULT_BLOCK_SIZE); |
450 | | |
451 | | let (tx, mut rx) = mpsc::channel(self.max_concurrent_uploads); |
452 | | let mut block_ids: Vec<Vec<u8>> = Vec::with_capacity(MAX_BLOCKS); |
453 | | let retrier = self.retrier.clone(); |
454 | | |
455 | | let read_stream_fut = { |
456 | | let tx = tx.clone(); |
457 | | let blob_path = blob_path.clone(); |
458 | 1 | async move { |
459 | 1 | let mut total_uploaded = 0; |
460 | 4 | for block_id in 0..MAX_BLOCKS1 { |
461 | 4 | let write_buf = reader |
462 | 4 | .consume(Some( |
463 | 4 | usize::try_from(block_size) |
464 | 4 | .err_tip(|| "Could not convert block_size to usize")?0 , |
465 | | )) |
466 | 4 | .await |
467 | 4 | .err_tip(|| "Failed to read chunk in azure_store")?0 ; |
468 | | |
469 | 4 | if write_buf.is_empty() { |
470 | 1 | break; |
471 | 3 | } |
472 | | |
473 | 3 | total_uploaded += write_buf.len() as u64; |
474 | | |
475 | | // Fixed-width, zero-padded ids keep the committed block list ordered |
476 | | // after a lexicographic sort. |
477 | 3 | let block_id = format!("{block_id:032}").into_bytes(); |
478 | 3 | let blob_path = blob_path.clone(); |
479 | | |
480 | 3 | tx.send(async move { |
481 | 3 | self.retrier |
482 | 3 | .retry(unfold( |
483 | 3 | (write_buf, block_id), |
484 | 3 | move |(write_buf, block_id)| { |
485 | 3 | let client = Arc::clone(&self.client); |
486 | 3 | let blob_path = blob_path.clone(); |
487 | 3 | async move { |
488 | 3 | let _permit = match fs::get_permit().await { |
489 | 3 | Ok(permit) => permit, |
490 | 0 | Err(e) => { |
491 | 0 | return Some(( |
492 | 0 | RetryResult::Retry(make_err!( |
493 | 0 | Code::Unavailable, |
494 | 0 | "Failed to acquire permit: {e}" |
495 | 0 | )), |
496 | 0 | (write_buf, block_id), |
497 | 0 | )); |
498 | | } |
499 | | }; |
500 | 3 | let content_length = write_buf.len() as u64; |
501 | 3 | let retry_result = client |
502 | 3 | .blob_client(&blob_path) |
503 | 3 | .block_blob_client() |
504 | 3 | .stage_block( |
505 | 3 | &block_id, |
506 | 3 | content_length, |
507 | 3 | RequestContent::from(write_buf.to_vec()), |
508 | 3 | None, |
509 | 3 | ) |
510 | 3 | .await |
511 | 3 | .map_or_else( |
512 | 0 | |e| { |
513 | 0 | RetryResult::Retry(make_err!( |
514 | 0 | Code::Aborted, |
515 | 0 | "Failed to upload block in Azure store: {e:?}" |
516 | 0 | )) |
517 | 0 | }, |
518 | 3 | |_| RetryResult::Ok(block_id.clone()), |
519 | | ); |
520 | 3 | Some((retry_result, (write_buf, block_id))) |
521 | 3 | } |
522 | 3 | }, |
523 | | )) |
524 | 3 | .await |
525 | 3 | }) |
526 | 3 | .await |
527 | 3 | .map_err(|err| {0 |
528 | 0 | Error::from_std_err(Code::Internal, &err) |
529 | 0 | .append("Failed to send block to channel") |
530 | 0 | })?; |
531 | | } |
532 | 1 | Ok::<_, Error>(total_uploaded) |
533 | 1 | } |
534 | | .fuse() |
535 | | }; |
536 | | |
537 | | let mut upload_futures = FuturesUnordered::new(); |
538 | | let mut total_uploaded = 0; |
539 | | |
540 | | tokio::pin!(read_stream_fut); |
541 | | |
542 | | loop { |
543 | | if read_stream_fut.is_terminated() && rx.is_empty() && upload_futures.is_empty() { |
544 | | break; |
545 | | } |
546 | | tokio::select! { |
547 | | result = &mut read_stream_fut => { |
548 | | total_uploaded = result?; |
549 | | }, |
550 | | Some(block_id) = upload_futures.next() => block_ids.push(block_id?), |
551 | | Some(fut) = rx.recv() => upload_futures.push(fut), |
552 | | } |
553 | | } |
554 | | |
555 | | // Sorting block IDs to ensure consistent ordering of the committed blob. |
556 | | block_ids.sort_unstable(); |
557 | | |
558 | | let block_list = BlockLookupList { |
559 | | latest: Some(block_ids), |
560 | | ..Default::default() |
561 | | }; |
562 | | |
563 | | retrier |
564 | 1 | .retry(unfold(block_list, move |block_list| { |
565 | 1 | let client = Arc::clone(&self.client); |
566 | 1 | let blob_path = blob_path.clone(); |
567 | | |
568 | 1 | async move { |
569 | 1 | let _permit = match fs::get_permit().await { |
570 | 1 | Ok(permit) => permit, |
571 | 0 | Err(e) => { |
572 | 0 | return Some(( |
573 | 0 | RetryResult::Retry(make_err!( |
574 | 0 | Code::Unavailable, |
575 | 0 | "Failed to acquire permit: {e}" |
576 | 0 | )), |
577 | 0 | block_list, |
578 | 0 | )); |
579 | | } |
580 | | }; |
581 | | |
582 | 1 | let blocks = match RequestContent::try_from(block_list.clone()) { |
583 | 1 | Ok(blocks) => blocks, |
584 | 0 | Err(e) => { |
585 | 0 | return Some(( |
586 | 0 | RetryResult::Err(make_err!( |
587 | 0 | Code::Internal, |
588 | 0 | "Failed to serialize block list in Azure store: {e:?}" |
589 | 0 | )), |
590 | 0 | block_list, |
591 | 0 | )); |
592 | | } |
593 | | }; |
594 | | |
595 | 1 | let retry_result = client |
596 | 1 | .blob_client(&blob_path) |
597 | 1 | .block_blob_client() |
598 | 1 | .commit_block_list(blocks, None) |
599 | 1 | .await |
600 | 1 | .map_or_else( |
601 | 0 | |e| { |
602 | 0 | RetryResult::Retry( |
603 | 0 | Error::from_std_err(Code::Aborted, &e) |
604 | 0 | .append("Failed to commit block list in Azure store:"), |
605 | 0 | ) |
606 | 0 | }, |
607 | 1 | |_| RetryResult::Ok(total_uploaded), |
608 | | ); |
609 | 1 | Some((retry_result, block_list)) |
610 | 1 | } |
611 | 1 | })) |
612 | | .await |
613 | 4 | } |
614 | | |
615 | | async fn get_part( |
616 | | self: Pin<&Self>, |
617 | | key: StoreKey<'_>, |
618 | | writer: &mut DropCloserWriteHalf, |
619 | | offset: u64, |
620 | | length: Option<u64>, |
621 | 5 | ) -> Result<(), Error> { |
622 | | if is_zero_digest(key.borrow()) { |
623 | | writer |
624 | | .send_eof() |
625 | | .err_tip(|| "Failed to send zero EOF in azure store get_part")?; |
626 | | return Ok(()); |
627 | | } |
628 | | |
629 | | let blob_path = self.make_blob_path(&key); |
630 | | |
631 | | let range = match length { |
632 | | Some(len) => Some(HttpRange::new(offset, len)), |
633 | | None if offset == 0 => None, |
634 | | None => Some(HttpRange::from_offset(offset)), |
635 | | }; |
636 | | |
637 | | self.retrier |
638 | 7 | .retry(unfold(writer, move |writer| { |
639 | 7 | let range = range.clone(); |
640 | 7 | let client = Arc::clone(&self.client); |
641 | 7 | let blob_path = blob_path.clone(); |
642 | 7 | async move { |
643 | 7 | let _permit = match fs::get_permit().await { |
644 | 7 | Ok(permit) => permit, |
645 | 0 | Err(e) => { |
646 | 0 | return Some(( |
647 | 0 | RetryResult::Retry(make_err!( |
648 | 0 | Code::Unavailable, |
649 | 0 | "Failed to acquire permit: {e}" |
650 | 0 | )), |
651 | 0 | writer, |
652 | 0 | )); |
653 | | } |
654 | | }; |
655 | | |
656 | 7 | let result: Result<(), Error> = async { |
657 | 7 | let options = BlobClientDownloadOptions { |
658 | 7 | range, |
659 | 7 | ..Default::default() |
660 | 7 | }; |
661 | 7 | let response3 = client |
662 | 7 | .blob_client(&blob_path) |
663 | 7 | .download(Some(options)) |
664 | 7 | .await |
665 | 7 | .map_err(|e| {4 |
666 | 4 | if e.http_status() == Some(StatusCode::NotFound) { |
667 | 1 | make_err!(Code::NotFound, "Blob not found in Azure: {e:?}") |
668 | | } else { |
669 | 3 | make_err!( |
670 | 3 | Code::Aborted, |
671 | | "Failed to start download from Azure: {e:?}" |
672 | | ) |
673 | | } |
674 | 4 | })?; |
675 | | |
676 | 3 | let mut body = response.body; |
677 | 6 | while let Some(chunk3 ) = body.try_next().await.map_err(|e| {0 |
678 | 0 | make_err!(Code::Aborted, "Error reading from Azure stream: {e:?}") |
679 | 0 | })? { |
680 | 3 | if chunk.is_empty() { |
681 | 0 | continue; |
682 | 3 | } |
683 | 3 | writer.send(chunk).await.map_err(|e| {0 |
684 | 0 | make_err!(Code::Aborted, "Failed to send data to writer: {e:?}") |
685 | 0 | })?; |
686 | | } |
687 | | |
688 | 3 | writer.send_eof().map_err(|e| {0 |
689 | 0 | make_err!(Code::Aborted, "Failed to send EOF to writer: {e:?}") |
690 | 0 | })?; |
691 | 3 | Ok(()) |
692 | 7 | } |
693 | 7 | .await; |
694 | | |
695 | 7 | match result { |
696 | 3 | Ok(()) => Some((RetryResult::Ok(()), writer)), |
697 | 4 | Err(e) => { |
698 | 4 | if e.code == Code::NotFound { |
699 | 1 | Some((RetryResult::Err(e), writer)) |
700 | | } else { |
701 | 3 | Some((RetryResult::Retry(e), writer)) |
702 | | } |
703 | | } |
704 | | } |
705 | 7 | } |
706 | 7 | })) |
707 | | .await |
708 | 5 | } |
709 | | |
710 | 0 | fn inner_store(&self, _digest: Option<StoreKey>) -> &'_ dyn StoreDriver { |
711 | 0 | self |
712 | 0 | } |
713 | | |
714 | 0 | fn as_any<'a>(&'a self) -> &'a (dyn core::any::Any + Sync + Send + 'static) { |
715 | 0 | self |
716 | 0 | } |
717 | | |
718 | 0 | fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send + 'static> { |
719 | 0 | self |
720 | 0 | } |
721 | | |
722 | 0 | fn register_health(self: Arc<Self>, registry: &mut HealthRegistryBuilder) { |
723 | 0 | registry.register_indicator(self); |
724 | 0 | } |
725 | | |
726 | 0 | fn register_remove_callback(self: Arc<Self>, _callback: RemoveCallback) -> Result<(), Error> { |
727 | | // Azure Blob Storage manages object lifecycle externally, |
728 | | // so we can safely ignore remove callbacks. |
729 | 0 | Ok(()) |
730 | 0 | } |
731 | | } |
732 | | |
733 | | #[async_trait] |
734 | | impl<I, NowFn> HealthStatusIndicator for AzureBlobStore<NowFn> |
735 | | where |
736 | | I: InstantWrapper, |
737 | | NowFn: Fn() -> I + Send + Sync + Unpin + 'static, |
738 | | { |
739 | 0 | fn get_name(&self) -> &'static str { |
740 | 0 | "AzureBlobStore" |
741 | 0 | } |
742 | | |
743 | 0 | async fn check_health(&self, namespace: Cow<'static, str>) -> HealthStatus { |
744 | | StoreDriver::check_health(Pin::new(self), namespace).await |
745 | 0 | } |
746 | | } |