/build/source/nativelink-store/src/ontap_s3_store.rs
Line | Count | Source |
1 | | // Copyright 2025 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::fs::File; |
20 | | use std::io::BufReader; |
21 | | use std::sync::Arc; |
22 | | |
23 | | use async_trait::async_trait; |
24 | | use aws_config::BehaviorVersion; |
25 | | use aws_config::default_provider::credentials::DefaultCredentialsChain; |
26 | | use aws_config::provider_config::ProviderConfig; |
27 | | use aws_sdk_s3::Client; |
28 | | use aws_sdk_s3::config::Region; |
29 | | use aws_sdk_s3::operation::create_multipart_upload::CreateMultipartUploadOutput; |
30 | | use aws_sdk_s3::operation::get_object::GetObjectError; |
31 | | use aws_sdk_s3::operation::head_object::HeadObjectError; |
32 | | use aws_sdk_s3::primitives::{ByteStream, SdkBody}; |
33 | | use aws_sdk_s3::types::builders::{CompletedMultipartUploadBuilder, CompletedPartBuilder}; |
34 | | use base64::Engine; |
35 | | use base64::prelude::BASE64_STANDARD_NO_PAD; |
36 | | use bytes::BytesMut; |
37 | | use futures::future::{Either, FusedFuture}; |
38 | | use futures::stream::{FuturesUnordered, unfold}; |
39 | | use futures::{FutureExt, StreamExt, TryFutureExt, TryStreamExt}; |
40 | | use hyper_rustls::ConfigBuilderExt; |
41 | | use nativelink_config::stores::ExperimentalOntapS3Spec; |
42 | | use nativelink_error::{Code, Error, ResultExt, make_err}; |
43 | | use nativelink_metric::MetricsComponent; |
44 | | use nativelink_util::buf_channel::{ |
45 | | DropCloserReadHalf, DropCloserWriteHalf, make_buf_channel_pair, |
46 | | }; |
47 | | use nativelink_util::health_utils::{HealthStatus, HealthStatusIndicator}; |
48 | | use nativelink_util::instant_wrapper::InstantWrapper; |
49 | | use nativelink_util::retry::{Retrier, RetryResult}; |
50 | | use nativelink_util::store_trait::{RemoveCallback, StoreDriver, StoreKey, UploadSizeInfo}; |
51 | | use parking_lot::Mutex; |
52 | | use rustls::{ClientConfig, RootCertStore}; |
53 | | use rustls_pki_types::CertificateDer; |
54 | | use rustls_pki_types::pem::PemObject; |
55 | | use sha2::{Digest, Sha256}; |
56 | | use tokio::time::sleep; |
57 | | use tracing::{Level, event, warn}; |
58 | | |
59 | | use crate::cas_utils::is_zero_digest; |
60 | | use crate::common_s3_utils::{TlsClient, install_default_rustls_crypto_provider}; |
61 | | |
62 | | // S3 parts cannot be smaller than this number |
63 | | const MIN_MULTIPART_SIZE: u64 = 5 * 1024 * 1024; // 5MB |
64 | | |
65 | | // S3 parts cannot be larger than this number |
66 | | const MAX_MULTIPART_SIZE: u64 = 5 * 1024 * 1024 * 1024; // 5GB |
67 | | |
68 | | // S3 parts cannot be more than this number |
69 | | const MAX_UPLOAD_PARTS: usize = 10_000; |
70 | | |
71 | | // Default max buffer size for retrying upload requests |
72 | | const DEFAULT_MAX_RETRY_BUFFER_PER_REQUEST: usize = 20 * 1024 * 1024; // 20MB |
73 | | |
74 | | // Default limit for concurrent part uploads per multipart upload |
75 | | const DEFAULT_MULTIPART_MAX_CONCURRENT_UPLOADS: usize = 10; |
76 | | |
77 | | #[derive(Debug, MetricsComponent)] |
78 | | pub struct OntapS3Store<NowFn> { |
79 | | s3_client: Arc<Client>, |
80 | | now_fn: NowFn, |
81 | | #[metric(help = "The bucket name for the ONTAP S3 store")] |
82 | | bucket: String, |
83 | | #[metric(help = "The key prefix for the ONTAP S3 store")] |
84 | | key_prefix: String, |
85 | | retrier: Retrier, |
86 | | #[metric(help = "The number of seconds to consider an object expired")] |
87 | | consider_expired_after_s: u64, |
88 | | #[metric(help = "The number of bytes to buffer for retrying requests")] |
89 | | max_retry_buffer_per_request: usize, |
90 | | #[metric(help = "The number of concurrent uploads allowed for multipart uploads")] |
91 | | multipart_max_concurrent_uploads: usize, |
92 | | |
93 | | remove_callbacks: Mutex<Vec<RemoveCallback>>, |
94 | | } |
95 | | |
96 | 0 | pub fn load_custom_certs(cert_path: &str) -> Result<Arc<ClientConfig>, Error> { |
97 | 0 | let mut root_store = RootCertStore::empty(); |
98 | | |
99 | | // Create a BufReader from the cert file |
100 | 0 | let mut cert_reader = BufReader::new( |
101 | 0 | File::open(cert_path) |
102 | 0 | .err_tip(|| format!("Failed to open CA certificate file {cert_path}"))?, |
103 | | ); |
104 | | |
105 | | // Parse certificates |
106 | 0 | let certs = CertificateDer::pem_reader_iter(&mut cert_reader).collect::<Result<Vec<_>, _>>()?; |
107 | | |
108 | | // Add each certificate to the root store |
109 | 0 | for cert in certs { |
110 | 0 | root_store.add(cert).map_err(|e| { |
111 | 0 | Error::from_std_err(Code::Internal, &e) |
112 | 0 | .append("Failed to add certificate to root store") |
113 | 0 | })?; |
114 | | } |
115 | | |
116 | | // Build the client config with the root store |
117 | 0 | let config = ClientConfig::builder() |
118 | 0 | .with_root_certificates(root_store) |
119 | 0 | .with_no_client_auth(); |
120 | | |
121 | 0 | Ok(Arc::new(config)) |
122 | 0 | } |
123 | | |
124 | | impl<I, NowFn> OntapS3Store<NowFn> |
125 | | where |
126 | | I: InstantWrapper, |
127 | | NowFn: Fn() -> I + Send + Sync + Unpin + 'static, |
128 | | { |
129 | 3 | pub async fn new(spec: &ExperimentalOntapS3Spec, now_fn: NowFn) -> Result<Arc<Self>, Error> { |
130 | 3 | install_default_rustls_crypto_provider(); |
131 | | |
132 | | // Load custom CA config |
133 | 3 | let ca_config = if let Some(cert_path0 ) = &spec.root_certificates { |
134 | 0 | load_custom_certs(cert_path)? |
135 | | } else { |
136 | 3 | Arc::new( |
137 | 3 | ClientConfig::builder() |
138 | 3 | .with_native_roots()?0 |
139 | 3 | .with_no_client_auth(), |
140 | | ) |
141 | | }; |
142 | | |
143 | 3 | let https_connector = hyper_rustls::HttpsConnectorBuilder::new() |
144 | 3 | .with_tls_config((*ca_config).clone()) |
145 | 3 | .https_only() |
146 | 3 | .enable_http1() |
147 | 3 | .enable_http2() |
148 | 3 | .build(); |
149 | | |
150 | 3 | let http_client = TlsClient::with_https_connector(&spec.common, https_connector); |
151 | | |
152 | 3 | let credentials_provider = DefaultCredentialsChain::builder() |
153 | 3 | .configure( |
154 | 3 | ProviderConfig::without_region() |
155 | 3 | .with_region(Some(Region::new(Cow::Owned(spec.vserver_name.clone())))) |
156 | 3 | .with_http_client(http_client.clone()), |
157 | 3 | ) |
158 | 3 | .build() |
159 | 3 | .await; |
160 | | |
161 | 3 | let config = aws_sdk_s3::Config::builder() |
162 | 3 | .credentials_provider(credentials_provider) |
163 | 3 | .endpoint_url(&spec.endpoint) |
164 | 3 | .region(Region::new(spec.vserver_name.clone())) |
165 | 3 | .app_name(aws_config::AppName::new("nativelink").expect("valid app name")) |
166 | 3 | .http_client(http_client) |
167 | 3 | .force_path_style(true) |
168 | 3 | .behavior_version(BehaviorVersion::latest()) |
169 | 3 | .timeout_config( |
170 | 3 | aws_config::timeout::TimeoutConfig::builder() |
171 | 3 | .connect_timeout(Duration::from_secs(30)) |
172 | 3 | .operation_timeout(Duration::from_mins(2)) |
173 | 3 | .build(), |
174 | | ) |
175 | 3 | .build(); |
176 | | |
177 | 3 | let s3_client = Client::from_conf(config); |
178 | | |
179 | 3 | Self::new_with_client_and_jitter( |
180 | 3 | spec, |
181 | 3 | s3_client, |
182 | 3 | spec.common.retry.make_jitter_fn(), |
183 | 3 | now_fn, |
184 | | ) |
185 | 3 | } |
186 | | |
187 | 15 | pub fn new_with_client_and_jitter( |
188 | 15 | spec: &ExperimentalOntapS3Spec, |
189 | 15 | s3_client: Client, |
190 | 15 | jitter_fn: Arc<dyn (Fn(Duration) -> Duration) + Send + Sync>, |
191 | 15 | now_fn: NowFn, |
192 | 15 | ) -> Result<Arc<Self>, Error> { |
193 | 15 | Ok(Arc::new(Self { |
194 | 15 | s3_client: Arc::new(s3_client), |
195 | 15 | now_fn, |
196 | 15 | bucket: spec.bucket.clone(), |
197 | 15 | key_prefix: spec |
198 | 15 | .common |
199 | 15 | .key_prefix |
200 | 15 | .as_ref() |
201 | 15 | .unwrap_or(&String::new()) |
202 | 15 | .clone(), |
203 | 15 | retrier: Retrier::new( |
204 | 15 | Arc::new(|duration| Box::pin1 (sleep1 (duration1 ))), |
205 | 15 | jitter_fn, |
206 | 15 | spec.common.retry.clone(), |
207 | | ), |
208 | 15 | consider_expired_after_s: u64::from(spec.common.consider_expired_after_s), |
209 | 15 | max_retry_buffer_per_request: spec |
210 | 15 | .common |
211 | 15 | .max_retry_buffer_per_request |
212 | 15 | .unwrap_or(DEFAULT_MAX_RETRY_BUFFER_PER_REQUEST), |
213 | 15 | multipart_max_concurrent_uploads: spec |
214 | 15 | .common |
215 | 15 | .multipart_max_concurrent_uploads |
216 | 15 | .unwrap_or(DEFAULT_MULTIPART_MAX_CONCURRENT_UPLOADS), |
217 | 15 | remove_callbacks: Mutex::new(vec![]), |
218 | | })) |
219 | 15 | } |
220 | | |
221 | 15 | fn make_s3_path(&self, key: &StoreKey<'_>) -> String { |
222 | 15 | format!("{}{}", self.key_prefix, key.as_str()) |
223 | 15 | } |
224 | | |
225 | 9 | async fn has(self: Pin<&Self>, digest: StoreKey<'_>) -> Result<Option<u64>, Error> { |
226 | 9 | let digest_clone = digest.into_owned(); |
227 | 9 | self.retrier |
228 | 10 | .retry9 (unfold9 (()9 , move |state| { |
229 | 10 | let local_digest = digest_clone.clone(); |
230 | 10 | async move { |
231 | 10 | let result = self |
232 | 10 | .s3_client |
233 | 10 | .head_object() |
234 | 10 | .bucket(&self.bucket) |
235 | 10 | .key(self.make_s3_path(&local_digest)) |
236 | 10 | .send() |
237 | 10 | .await; |
238 | | |
239 | 10 | match result { |
240 | 8 | Ok(head_object_output) => { |
241 | 8 | if self.consider_expired_after_s != 0 |
242 | 2 | && let Some(last_modified) = head_object_output.last_modified |
243 | | { |
244 | 2 | let now_s = (self.now_fn)().unix_timestamp(); |
245 | 2 | if TryInto::<u64>::try_into(last_modified.secs()) |
246 | 2 | .unwrap_or(u64::MAX) |
247 | 2 | + self.consider_expired_after_s |
248 | 2 | <= now_s |
249 | | { |
250 | 1 | let remove_callbacks = self.remove_callbacks.lock().clone(); |
251 | 1 | let mut callbacks: FuturesUnordered<_> = remove_callbacks |
252 | 1 | .into_iter() |
253 | 1 | .map(|callback| {0 |
254 | 0 | let store_key = local_digest.borrow(); |
255 | 0 | async move { callback.callback(store_key).await } |
256 | 0 | }) |
257 | 1 | .collect(); |
258 | 1 | while callbacks.next().await.is_some() {}0 |
259 | 1 | return Some((RetryResult::Ok(None), state)); |
260 | 1 | } |
261 | 6 | } |
262 | 7 | let Some(length) = head_object_output.content_length else { |
263 | 0 | return Some((RetryResult::Ok(None), state)); |
264 | | }; |
265 | 7 | if length >= 0 { |
266 | 7 | return Some((RetryResult::Ok(Some(length as u64)), state)); |
267 | 0 | } |
268 | 0 | Some(( |
269 | 0 | RetryResult::Err(make_err!( |
270 | 0 | Code::InvalidArgument, |
271 | 0 | "Negative content length in ONTAP S3: {length:?}" |
272 | 0 | )), |
273 | 0 | state, |
274 | 0 | )) |
275 | | } |
276 | 2 | Err(sdk_error) => match sdk_error.into_service_error() { |
277 | 1 | HeadObjectError::NotFound(_) => Some((RetryResult::Ok(None), state)), |
278 | 1 | other => Some(( |
279 | 1 | RetryResult::Retry(make_err!( |
280 | 1 | Code::Unavailable, |
281 | 1 | "Unhandled HeadObjectError in ONTAP S3: {other:?}" |
282 | 1 | )), |
283 | 1 | state, |
284 | 1 | )), |
285 | | }, |
286 | | } |
287 | 10 | } |
288 | 10 | })) |
289 | 9 | .await |
290 | 9 | } |
291 | | } |
292 | | |
293 | | #[async_trait] |
294 | | impl<I, NowFn> StoreDriver for OntapS3Store<NowFn> |
295 | | where |
296 | | I: InstantWrapper, |
297 | | NowFn: Fn() -> I + Send + Sync + Unpin + 'static, |
298 | | { |
299 | 0 | async fn post_init(self: Arc<Self>) -> Result<(), Error> { |
300 | | Ok(()) |
301 | 0 | } |
302 | | |
303 | | async fn has_with_results( |
304 | | self: Pin<&Self>, |
305 | | keys: &[StoreKey<'_>], |
306 | | results: &mut [Option<u64>], |
307 | 10 | ) -> Result<(), Error> { |
308 | | keys.iter() |
309 | | .zip(results.iter_mut()) |
310 | 10 | .map(|(key, result)| async move { |
311 | 10 | if is_zero_digest(key.borrow()) { |
312 | 1 | *result = Some(0); |
313 | 1 | return Ok::<_, Error>(()); |
314 | 9 | } |
315 | | |
316 | 9 | match self.has(key.borrow()).await { |
317 | 9 | Ok(size) => { |
318 | 9 | *result = size; |
319 | 9 | if size.is_none() { |
320 | 2 | event!( |
321 | 2 | Level::INFO, |
322 | 2 | key = %key.as_str(), |
323 | | "Object not found in ONTAP S3" |
324 | | ); |
325 | 7 | } |
326 | 9 | Ok(()) |
327 | | } |
328 | 0 | Err(err) => { |
329 | 0 | event!( |
330 | 0 | Level::ERROR, |
331 | 0 | key = %key.as_str(), |
332 | | error = ?err, |
333 | | "Error checking object existence" |
334 | | ); |
335 | 0 | Err(err) |
336 | | } |
337 | | } |
338 | 20 | }) |
339 | | .collect::<FuturesUnordered<_>>() |
340 | | .try_collect() |
341 | | .await |
342 | 10 | } |
343 | | |
344 | | async fn update( |
345 | | self: Pin<&Self>, |
346 | | key: StoreKey<'_>, |
347 | | mut reader: DropCloserReadHalf, |
348 | | size_info: UploadSizeInfo, |
349 | 2 | ) -> Result<u64, Error> { |
350 | | let s3_path = &self.make_s3_path(&key); |
351 | | |
352 | | let max_size = match size_info { |
353 | | UploadSizeInfo::ExactSize(sz) | UploadSizeInfo::MaxSize(sz) => sz, |
354 | | }; |
355 | | |
356 | | // For small files, use simple upload. For larger files or unknown size, use multipart |
357 | | if max_size < MIN_MULTIPART_SIZE && matches!(size_info, UploadSizeInfo::ExactSize(_)) { |
358 | | let UploadSizeInfo::ExactSize(sz) = size_info else { |
359 | | unreachable!("upload_size must be UploadSizeInfo::ExactSize here"); |
360 | | }; |
361 | | reader.set_max_recent_data_size( |
362 | | u64::try_from(self.max_retry_buffer_per_request) |
363 | | .err_tip(|| "Could not convert max_retry_buffer_per_request to u64")?, |
364 | | ); |
365 | | return self.retrier.retry( |
366 | 1 | unfold(reader, move |mut reader| async move { |
367 | 1 | let (mut tx, mut rx) = make_buf_channel_pair(); |
368 | | |
369 | 1 | let result = { |
370 | 1 | let reader_ref = &mut reader; |
371 | 1 | let (upload_res, bind_res): (Result<u64, Error>, Result<(), Error>) = tokio::join!(async move { |
372 | 1 | let raw_body_bytes = { |
373 | 1 | let mut raw_body_chunks = BytesMut::new(); |
374 | | loop { |
375 | 51 | match rx.recv().await { |
376 | 51 | Ok(chunk) => { |
377 | 51 | if chunk.is_empty() { |
378 | 1 | break Ok(raw_body_chunks.freeze()); |
379 | 50 | } |
380 | 50 | raw_body_chunks.extend_from_slice(&chunk); |
381 | | } |
382 | 0 | Err(err) => { |
383 | 0 | break Err(err); |
384 | | } |
385 | | } |
386 | | } |
387 | | }; |
388 | 1 | let internal_res = match raw_body_bytes { |
389 | 1 | Ok(body_bytes) => { |
390 | 1 | let hash = Sha256::digest(&body_bytes); |
391 | 1 | let body_len: u64 = body_bytes.len().try_into().unwrap_or(0); |
392 | 1 | let send_res = self.s3_client |
393 | 1 | .put_object() |
394 | 1 | .bucket(&self.bucket) |
395 | 1 | .key(s3_path.clone()) |
396 | 1 | .content_length(sz.try_into().unwrap_or(i64::MAX)) |
397 | 1 | .body( |
398 | 1 | ByteStream::from(body_bytes) |
399 | | ) |
400 | 1 | .set_checksum_algorithm(Some(aws_sdk_s3::types::ChecksumAlgorithm::Sha256)) |
401 | 1 | .set_checksum_sha256(Some(BASE64_STANDARD_NO_PAD.encode(hash))) |
402 | 1 | .customize() |
403 | 1 | .mutate_request(|req| {req.headers_mut().insert("x-amz-content-sha256", "UNSIGNED-PAYLOAD");}) |
404 | 1 | .send(); |
405 | 1 | Either::Left(send_res.map_ok_or_else(|e| Err( |
406 | 1 | Error::from_std_err0 (Code::Aborted0 , &e0 )), move |_| Ok(body_len))) |
407 | | } |
408 | 0 | Err(collect_err) => { |
409 | 0 | async fn make_collect_err(collect_err: Error) -> Result<u64, Error> { |
410 | 0 | Err(collect_err) |
411 | 0 | } |
412 | | |
413 | 0 | warn!( |
414 | | ?collect_err, |
415 | | "Failed to get body"); |
416 | 0 | let future_err = make_collect_err(collect_err); |
417 | 0 | Either::Right(future_err) |
418 | | } |
419 | | }; |
420 | 1 | internal_res.await |
421 | 1 | }, |
422 | 1 | tx.bind_buffered(reader_ref) |
423 | | ); |
424 | 1 | match (upload_res, bind_res) { |
425 | 1 | (Ok(size), Ok(())) => Ok(size), |
426 | 0 | (Err(e), _) | (_, Err(e)) => Err(e), |
427 | | } |
428 | 1 | .err_tip(|| "Failed to upload file to ONTAP S3 in single chunk") |
429 | | }; |
430 | | |
431 | 1 | let retry_result = result.map_or_else( |
432 | 0 | |mut err| { |
433 | 0 | err.code = Code::Aborted; |
434 | 0 | let bytes_received = reader.get_bytes_received(); |
435 | 0 | if let Err(try_reset_err) = reader.try_reset_stream() { |
436 | 0 | event!( |
437 | 0 | Level::ERROR, |
438 | | ?bytes_received, |
439 | | err = ?try_reset_err, |
440 | | "Unable to reset stream after failed upload in OntapS3Store::update" |
441 | | ); |
442 | 0 | return RetryResult::Err( |
443 | 0 | err |
444 | 0 | .merge(try_reset_err) |
445 | 0 | .append( |
446 | 0 | format!( |
447 | 0 | "Failed to retry upload with {bytes_received} bytes received in OntapS3Store::update" |
448 | 0 | ) |
449 | 0 | ) |
450 | 0 | ); |
451 | 0 | } |
452 | 0 | let err = err.append( |
453 | 0 | format!( |
454 | | "Retry on upload happened with {bytes_received} bytes received in OntapS3Store::update" |
455 | | ) |
456 | | ); |
457 | 0 | event!(Level::INFO, ?err, ?bytes_received, "Retryable ONTAP S3 error"); |
458 | 0 | RetryResult::Retry(err) |
459 | 0 | }, |
460 | | RetryResult::Ok |
461 | | ); |
462 | 1 | Some((retry_result, reader)) |
463 | 2 | }) |
464 | | ).await; |
465 | | } |
466 | | |
467 | | // Handle multipart upload for large files |
468 | | let upload_id = &self |
469 | | .retrier |
470 | 1 | .retry(unfold((), move |()| async move { |
471 | 1 | let retry_result = self |
472 | 1 | .s3_client |
473 | 1 | .create_multipart_upload() |
474 | 1 | .bucket(&self.bucket) |
475 | 1 | .key(s3_path) |
476 | 1 | .send() |
477 | 1 | .await |
478 | 1 | .map_or_else( |
479 | 0 | |e| { |
480 | 0 | RetryResult::Retry( |
481 | 0 | Error::from_std_err(Code::Aborted, &e) |
482 | 0 | .append("Failed to create multipart upload to ONTAP S3"), |
483 | 0 | ) |
484 | 0 | }, |
485 | 1 | |CreateMultipartUploadOutput { upload_id, .. }| { |
486 | 1 | upload_id.map_or_else( |
487 | 0 | || { |
488 | 0 | RetryResult::Err(make_err!( |
489 | 0 | Code::Internal, |
490 | 0 | "Expected upload_id to be set by ONTAP S3 response" |
491 | 0 | )) |
492 | 0 | }, |
493 | | RetryResult::Ok, |
494 | | ) |
495 | 1 | }, |
496 | | ); |
497 | 1 | Some((retry_result, ())) |
498 | 2 | })) |
499 | | .await?; |
500 | | |
501 | | let bytes_per_upload_part = |
502 | | (max_size / (MIN_MULTIPART_SIZE - 1)).clamp(MIN_MULTIPART_SIZE, MAX_MULTIPART_SIZE); |
503 | | |
504 | 1 | let upload_parts = move || async move { |
505 | 1 | let (tx, mut rx) = tokio::sync::mpsc::channel(self.multipart_max_concurrent_uploads); |
506 | | |
507 | 1 | let read_stream_fut = ( |
508 | 1 | async move { |
509 | 1 | let retrier = &Pin::get_ref(self).retrier; |
510 | 1 | let mut total_uploaded = 0; |
511 | 4 | for part_number in 1..i32::MAX1 { |
512 | 4 | let write_buf = reader |
513 | 4 | .consume( |
514 | | Some( |
515 | 4 | usize |
516 | 4 | ::try_from(bytes_per_upload_part) |
517 | 4 | .err_tip( |
518 | | || "Could not convert bytes_per_upload_part to usize" |
519 | 0 | )? |
520 | | ) |
521 | 4 | ).await |
522 | 4 | .err_tip(|| "Failed to read chunk in ontap_s3_store")?0 ; |
523 | 4 | if write_buf.is_empty() { |
524 | 1 | break; |
525 | 3 | } |
526 | | |
527 | 3 | total_uploaded += write_buf.len() as u64; |
528 | | |
529 | 3 | tx |
530 | 3 | .send( |
531 | 3 | retrier.retry( |
532 | 3 | unfold(write_buf, move |write_buf| async move { |
533 | 3 | let retry_result = self.s3_client |
534 | 3 | .upload_part() |
535 | 3 | .bucket(&self.bucket) |
536 | 3 | .key(s3_path) |
537 | 3 | .upload_id(upload_id) |
538 | 3 | .body(ByteStream::new(SdkBody::from(write_buf.clone()))) |
539 | 3 | .part_number(part_number) |
540 | 3 | .send().await |
541 | 3 | .map_or_else( |
542 | 0 | |e| { |
543 | 0 | RetryResult::Retry( |
544 | 0 | Error::from_std_err( |
545 | 0 | Code::Aborted,&e).append( |
546 | 0 | "Failed to upload part {part_number} in ONTAP S3 store" |
547 | 0 | ) |
548 | 0 | ) |
549 | 0 | }, |
550 | 3 | |mut response| { |
551 | 3 | RetryResult::Ok( |
552 | 3 | CompletedPartBuilder::default() |
553 | 3 | .set_e_tag(response.e_tag.take()) |
554 | 3 | .part_number(part_number) |
555 | 3 | .build() |
556 | 3 | ) |
557 | 3 | } |
558 | | ); |
559 | 3 | Some((retry_result, write_buf)) |
560 | 6 | }) |
561 | | ) |
562 | 3 | ).await |
563 | 3 | .map_err(|err| {0 |
564 | 0 | Error::from_std_err(Code::Internal, &err).append("Failed to send part to channel in ontap_s3_store") |
565 | 0 | })?; |
566 | | } |
567 | 1 | Result::<_, Error>::Ok(total_uploaded) |
568 | | } |
569 | 1 | ).fuse(); |
570 | | |
571 | 1 | let mut upload_futures = FuturesUnordered::new(); |
572 | 1 | let mut total_uploaded = 0; |
573 | 1 | let mut completed_parts = Vec::with_capacity( |
574 | 1 | usize::try_from(cmp::min( |
575 | 1 | MAX_UPLOAD_PARTS as u64, |
576 | 1 | max_size / bytes_per_upload_part + 1, |
577 | | )) |
578 | 1 | .err_tip(|| "Could not convert u64 to usize")?0 , |
579 | | ); |
580 | | |
581 | 1 | tokio::pin!(read_stream_fut); |
582 | | loop { |
583 | 8 | if read_stream_fut.is_terminated() && rx7 .is_empty7 () && upload_futures2 .is_empty2 () { |
584 | 1 | break; |
585 | 7 | } |
586 | 7 | tokio::select! { |
587 | 7 | result1 = &mut read_stream_fut => { |
588 | 1 | total_uploaded = result?0 ; |
589 | | }, |
590 | 7 | Some(upload_result3 ) = upload_futures.next() => completed_parts3 .push3 (upload_result3 ?0 ), |
591 | 7 | Some(fut3 ) = rx.recv() => upload_futures3 .push3 (fut3 ), |
592 | | } |
593 | | } |
594 | | |
595 | 1 | completed_parts.sort_unstable_by_key(|part| part.part_number); |
596 | | |
597 | 1 | self.retrier |
598 | 1 | .retry(unfold(completed_parts, move |completed_parts| async move { |
599 | | Some(( |
600 | 1 | self.s3_client |
601 | 1 | .complete_multipart_upload() |
602 | 1 | .bucket(&self.bucket) |
603 | 1 | .key(s3_path) |
604 | 1 | .multipart_upload( |
605 | 1 | CompletedMultipartUploadBuilder::default() |
606 | 1 | .set_parts(Some(completed_parts.clone())) |
607 | 1 | .build(), |
608 | 1 | ) |
609 | 1 | .upload_id(upload_id) |
610 | 1 | .send() |
611 | 1 | .await |
612 | 1 | .map_or_else( |
613 | 0 | |e| { |
614 | 0 | RetryResult::Retry( |
615 | 0 | Error::from_std_err(Code::Aborted, &e).append( |
616 | 0 | "Failed to complete multipart upload in ONTAP S3 store", |
617 | 0 | ), |
618 | 0 | ) |
619 | 0 | }, |
620 | 1 | |_| RetryResult::Ok(total_uploaded), |
621 | | ), |
622 | 1 | completed_parts, |
623 | | )) |
624 | 2 | })) |
625 | 1 | .await |
626 | 2 | }; |
627 | | |
628 | | upload_parts() |
629 | 0 | .or_else(move |mut e| async move { |
630 | 0 | let abort_res = self |
631 | 0 | .s3_client |
632 | 0 | .abort_multipart_upload() |
633 | 0 | .bucket(&self.bucket) |
634 | 0 | .key(s3_path) |
635 | 0 | .upload_id(upload_id) |
636 | 0 | .send() |
637 | 0 | .await; |
638 | 0 | if let Err(abort_err) = abort_res { |
639 | 0 | let err = Error::from_std_err(Code::Aborted, &abort_err) |
640 | 0 | .append("Failed to abort multipart upload in ONTAP S3 store"); |
641 | 0 | event!(Level::INFO, ?err, "Multipart upload error"); |
642 | 0 | e = e.merge(err); |
643 | 0 | } |
644 | 0 | Err(e) |
645 | 0 | }) |
646 | | .await |
647 | 2 | } |
648 | | |
649 | | async fn get_part( |
650 | | self: Pin<&Self>, |
651 | | key: StoreKey<'_>, |
652 | | writer: &mut DropCloserWriteHalf, |
653 | | offset: u64, |
654 | | length: Option<u64>, |
655 | 4 | ) -> Result<(), Error> { |
656 | | if is_zero_digest(key.borrow()) { |
657 | | writer |
658 | | .send_eof() |
659 | | .err_tip(|| "Failed to send zero EOF in ONTAP S3 store get_part")?; |
660 | | return Ok(()); |
661 | | } |
662 | | |
663 | | let s3_path = &self.make_s3_path(&key); |
664 | | let end_read_byte = length |
665 | 2 | .map_or(Some(None), |length| Some(offset.checked_add(length))) |
666 | | .err_tip(|| "Integer overflow protection triggered")?; |
667 | | |
668 | | self.retrier |
669 | 3 | .retry(unfold(writer, move |writer| async move { |
670 | 3 | let result = self |
671 | 3 | .s3_client |
672 | 3 | .get_object() |
673 | 3 | .bucket(&self.bucket) |
674 | 3 | .key(s3_path) |
675 | 3 | .range(format!( |
676 | | "bytes={}-{}", |
677 | 3 | offset + writer.get_bytes_written(), |
678 | 3 | end_read_byte.map_or_else(String::new, |v| v2 .to_string2 ()) |
679 | | )) |
680 | 3 | .send() |
681 | 3 | .await; |
682 | | |
683 | 3 | match result { |
684 | 3 | Ok(head_object_output) => { |
685 | 3 | let mut s3_in_stream = head_object_output.body; |
686 | 3 | let _bytes_sent = 0; |
687 | | |
688 | 7 | while let Some(maybe_bytes4 ) = s3_in_stream.next().await { |
689 | 4 | match maybe_bytes { |
690 | 4 | Ok(bytes) => { |
691 | 4 | if bytes.is_empty() { |
692 | 1 | continue; |
693 | 3 | } |
694 | | |
695 | | // Clone bytes before sending |
696 | 3 | let bytes_clone = bytes.clone(); |
697 | | |
698 | | // More robust sending mechanism |
699 | 3 | match writer.send(bytes).await { |
700 | 3 | Ok(()) => { |
701 | 3 | let _ = bytes_clone.len(); |
702 | 3 | } |
703 | 0 | Err(e) => { |
704 | 0 | return Some(( |
705 | 0 | RetryResult::Err(Error::from_std_err( |
706 | 0 | Code::Aborted,&e).append( |
707 | 0 | "Error sending bytes to consumer in ONTAP S3" |
708 | 0 | )), |
709 | 0 | writer, |
710 | 0 | )); |
711 | | } |
712 | | } |
713 | | } |
714 | 0 | Err(e) => { |
715 | 0 | return Some(( |
716 | 0 | RetryResult::Retry( |
717 | 0 | Error::from_std_err(Code::Aborted, &e) |
718 | 0 | .append("Bad bytestream element in ONTAP S3"), |
719 | 0 | ), |
720 | 0 | writer, |
721 | 0 | )); |
722 | | } |
723 | | } |
724 | | } |
725 | | |
726 | | // EOF handling |
727 | 3 | if let Err(e0 ) = writer.send_eof() { |
728 | 0 | return Some(( |
729 | 0 | RetryResult::Err( |
730 | 0 | Error::from_std_err(Code::Aborted, &e) |
731 | 0 | .append("Failed to send EOF to consumer in ONTAP S3"), |
732 | 0 | ), |
733 | 0 | writer, |
734 | 0 | )); |
735 | 3 | } |
736 | | |
737 | 3 | Some((RetryResult::Ok(()), writer)) |
738 | | } |
739 | 0 | Err(sdk_error) => match sdk_error.into_service_error() { |
740 | 0 | GetObjectError::NoSuchKey(e) => Some(( |
741 | 0 | RetryResult::Err( |
742 | 0 | Error::from_std_err(Code::NotFound, &e) |
743 | 0 | .append("No such key in ONTAP S3"), |
744 | 0 | ), |
745 | 0 | writer, |
746 | 0 | )), |
747 | 0 | other => Some(( |
748 | 0 | RetryResult::Retry( |
749 | 0 | Error::from_std_err(Code::Unavailable, &other) |
750 | 0 | .append("Unhandled GetObjectError in ONTAP S3"), |
751 | 0 | ), |
752 | 0 | writer, |
753 | 0 | )), |
754 | | }, |
755 | | } |
756 | 6 | })) |
757 | | .await |
758 | 4 | } |
759 | | |
760 | 0 | fn inner_store(&self, _digest: Option<StoreKey>) -> &dyn StoreDriver { |
761 | 0 | self |
762 | 0 | } |
763 | | |
764 | 0 | fn as_any<'a>(&'a self) -> &'a (dyn core::any::Any + Sync + Send + 'static) { |
765 | 0 | self |
766 | 0 | } |
767 | | |
768 | 0 | fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send + 'static> { |
769 | 0 | self |
770 | 0 | } |
771 | | |
772 | 3 | fn register_remove_callback(self: Arc<Self>, callback: RemoveCallback) -> Result<(), Error> { |
773 | 3 | self.remove_callbacks.lock().push(callback); |
774 | 3 | Ok(()) |
775 | 3 | } |
776 | | } |
777 | | |
778 | | #[async_trait] |
779 | | impl<I, NowFn> HealthStatusIndicator for OntapS3Store<NowFn> |
780 | | where |
781 | | I: InstantWrapper, |
782 | | NowFn: Fn() -> I + Send + Sync + Unpin + 'static, |
783 | | { |
784 | 0 | fn get_name(&self) -> &'static str { |
785 | 0 | "OntapS3Store" |
786 | 0 | } |
787 | | |
788 | 0 | async fn check_health(&self, namespace: Cow<'static, str>) -> HealthStatus { |
789 | | StoreDriver::check_health(Pin::new(self), namespace).await |
790 | 0 | } |
791 | | } |