Coverage Report

Created: 2026-11-13 23:29

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