Coverage Report

Created: 2026-07-29 18:50

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/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
}