/build/source/nativelink-store/src/s3_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 aws_config::default_provider::credentials; |
23 | | use aws_config::provider_config::ProviderConfig; |
24 | | use aws_config::{AppName, BehaviorVersion}; |
25 | | use aws_sdk_s3::Client; |
26 | | use aws_sdk_s3::config::Region; |
27 | | use aws_sdk_s3::operation::create_multipart_upload::CreateMultipartUploadOutput; |
28 | | use aws_sdk_s3::operation::get_object::GetObjectError; |
29 | | use aws_sdk_s3::operation::head_object::HeadObjectError; |
30 | | use aws_sdk_s3::primitives::ByteStream; // SdkBody |
31 | | use aws_sdk_s3::types::builders::{CompletedMultipartUploadBuilder, CompletedPartBuilder}; |
32 | | use aws_smithy_types::body::SdkBody; |
33 | | use futures::future::FusedFuture; |
34 | | use futures::stream::{FuturesUnordered, unfold}; |
35 | | use futures::{FutureExt, StreamExt, TryFutureExt, TryStreamExt}; |
36 | | use nativelink_config::stores::ExperimentalAwsSpec; |
37 | | // Note: S3 store should be very careful about the error codes it returns |
38 | | // when in a retryable wrapper. Always prefer Code::Aborted or another |
39 | | // retryable code over Code::InvalidArgument or make_input_err!(). |
40 | | // ie: Don't import make_input_err!() to help prevent this. |
41 | | use nativelink_error::{Code, Error, ResultExt, make_err}; |
42 | | use nativelink_metric::MetricsComponent; |
43 | | use nativelink_util::buf_channel::{ |
44 | | DropCloserReadHalf, DropCloserWriteHalf, make_buf_channel_pair, |
45 | | }; |
46 | | use nativelink_util::health_utils::{HealthRegistryBuilder, HealthStatus, HealthStatusIndicator}; |
47 | | use nativelink_util::instant_wrapper::InstantWrapper; |
48 | | use nativelink_util::retry::{Retrier, RetryResult}; |
49 | | use nativelink_util::store_trait::{ |
50 | | RemoveCallback, StoreDriver, StoreKey, StoreOptimizations, UploadSizeInfo, |
51 | | }; |
52 | | use parking_lot::Mutex; |
53 | | use tokio::sync::mpsc; |
54 | | use tokio::time::sleep; |
55 | | use tracing::{error, info}; |
56 | | |
57 | | use crate::cas_utils::is_zero_digest; |
58 | | use crate::common_s3_utils::{BodyWrapper, TlsClient}; |
59 | | |
60 | | // S3 object cannot be larger than this number. See: |
61 | | // https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html |
62 | | const MAX_UPLOAD_SIZE: u64 = 48 * 1024 * 1024 * 1024 * 1024; // 48TiB (technically should be 48.8 TiB, but close enough) |
63 | | |
64 | | // S3 parts cannot be smaller than this number. See: |
65 | | // https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html |
66 | | const MIN_MULTIPART_SIZE: u64 = 5 * 1024 * 1024; // 5MB. |
67 | | |
68 | | // S3 parts cannot be larger than this number. See: |
69 | | // https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html |
70 | | const MAX_MULTIPART_SIZE: u64 = 5 * 1024 * 1024 * 1024; // 5GB. |
71 | | |
72 | | // S3 parts cannot be more than this number. See: |
73 | | // https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html |
74 | | // Note: Type 'u64' chosen to simplify calculations |
75 | | const MAX_UPLOAD_PARTS: u64 = 10_000; |
76 | | |
77 | | // Default max buffer size for retrying upload requests. |
78 | | // Note: If you change this, adjust the docs in the config. |
79 | | const DEFAULT_MAX_RETRY_BUFFER_PER_REQUEST: usize = 5 * 1024 * 1024; // 5MB. |
80 | | |
81 | | // Default limit for concurrent part uploads per multipart upload. |
82 | | // Note: If you change this, adjust the docs in the config. |
83 | | const DEFAULT_MULTIPART_MAX_CONCURRENT_UPLOADS: usize = 10; |
84 | | |
85 | | #[derive(Debug, MetricsComponent)] |
86 | | pub struct S3Store<NowFn> { |
87 | | s3_client: Arc<Client>, |
88 | | now_fn: NowFn, |
89 | | #[metric(help = "The bucket name for the S3 store")] |
90 | | bucket: String, |
91 | | #[metric(help = "The key prefix for the S3 store")] |
92 | | key_prefix: String, |
93 | | retrier: Retrier, |
94 | | #[metric(help = "The number of seconds to consider an object expired")] |
95 | | consider_expired_after_s: i64, |
96 | | #[metric(help = "The number of bytes to buffer for retrying requests")] |
97 | | max_retry_buffer_per_request: usize, |
98 | | #[metric(help = "The number of concurrent uploads allowed for multipart uploads")] |
99 | | multipart_max_concurrent_uploads: usize, |
100 | | |
101 | | remove_callbacks: Mutex<Vec<RemoveCallback>>, |
102 | | } |
103 | | |
104 | | impl<I, NowFn> S3Store<NowFn> |
105 | | where |
106 | | I: InstantWrapper, |
107 | | NowFn: Fn() -> I + Send + Sync + Unpin + 'static, |
108 | | { |
109 | 0 | pub async fn new(spec: &ExperimentalAwsSpec, now_fn: NowFn) -> Result<Arc<Self>, Error> { |
110 | 0 | let jitter_fn = spec.common.retry.make_jitter_fn(); |
111 | 0 | let s3_client = { |
112 | 0 | let http_client = TlsClient::new(&spec.common.clone()); |
113 | | |
114 | 0 | let credential_provider = credentials::DefaultCredentialsChain::builder() |
115 | 0 | .configure( |
116 | 0 | ProviderConfig::without_region() |
117 | 0 | .with_region(Some(Region::new(Cow::Owned(spec.region.clone())))) |
118 | 0 | .with_http_client(http_client.clone()), |
119 | 0 | ) |
120 | 0 | .build() |
121 | 0 | .await; |
122 | | |
123 | 0 | let config = aws_config::defaults(BehaviorVersion::latest()) |
124 | 0 | .credentials_provider(credential_provider) |
125 | 0 | .app_name(AppName::new("nativelink").expect("valid app name")) |
126 | 0 | .timeout_config( |
127 | 0 | aws_config::timeout::TimeoutConfig::builder() |
128 | 0 | .connect_timeout(Duration::from_secs(15)) |
129 | 0 | .build(), |
130 | 0 | ) |
131 | 0 | .region(Region::new(Cow::Owned(spec.region.clone()))) |
132 | 0 | .http_client(http_client) |
133 | 0 | .load() |
134 | 0 | .await; |
135 | | |
136 | 0 | Client::new(&config) |
137 | | }; |
138 | 0 | Self::new_with_client_and_jitter(spec, s3_client, jitter_fn, now_fn) |
139 | 0 | } |
140 | | |
141 | 19 | pub fn new_with_client_and_jitter( |
142 | 19 | spec: &ExperimentalAwsSpec, |
143 | 19 | s3_client: Client, |
144 | 19 | jitter_fn: Arc<dyn Fn(Duration) -> Duration + Send + Sync>, |
145 | 19 | now_fn: NowFn, |
146 | 19 | ) -> Result<Arc<Self>, Error> { |
147 | 19 | Ok(Arc::new(Self { |
148 | 19 | s3_client: Arc::new(s3_client), |
149 | 19 | now_fn, |
150 | 19 | bucket: spec.bucket.clone(), |
151 | 19 | key_prefix: spec |
152 | 19 | .common |
153 | 19 | .key_prefix |
154 | 19 | .as_ref() |
155 | 19 | .unwrap_or(&String::new()) |
156 | 19 | .clone(), |
157 | 19 | retrier: Retrier::new( |
158 | 19 | Arc::new(|duration| Box::pin2 (sleep2 (duration2 ))), |
159 | 19 | jitter_fn, |
160 | 19 | spec.common.retry.clone(), |
161 | | ), |
162 | 19 | consider_expired_after_s: i64::from(spec.common.consider_expired_after_s), |
163 | 19 | max_retry_buffer_per_request: spec |
164 | 19 | .common |
165 | 19 | .max_retry_buffer_per_request |
166 | 19 | .unwrap_or(DEFAULT_MAX_RETRY_BUFFER_PER_REQUEST), |
167 | 19 | multipart_max_concurrent_uploads: spec |
168 | 19 | .common |
169 | 19 | .multipart_max_concurrent_uploads |
170 | 19 | .map_or(DEFAULT_MULTIPART_MAX_CONCURRENT_UPLOADS, |v| v), |
171 | 19 | remove_callbacks: Mutex::new(Vec::new()), |
172 | | })) |
173 | 19 | } |
174 | | |
175 | 19 | fn make_s3_path(&self, key: &StoreKey<'_>) -> String { |
176 | 19 | format!("{}{}", self.key_prefix, key.as_str()) |
177 | 19 | } |
178 | | |
179 | 7 | async fn has(self: Pin<&Self>, digest: StoreKey<'_>) -> Result<Option<u64>, Error> { |
180 | 7 | let digest_clone = digest.into_owned(); |
181 | 7 | self.retrier |
182 | 8 | .retry7 (unfold7 (()7 , move |state| { |
183 | 8 | let local_digest = digest_clone.clone(); |
184 | 8 | async move { |
185 | 8 | let result = self |
186 | 8 | .s3_client |
187 | 8 | .head_object() |
188 | 8 | .bucket(&self.bucket) |
189 | 8 | .key(self.make_s3_path(&local_digest)) |
190 | 8 | .send() |
191 | 8 | .await; |
192 | | |
193 | 8 | match result { |
194 | 5 | Ok(head_object_output) => { |
195 | 5 | if self.consider_expired_after_s != 0 |
196 | 2 | && let Some(last_modified) = head_object_output.last_modified |
197 | | { |
198 | 2 | let now_s = (self.now_fn)() |
199 | 2 | .unix_timestamp() |
200 | 2 | .try_into() |
201 | 2 | .unwrap_or(i64::MAX); |
202 | 2 | if last_modified.secs() + self.consider_expired_after_s <= now_s { |
203 | 1 | let remove_callbacks = self.remove_callbacks.lock().clone(); |
204 | 1 | let mut callbacks: FuturesUnordered<_> = remove_callbacks |
205 | 1 | .iter() |
206 | 1 | .map(|callback| callback0 .callback0 (local_digest0 .borrow0 ())) |
207 | 1 | .collect(); |
208 | 1 | while callbacks.next().await.is_some() {}0 |
209 | 1 | return Some((RetryResult::Ok(None), state)); |
210 | 1 | } |
211 | 3 | } |
212 | 4 | let Some(length) = head_object_output.content_length else { |
213 | 0 | return Some((RetryResult::Ok(None), state)); |
214 | | }; |
215 | 4 | if length >= 0 { |
216 | 4 | return Some((RetryResult::Ok(Some(length as u64)), state)); |
217 | 0 | } |
218 | 0 | Some(( |
219 | 0 | RetryResult::Err(make_err!( |
220 | 0 | Code::InvalidArgument, |
221 | 0 | "Negative content length in S3: {length:?}", |
222 | 0 | )), |
223 | 0 | state, |
224 | 0 | )) |
225 | | } |
226 | 3 | Err(sdk_error) => match sdk_error.into_service_error() { |
227 | 2 | HeadObjectError::NotFound(_) => Some((RetryResult::Ok(None), state)), |
228 | 1 | other => Some(( |
229 | 1 | RetryResult::Retry( |
230 | 1 | Error::from_std_err(Code::Unavailable, &other) |
231 | 1 | .append("Unhandled HeadObjectError in S3"), |
232 | 1 | ), |
233 | 1 | state, |
234 | 1 | )), |
235 | | }, |
236 | | } |
237 | 8 | } |
238 | 8 | })) |
239 | 7 | .await |
240 | 7 | } |
241 | | } |
242 | | |
243 | | #[async_trait] |
244 | | impl<I, NowFn> StoreDriver for S3Store<NowFn> |
245 | | where |
246 | | I: InstantWrapper, |
247 | | NowFn: Fn() -> I + Send + Sync + Unpin + 'static, |
248 | | { |
249 | 0 | async fn post_init(self: Arc<Self>) -> Result<(), Error> { |
250 | | Ok(()) |
251 | 0 | } |
252 | | |
253 | | async fn has_with_results( |
254 | | self: Pin<&Self>, |
255 | | keys: &[StoreKey<'_>], |
256 | | results: &mut [Option<u64>], |
257 | 8 | ) -> Result<(), Error> { |
258 | | keys.iter() |
259 | | .zip(results.iter_mut()) |
260 | 8 | .map(|(key, result)| async move { |
261 | | // We need to do a special pass to ensure our zero key exist. |
262 | 8 | if is_zero_digest(key.borrow()) { |
263 | 1 | *result = Some(0); |
264 | 1 | return Ok::<_, Error>(()); |
265 | 7 | } |
266 | 7 | *result = self.has(key.borrow()).await?0 ; |
267 | 7 | Ok::<_, Error>(()) |
268 | 16 | }) |
269 | | .collect::<FuturesUnordered<_>>() |
270 | | .try_collect() |
271 | | .await |
272 | 8 | } |
273 | | |
274 | 0 | fn optimized_for(&self, optimization: StoreOptimizations) -> bool { |
275 | 0 | matches!(optimization, StoreOptimizations::LazyExistenceOnSync) |
276 | 0 | } |
277 | | |
278 | | async fn update( |
279 | | self: Pin<&Self>, |
280 | | digest: StoreKey<'_>, |
281 | | mut reader: DropCloserReadHalf, |
282 | | upload_size: UploadSizeInfo, |
283 | 4 | ) -> Result<u64, Error> { |
284 | | let s3_path = &self.make_s3_path(&digest); |
285 | | |
286 | | let max_size = match upload_size { |
287 | | UploadSizeInfo::ExactSize(sz) | UploadSizeInfo::MaxSize(sz) => sz, |
288 | | }; |
289 | | |
290 | | // Sanity check S3 maximum upload size. |
291 | | if max_size > MAX_UPLOAD_SIZE { |
292 | | return Err(make_err!( |
293 | | Code::FailedPrecondition, |
294 | | "File size exceeds max of {MAX_UPLOAD_SIZE}" |
295 | | )); |
296 | | } |
297 | | |
298 | | // Note(aaronmondal) It might be more optimal to use a different |
299 | | // heuristic here, but for simplicity we use a hard coded value. |
300 | | // Anything going down this if-statement will have the advantage of only |
301 | | // 1 network request for the upload instead of minimum of 3 required for |
302 | | // multipart upload requests. |
303 | | // |
304 | | // Note(aaronmondal) If the upload size is not known, we go down the multipart upload path. |
305 | | // This is not very efficient, but it greatly reduces the complexity of the code. |
306 | | if max_size < MIN_MULTIPART_SIZE && matches!(upload_size, UploadSizeInfo::ExactSize(_)) { |
307 | | let UploadSizeInfo::ExactSize(sz) = upload_size else { |
308 | | unreachable!("upload_size must be UploadSizeInfo::ExactSize here"); |
309 | | }; |
310 | | reader.set_max_recent_data_size( |
311 | | u64::try_from(self.max_retry_buffer_per_request) |
312 | | .err_tip(|| "Could not convert max_retry_buffer_per_request to u64")?, |
313 | | ); |
314 | | return self |
315 | | .retrier |
316 | 2 | .retry(unfold(reader, move |mut reader| async move { |
317 | | // We need to make a new pair here because the aws sdk does not give us |
318 | | // back the body after we send it in order to retry. |
319 | 2 | let (mut tx, rx) = make_buf_channel_pair(); |
320 | | |
321 | | // Upload the data to the S3 backend. |
322 | 2 | let result = { |
323 | 2 | let reader_ref = &mut reader; |
324 | 2 | let (upload_res, bind_res) = tokio::join!( |
325 | 2 | self.s3_client |
326 | 2 | .put_object() |
327 | 2 | .bucket(&self.bucket) |
328 | 2 | .key(s3_path.clone()) |
329 | 2 | .content_length(sz.try_into().unwrap_or(i64::MAX)) |
330 | 2 | .body(ByteStream::from_body_1_x(BodyWrapper { |
331 | 2 | reader: rx, |
332 | 2 | size: sz, |
333 | 2 | })) |
334 | 2 | .send() |
335 | 2 | .map_ok_or_else(|e| Err(Error::from_std_err0 (Code::Aborted0 , &e0 )), |_| Ok(sz)), |
336 | | // Stream all data from the reader channel to the writer channel. |
337 | 2 | tx.bind_buffered(reader_ref) |
338 | | ); |
339 | 2 | match (upload_res, bind_res) { |
340 | 2 | (Ok(size), Ok(())) => Ok(size), |
341 | 0 | (Err(e), _) | (_, Err(e)) => Err(e), |
342 | | } |
343 | 2 | .err_tip(|| "Failed to upload file to s3 in single chunk") |
344 | | }; |
345 | | |
346 | | // If we failed to upload the file, check to see if we can retry. |
347 | 2 | let retry_result = result.map_or_else(|mut err| {0 |
348 | | // Ensure our code is Code::Aborted, so the client can retry if possible. |
349 | 0 | err.code = Code::Aborted; |
350 | 0 | let bytes_received = reader.get_bytes_received(); |
351 | 0 | if let Err(try_reset_err) = reader.try_reset_stream() { |
352 | 0 | error!( |
353 | | ?bytes_received, |
354 | | err = ?try_reset_err, |
355 | | "Unable to reset stream after failed upload in S3Store::update" |
356 | | ); |
357 | 0 | return RetryResult::Err(err |
358 | 0 | .merge(try_reset_err) |
359 | 0 | .append(format!("Failed to retry upload with {bytes_received} bytes received in S3Store::update"))); |
360 | 0 | } |
361 | 0 | let err = err.append(format!("Retry on upload happened with {bytes_received} bytes received in S3Store::update")); |
362 | 0 | info!( |
363 | | ?err, |
364 | | ?bytes_received, |
365 | | "Retryable S3 error" |
366 | | ); |
367 | 0 | RetryResult::Retry(err) |
368 | 0 | }, RetryResult::Ok); |
369 | 2 | Some((retry_result, reader)) |
370 | 4 | })) |
371 | | .await; |
372 | | } |
373 | | |
374 | | let upload_id = &self |
375 | | .retrier |
376 | 2 | .retry(unfold((), move |()| async move { |
377 | 2 | let retry_result = self |
378 | 2 | .s3_client |
379 | 2 | .create_multipart_upload() |
380 | 2 | .bucket(&self.bucket) |
381 | 2 | .key(s3_path) |
382 | 2 | .send() |
383 | 2 | .await |
384 | 2 | .map_or_else( |
385 | 0 | |e| { |
386 | 0 | RetryResult::Retry( |
387 | 0 | Error::from_std_err(Code::Aborted, &e) |
388 | 0 | .append("Failed to create multipart upload to s3"), |
389 | 0 | ) |
390 | 0 | }, |
391 | 2 | |CreateMultipartUploadOutput { upload_id, .. }| { |
392 | 2 | upload_id.map_or_else( |
393 | 0 | || { |
394 | 0 | RetryResult::Err(make_err!( |
395 | 0 | Code::Internal, |
396 | 0 | "Expected upload_id to be set by s3 response" |
397 | 0 | )) |
398 | 0 | }, |
399 | | RetryResult::Ok, |
400 | | ) |
401 | 2 | }, |
402 | | ); |
403 | 2 | Some((retry_result, ())) |
404 | 4 | })) |
405 | | .await?; |
406 | | |
407 | | // S3 requires us to upload in parts if the size is greater than 5GB. The part size must be at least |
408 | | // 5MB (except last part) and can have up to 10,000 parts. |
409 | | |
410 | | // Calculate of number of chunks if we upload in 5MB chucks (min chunk size), clamping to |
411 | | // 10,000 parts and correcting for lossy integer division. This provides the |
412 | | let chunk_count = (max_size / MIN_MULTIPART_SIZE).clamp(0, MAX_UPLOAD_PARTS - 1) + 1; |
413 | | |
414 | | // Using clamped first approximation of number of chunks, calculate byte count of each |
415 | | // chunk, excluding last chunk, clamping to min/max upload size 5MB, 5GB. |
416 | | let bytes_per_upload_part = |
417 | | (max_size / chunk_count).clamp(MIN_MULTIPART_SIZE, MAX_MULTIPART_SIZE); |
418 | | |
419 | | // Sanity check before continuing. |
420 | | if !(MIN_MULTIPART_SIZE..MAX_MULTIPART_SIZE).contains(&bytes_per_upload_part) { |
421 | | return Err(make_err!( |
422 | | Code::FailedPrecondition, |
423 | | "Failed to calculate file chuck size (min, max, calc): {MIN_MULTIPART_SIZE}, {MAX_MULTIPART_SIZE}, {bytes_per_upload_part}", |
424 | | )); |
425 | | } |
426 | | |
427 | 2 | let upload_parts = move || async move { |
428 | | // This will ensure we only have `multipart_max_concurrent_uploads` * `bytes_per_upload_part` |
429 | | // bytes in memory at any given time waiting to be uploaded. |
430 | 2 | let (tx, mut rx) = mpsc::channel(self.multipart_max_concurrent_uploads); |
431 | | |
432 | 2 | let read_stream_fut = async move { |
433 | 2 | let retrier = &Pin::get_ref(self).retrier; |
434 | 2 | let mut total_uploaded = 0; |
435 | | // Note: Our break condition is when we reach EOF. |
436 | 9 | for part_number in 1..i32::MAX2 { |
437 | 9 | let write_buf = reader |
438 | 9 | .consume(Some(usize::try_from(bytes_per_upload_part).err_tip( |
439 | | || "Could not convert bytes_per_upload_part to usize", |
440 | 0 | )?)) |
441 | 9 | .await |
442 | 9 | .err_tip(|| "Failed to read chunk in s3_store")?0 ; |
443 | 9 | if write_buf.is_empty() { |
444 | 2 | break; // Reached EOF. |
445 | 7 | } |
446 | | |
447 | 7 | total_uploaded += write_buf.len() as u64; |
448 | | |
449 | 7 | tx.send(retrier.retry(unfold(write_buf, move |write_buf| { |
450 | 7 | async move { |
451 | 7 | let retry_result = self |
452 | 7 | .s3_client |
453 | 7 | .upload_part() |
454 | 7 | .bucket(&self.bucket) |
455 | 7 | .key(s3_path) |
456 | 7 | .upload_id(upload_id) |
457 | 7 | .body(ByteStream::new(SdkBody::from(write_buf.clone()))) |
458 | 7 | .part_number(part_number) |
459 | 7 | .send() |
460 | 7 | .await |
461 | 7 | .map_or_else( |
462 | 0 | |e| { |
463 | 0 | RetryResult::Retry( |
464 | 0 | Error::from_std_err(Code::Aborted, &e).append(format!( |
465 | 0 | "Failed to upload part {part_number} in S3 store" |
466 | 0 | )), |
467 | 0 | ) |
468 | 0 | }, |
469 | 7 | |mut response| { |
470 | 7 | RetryResult::Ok( |
471 | 7 | CompletedPartBuilder::default() |
472 | 7 | // Only set an entity tag if it exists. This saves |
473 | 7 | // 13 bytes per part on the final request if it can |
474 | 7 | // omit the `<ETAG><ETAG/>` string. |
475 | 7 | .set_e_tag(response.e_tag.take()) |
476 | 7 | .part_number(part_number) |
477 | 7 | .build(), |
478 | 7 | ) |
479 | 7 | }, |
480 | | ); |
481 | 7 | Some((retry_result, write_buf)) |
482 | 7 | } |
483 | 7 | }))) |
484 | 7 | .await |
485 | 7 | .map_err(|err| {0 |
486 | 0 | Error::from_std_err(Code::Internal, &err) |
487 | 0 | .append("Failed to send part to channel in s3_store") |
488 | 0 | })?; |
489 | | } |
490 | 2 | Result::<_, Error>::Ok(total_uploaded) |
491 | 2 | } |
492 | 2 | .fuse(); |
493 | | |
494 | 2 | let mut upload_futures = FuturesUnordered::new(); |
495 | 2 | let mut total_uploaded = 0; |
496 | | |
497 | 2 | let mut completed_parts = Vec::with_capacity( |
498 | 2 | usize::try_from(cmp::min(MAX_UPLOAD_PARTS, chunk_count)) |
499 | 2 | .err_tip(|| "Could not convert u64 to usize")?0 , |
500 | | ); |
501 | 2 | tokio::pin!(read_stream_fut); |
502 | | loop { |
503 | 18 | if read_stream_fut.is_terminated() && rx16 .is_empty16 () && upload_futures5 .is_empty5 () { |
504 | 2 | break; // No more data to process. |
505 | 16 | } |
506 | 16 | tokio::select! { |
507 | 16 | result2 = &mut read_stream_fut => { |
508 | 2 | total_uploaded = result?0 ; |
509 | | }, // Return error or wait for other futures. |
510 | 16 | Some(upload_result7 ) = upload_futures.next() => completed_parts7 .push7 (upload_result7 ?0 ), |
511 | 16 | Some(fut7 ) = rx.recv() => upload_futures7 .push7 (fut7 ), |
512 | | } |
513 | | } |
514 | | |
515 | | // Even though the spec does not require parts to be sorted by number, we do it just in case |
516 | | // there's an S3 implementation that requires it. |
517 | 2 | completed_parts.sort_unstable_by_key(|part| part.part_number); |
518 | | |
519 | 2 | self.retrier |
520 | 2 | .retry(unfold(completed_parts, move |completed_parts| async move { |
521 | | Some(( |
522 | 2 | self.s3_client |
523 | 2 | .complete_multipart_upload() |
524 | 2 | .bucket(&self.bucket) |
525 | 2 | .key(s3_path) |
526 | 2 | .multipart_upload( |
527 | 2 | CompletedMultipartUploadBuilder::default() |
528 | 2 | .set_parts(Some(completed_parts.clone())) |
529 | 2 | .build(), |
530 | 2 | ) |
531 | 2 | .upload_id(upload_id) |
532 | 2 | .send() |
533 | 2 | .await |
534 | 2 | .map_or_else( |
535 | 0 | |e| { |
536 | 0 | RetryResult::Retry( |
537 | 0 | Error::from_std_err(Code::Aborted, &e).append( |
538 | 0 | "Failed to complete multipart upload in S3 store", |
539 | 0 | ), |
540 | 0 | ) |
541 | 0 | }, |
542 | 2 | |_| RetryResult::Ok(total_uploaded), |
543 | | ), |
544 | 2 | completed_parts, |
545 | | )) |
546 | 4 | })) |
547 | 2 | .await |
548 | 4 | }; |
549 | | // Upload our parts and complete the multipart upload. |
550 | | // If we fail attempt to abort the multipart upload (cleanup). |
551 | | upload_parts() |
552 | 0 | .or_else(move |mut e| async move { |
553 | 0 | let abort_res = self |
554 | 0 | .s3_client |
555 | 0 | .abort_multipart_upload() |
556 | 0 | .bucket(&self.bucket) |
557 | 0 | .key(s3_path) |
558 | 0 | .upload_id(upload_id) |
559 | 0 | .send() |
560 | 0 | .await; |
561 | 0 | if let Err(abort_err) = abort_res { |
562 | 0 | let err = Error::from_std_err(Code::Aborted, &abort_err) |
563 | 0 | .append("Failed to abort multipart upload in S3 store"); |
564 | 0 | info!(?err, "Multipart upload error"); |
565 | 0 | e = e.merge(err); |
566 | 0 | } |
567 | 0 | Err(e) |
568 | 0 | }) |
569 | | .await |
570 | 4 | } |
571 | | |
572 | | async fn get_part( |
573 | | self: Pin<&Self>, |
574 | | key: StoreKey<'_>, |
575 | | writer: &mut DropCloserWriteHalf, |
576 | | offset: u64, |
577 | | length: Option<u64>, |
578 | 8 | ) -> Result<(), Error> { |
579 | | if is_zero_digest(key.borrow()) { |
580 | | writer |
581 | | .send_eof() |
582 | | .err_tip(|| "Failed to send zero EOF in filesystem store get_part")?; |
583 | | return Ok(()); |
584 | | } |
585 | | |
586 | | let s3_path = &self.make_s3_path(&key); |
587 | | let end_read_byte = length |
588 | 3 | .map_or(Some(None), |length| Some(offset.checked_add(length))) |
589 | | .err_tip(|| "Integer overflow protection triggered")?; |
590 | | |
591 | | self.retrier |
592 | 8 | .retry(unfold(writer, move |writer| async move { |
593 | 8 | let result = self |
594 | 8 | .s3_client |
595 | 8 | .get_object() |
596 | 8 | .bucket(&self.bucket) |
597 | 8 | .key(s3_path) |
598 | 8 | .range(format!( |
599 | | "bytes={}-{}", |
600 | 8 | offset + writer.get_bytes_written(), |
601 | 8 | end_read_byte.map_or_else(String::new, |v| v3 .to_string3 ()) |
602 | | )) |
603 | 8 | .send() |
604 | 8 | .await; |
605 | | |
606 | 8 | let mut s3_in_stream6 = match result { |
607 | 6 | Ok(head_object_output) => head_object_output.body, |
608 | 2 | Err(sdk_error) => match sdk_error.into_service_error() { |
609 | 1 | GetObjectError::NoSuchKey(e) => { |
610 | 1 | return Some(( |
611 | 1 | RetryResult::Err( |
612 | 1 | Error::from_std_err(Code::NotFound, &e) |
613 | 1 | .append("No such key in S3"), |
614 | 1 | ), |
615 | 1 | writer, |
616 | 1 | )); |
617 | | } |
618 | 1 | other => { |
619 | 1 | return Some(( |
620 | 1 | RetryResult::Retry( |
621 | 1 | Error::from_std_err(Code::Unavailable, &other) |
622 | 1 | .append("Unhandled GetObjectError in S3"), |
623 | 1 | ), |
624 | 1 | writer, |
625 | 1 | )); |
626 | | } |
627 | | }, |
628 | | }; |
629 | | |
630 | | // Copy data from s3 input stream to the writer stream. |
631 | 11 | while let Some(maybe_bytes5 ) = s3_in_stream.next().await { |
632 | 5 | match maybe_bytes { |
633 | 5 | Ok(bytes) => { |
634 | 5 | if bytes.is_empty() { |
635 | | // Ignore possible EOF. Different implementations of S3 may or may not |
636 | | // send EOF this way. |
637 | 1 | continue; |
638 | 4 | } |
639 | 4 | if let Err(e0 ) = writer.send(bytes).await { |
640 | 0 | return Some(( |
641 | 0 | RetryResult::Err( |
642 | 0 | Error::from_std_err(Code::Aborted, &e) |
643 | 0 | .append("Error sending bytes to consumer in S3"), |
644 | 0 | ), |
645 | 0 | writer, |
646 | 0 | )); |
647 | 4 | } |
648 | | } |
649 | 0 | Err(e) => { |
650 | 0 | return Some(( |
651 | 0 | RetryResult::Retry( |
652 | 0 | Error::from_std_err(Code::Aborted, &e) |
653 | 0 | .append("Bad bytestream element in S3"), |
654 | 0 | ), |
655 | 0 | writer, |
656 | 0 | )); |
657 | | } |
658 | | } |
659 | | } |
660 | 6 | if let Err(e0 ) = writer.send_eof() { |
661 | 0 | return Some(( |
662 | 0 | RetryResult::Err( |
663 | 0 | Error::from_std_err(Code::Aborted, &e) |
664 | 0 | .append("Failed to send EOF to consumer in S3"), |
665 | 0 | ), |
666 | 0 | writer, |
667 | 0 | )); |
668 | 6 | } |
669 | 6 | Some((RetryResult::Ok(()), writer)) |
670 | 16 | })) |
671 | | .await |
672 | 8 | } |
673 | | |
674 | 0 | fn inner_store(&self, _digest: Option<StoreKey>) -> &'_ dyn StoreDriver { |
675 | 0 | self |
676 | 0 | } |
677 | | |
678 | 0 | fn as_any<'a>(&'a self) -> &'a (dyn core::any::Any + Sync + Send + 'static) { |
679 | 0 | self |
680 | 0 | } |
681 | | |
682 | 0 | fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send + 'static> { |
683 | 0 | self |
684 | 0 | } |
685 | | |
686 | 0 | fn register_health(self: Arc<Self>, registry: &mut HealthRegistryBuilder) { |
687 | 0 | registry.register_indicator(self); |
688 | 0 | } |
689 | | |
690 | 0 | fn register_remove_callback(self: Arc<Self>, callback: RemoveCallback) -> Result<(), Error> { |
691 | 0 | self.remove_callbacks.lock().push(callback); |
692 | 0 | Ok(()) |
693 | 0 | } |
694 | | } |
695 | | |
696 | | #[async_trait] |
697 | | impl<I, NowFn> HealthStatusIndicator for S3Store<NowFn> |
698 | | where |
699 | | I: InstantWrapper, |
700 | | NowFn: Fn() -> I + Send + Sync + Unpin + 'static, |
701 | | { |
702 | 0 | fn get_name(&self) -> &'static str { |
703 | 0 | "S3Store" |
704 | 0 | } |
705 | | |
706 | 0 | async fn check_health(&self, namespace: Cow<'static, str>) -> HealthStatus { |
707 | | StoreDriver::check_health(Pin::new(self), namespace).await |
708 | 0 | } |
709 | | } |