Skip to main content

risingwave_object_store/object/
s3.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (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//     http://www.apache.org/licenses/LICENSE-2.0
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
15use std::borrow::BorrowMut;
16use std::cmp;
17use std::collections::VecDeque;
18use std::pin::Pin;
19use std::sync::Arc;
20use std::task::{Context, Poll, ready};
21use std::time::Duration;
22
23use await_tree::{InstrumentAwait, SpanExt};
24use aws_sdk_s3::Client;
25use aws_sdk_s3::config::{Credentials, Region};
26use aws_sdk_s3::error::BoxError;
27use aws_sdk_s3::operation::abort_multipart_upload::AbortMultipartUploadError;
28use aws_sdk_s3::operation::complete_multipart_upload::CompleteMultipartUploadError;
29use aws_sdk_s3::operation::create_multipart_upload::CreateMultipartUploadError;
30use aws_sdk_s3::operation::delete_object::DeleteObjectError;
31use aws_sdk_s3::operation::delete_objects::DeleteObjectsError;
32use aws_sdk_s3::operation::get_object::GetObjectError;
33use aws_sdk_s3::operation::get_object::builders::GetObjectFluentBuilder;
34use aws_sdk_s3::operation::head_object::HeadObjectError;
35use aws_sdk_s3::operation::list_objects_v2::ListObjectsV2Error;
36use aws_sdk_s3::operation::put_object::PutObjectError;
37use aws_sdk_s3::operation::upload_part::UploadPartOutput;
38use aws_sdk_s3::primitives::ByteStream;
39use aws_sdk_s3::types::{
40    AbortIncompleteMultipartUpload, BucketLifecycleConfiguration, CompletedMultipartUpload,
41    CompletedPart, Delete, ExpirationStatus, LifecycleRule, LifecycleRuleFilter, ObjectIdentifier,
42};
43use aws_smithy_http::futures_stream_adapter::FuturesStreamCompatByteStream;
44use aws_smithy_http_client::{
45    Builder as SmithyHttpClientBuilder, Connector as SmithyConnector, tls,
46};
47use aws_smithy_runtime_api::client::http::HttpClient;
48use aws_smithy_runtime_api::client::result::SdkError;
49use aws_smithy_types::body::SdkBody;
50use aws_smithy_types::error::metadata::ProvideErrorMetadata;
51use bytes::BytesMut;
52use fail::fail_point;
53use futures::future::{BoxFuture, FutureExt, try_join_all};
54use futures::{Stream, StreamExt, TryStreamExt};
55use itertools::Itertools;
56use risingwave_common::config::ObjectStoreConfig;
57use risingwave_common::range::RangeBoundsExt;
58use thiserror_ext::AsReport;
59use tokio::task::JoinHandle;
60
61use super::object_metrics::ObjectStoreMetrics;
62use super::{
63    Bytes, ObjectError, ObjectErrorInner, ObjectMetadata, ObjectRangeBounds, ObjectResult,
64    ObjectStore, StreamingUploader, prefix, retry_request,
65};
66use crate::object::{
67    ObjectDataStream, ObjectMetadataIter, OperationType, try_update_failure_metric,
68};
69
70type PartId = i32;
71
72/// MinIO and S3 share the same minimum part ID and part size.
73const MIN_PART_ID: PartId = 1;
74/// Stop multipart uploads that don't complete within a specified number of days after being
75/// initiated. (Day is the smallest granularity)
76const S3_INCOMPLETE_MULTIPART_UPLOAD_RETENTION_DAYS: i32 = 1;
77
78/// S3 multipart upload handle. The multipart upload is not initiated until the first part is
79/// available for upload.
80///
81/// Reference: <https://docs.aws.amazon.com/AmazonS3/latest/userguide/mpuoverview.html>
82pub struct S3StreamingUploader {
83    client: Client,
84    part_size: usize,
85    bucket: String,
86    /// The key of the object.
87    key: String,
88    /// The identifier of multipart upload task for S3.
89    upload_id: Option<String>,
90    /// Next part ID.
91    next_part_id: PartId,
92    /// Join handles for part uploads.
93    join_handles: Vec<JoinHandle<ObjectResult<(PartId, UploadPartOutput)>>>,
94    /// Buffer for data. It will store at least `part_size` bytes of data before wrapping itself
95    /// into a stream and upload to object store as a part.
96    buf: Vec<Bytes>,
97    /// Length of the data that have not been uploaded to S3.
98    not_uploaded_len: usize,
99    /// To record metrics for uploading part.
100    metrics: Arc<ObjectStoreMetrics>,
101
102    config: Arc<ObjectStoreConfig>,
103}
104
105impl S3StreamingUploader {
106    const MEDIA_TYPE: &'static str = "s3";
107
108    pub fn new(
109        client: Client,
110        bucket: String,
111        key: String,
112        metrics: Arc<ObjectStoreMetrics>,
113        config: Arc<ObjectStoreConfig>,
114    ) -> S3StreamingUploader {
115        /// The minimum number of bytes that is buffered before they are uploaded as a part.
116        /// Its value must be greater than the minimum part size of 5MiB.
117        ///
118        /// Reference: <https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html>
119        const MIN_PART_SIZE: usize = 5 * 1024 * 1024;
120        const MAX_PART_SIZE: usize = 5 * 1024 * 1024 * 1024;
121        let part_size = config.upload_part_size.clamp(MIN_PART_SIZE, MAX_PART_SIZE);
122
123        Self {
124            client,
125            bucket,
126            part_size,
127            key,
128            upload_id: None,
129            next_part_id: MIN_PART_ID,
130            join_handles: Default::default(),
131            buf: Default::default(),
132            not_uploaded_len: 0,
133            metrics,
134            config,
135        }
136    }
137
138    async fn upload_next_part(&mut self) -> ObjectResult<()> {
139        let operation_type = OperationType::StreamingUpload;
140        let operation_type_str = operation_type.as_str();
141
142        // Lazily create multipart upload.
143        if self.upload_id.is_none() {
144            let builder = || async {
145                self.client
146                    .create_multipart_upload()
147                    .bucket(&self.bucket)
148                    .key(&self.key)
149                    .send()
150                    .await
151                    .map_err(|err| {
152                        set_error_should_retry::<CreateMultipartUploadError>(
153                            self.config.clone(),
154                            err.into(),
155                        )
156                    })
157            };
158
159            let resp = retry_request(
160                builder,
161                &self.config,
162                OperationType::StreamingUploadInit,
163                self.metrics.clone(),
164                Self::MEDIA_TYPE,
165            )
166            .await;
167
168            try_update_failure_metric(
169                &self.metrics,
170                &resp,
171                OperationType::StreamingUploadInit.as_str(),
172            );
173
174            self.upload_id = Some(resp?.upload_id.unwrap());
175        }
176
177        // Get the data to upload for the next part.
178        let data = self.buf.drain(..).collect_vec();
179        let len = self.not_uploaded_len;
180        debug_assert_eq!(
181            data.iter().map(|b| b.len()).sum::<usize>(),
182            self.not_uploaded_len
183        );
184
185        // Update part id.
186        let part_id = self.next_part_id;
187        self.next_part_id += 1;
188
189        // Clone the variables to be passed into the upload join handle.
190        let client_cloned = self.client.clone();
191        let bucket = self.bucket.clone();
192        let key = self.key.clone();
193        let upload_id = self.upload_id.clone().unwrap();
194
195        let metrics = self.metrics.clone();
196        metrics
197            .operation_size
198            .with_label_values(&[operation_type_str])
199            .observe(len as f64);
200        let config = self.config.clone();
201
202        self.join_handles.push(tokio::spawn(async move {
203            let _timer = metrics
204                .operation_latency
205                .with_label_values(&["s3", operation_type_str])
206                .start_timer();
207
208            let builder = || async {
209                client_cloned
210                    .upload_part()
211                    .bucket(bucket.clone())
212                    .key(key.clone())
213                    .upload_id(upload_id.clone())
214                    .part_number(part_id)
215                    .body(get_upload_body(data.clone()))
216                    .content_length(len as i64)
217                    .send()
218                    .await
219                    .map_err(|err| {
220                        set_error_should_retry::<CreateMultipartUploadError>(
221                            config.clone(),
222                            err.into(),
223                        )
224                    })
225            };
226
227            let res = retry_request(
228                builder,
229                &config,
230                operation_type,
231                metrics.clone(),
232                Self::MEDIA_TYPE,
233            )
234            .await;
235            try_update_failure_metric(&metrics, &res, operation_type_str);
236            Ok((part_id, res?))
237        }));
238
239        Ok(())
240    }
241
242    async fn flush_multipart_and_complete(&mut self) -> ObjectResult<()> {
243        let operation_type = OperationType::StreamingUploadFinish;
244
245        if !self.buf.is_empty() {
246            self.upload_next_part().await?;
247        }
248
249        // If any part fails to upload, abort the upload.
250        let join_handles = self.join_handles.drain(..).collect_vec();
251
252        let mut uploaded_parts = Vec::with_capacity(join_handles.len());
253        for result in try_join_all(join_handles)
254            .await
255            .map_err(ObjectError::internal)?
256        {
257            uploaded_parts.push(result?);
258        }
259
260        let completed_parts = Some(
261            uploaded_parts
262                .iter()
263                .map(|(part_id, output)| {
264                    CompletedPart::builder()
265                        .set_e_tag(output.e_tag.clone())
266                        .set_part_number(Some(*part_id))
267                        .build()
268                })
269                .collect_vec(),
270        );
271
272        let builder = || async {
273            self.client
274                .complete_multipart_upload()
275                .bucket(&self.bucket)
276                .key(&self.key)
277                .upload_id(self.upload_id.as_ref().unwrap())
278                .multipart_upload(
279                    CompletedMultipartUpload::builder()
280                        .set_parts(completed_parts.clone())
281                        .build(),
282                )
283                .send()
284                .await
285                .map_err(|err| {
286                    set_error_should_retry::<CompleteMultipartUploadError>(
287                        self.config.clone(),
288                        err.into(),
289                    )
290                })
291        };
292
293        let res = retry_request(
294            builder,
295            &self.config,
296            operation_type,
297            self.metrics.clone(),
298            Self::MEDIA_TYPE,
299        )
300        .await;
301        try_update_failure_metric(&self.metrics, &res, operation_type.as_str());
302        let _res = res?;
303
304        Ok(())
305    }
306
307    async fn abort_multipart_upload(&self) -> ObjectResult<()> {
308        self.client
309            .abort_multipart_upload()
310            .bucket(&self.bucket)
311            .key(&self.key)
312            .upload_id(self.upload_id.as_ref().unwrap())
313            .send()
314            .await
315            .map_err(|err| {
316                set_error_should_retry::<AbortMultipartUploadError>(self.config.clone(), err.into())
317            })?;
318        Ok(())
319    }
320}
321
322impl StreamingUploader for S3StreamingUploader {
323    async fn write_bytes(&mut self, data: Bytes) -> ObjectResult<()> {
324        fail_point!("s3_write_bytes_err", |_| Err(ObjectError::internal(
325            "s3 write bytes error"
326        )));
327        let data_len = data.len();
328        self.not_uploaded_len += data_len;
329        self.buf.push(data);
330
331        if self.not_uploaded_len >= self.part_size {
332            self.upload_next_part()
333                .instrument_await("s3_upload_next_part".verbose())
334                .await?;
335            self.not_uploaded_len = 0;
336        }
337        Ok(())
338    }
339
340    /// If the multipart upload has not been initiated, we can use `PutObject` instead to save the
341    /// `CreateMultipartUpload` and `CompleteMultipartUpload` requests. Otherwise flush the
342    /// remaining data of the buffer to S3 as a new part.
343    async fn finish(mut self) -> ObjectResult<()> {
344        fail_point!("s3_finish_streaming_upload_err", |_| Err(
345            ObjectError::internal("s3 finish streaming upload error")
346        ));
347
348        if self.upload_id.is_none() {
349            debug_assert!(self.join_handles.is_empty());
350            if self.buf.is_empty() {
351                debug_assert_eq!(self.not_uploaded_len, 0);
352                Err(ObjectError::internal("upload empty object"))
353            } else {
354                let operation_type = OperationType::Upload;
355                let builder = || async {
356                    self.client
357                        .put_object()
358                        .bucket(&self.bucket)
359                        .body(get_upload_body(self.buf.clone()))
360                        .content_length(self.not_uploaded_len as i64)
361                        .key(&self.key)
362                        .send()
363                        .instrument_await("s3_put_object".verbose())
364                        .await
365                        .map_err(|err| {
366                            set_error_should_retry::<PutObjectError>(
367                                self.config.clone(),
368                                err.into(),
369                            )
370                        })
371                };
372
373                let res = retry_request(
374                    builder,
375                    &self.config,
376                    operation_type,
377                    self.metrics.clone(),
378                    Self::MEDIA_TYPE,
379                )
380                .await;
381                try_update_failure_metric(&self.metrics, &res, operation_type.as_str());
382                res?;
383                Ok(())
384            }
385        } else {
386            match self
387                .flush_multipart_and_complete()
388                .instrument_await("s3_flush_multipart_and_complete".verbose())
389                .await
390            {
391                Err(e) => {
392                    tracing::warn!(key = self.key, error = %e.as_report(), "Failed to upload object");
393                    self.abort_multipart_upload().await?;
394                    Err(e)
395                }
396                _ => Ok(()),
397            }
398        }
399    }
400
401    fn get_memory_usage(&self) -> u64 {
402        self.part_size as u64
403    }
404}
405
406fn get_upload_body(data: Vec<Bytes>) -> ByteStream {
407    // `ByteStream` is retryable when created from in-memory data.
408    // This code path is used for non-multipart uploads, so a copy is acceptable.
409    let total_len: usize = data.iter().map(|b| b.len()).sum();
410    let mut buf = BytesMut::with_capacity(total_len);
411    for chunk in data {
412        buf.extend_from_slice(&chunk);
413    }
414    ByteStream::from(buf.freeze())
415}
416
417/// Object store with S3 backend
418/// The full path to a file on S3 would be `s3://bucket/<data_directory>/prefix/file`
419#[derive(Clone)]
420pub struct S3ObjectStore {
421    client: Client,
422    bucket: String,
423    /// For S3 specific metrics.
424    metrics: Arc<ObjectStoreMetrics>,
425
426    config: Arc<ObjectStoreConfig>,
427}
428
429#[async_trait::async_trait]
430impl ObjectStore for S3ObjectStore {
431    type StreamingUploader = S3StreamingUploader;
432
433    fn get_object_prefix(&self, obj_id: u64, _use_new_object_prefix_strategy: bool) -> String {
434        // Delegate to static method to avoid creating an `S3ObjectStore` in unit test.
435        // Using aws s3 sdk as object storage, the object prefix will be divided by default.
436        prefix::s3::get_object_prefix(obj_id)
437    }
438
439    async fn upload(&self, path: &str, obj: Bytes) -> ObjectResult<()> {
440        fail_point!("s3_upload_err", |_| Err(ObjectError::internal(
441            "s3 upload error"
442        )));
443        if obj.is_empty() {
444            Err(ObjectError::internal("upload empty object"))
445        } else {
446            self.client
447                .put_object()
448                .bucket(&self.bucket)
449                .body(ByteStream::from(obj))
450                .key(path)
451                .send()
452                .await
453                .map_err(|err| {
454                    set_error_should_retry::<PutObjectError>(self.config.clone(), err.into())
455                })?;
456            Ok(())
457        }
458    }
459
460    async fn streaming_upload(&self, path: &str) -> ObjectResult<Self::StreamingUploader> {
461        fail_point!("s3_streaming_upload_err", |_| Err(ObjectError::internal(
462            "s3 streaming upload error"
463        )));
464        Ok(S3StreamingUploader::new(
465            self.client.clone(),
466            self.bucket.clone(),
467            path.to_owned(),
468            self.metrics.clone(),
469            self.config.clone(),
470        ))
471    }
472
473    /// Amazon S3 doesn't support retrieving multiple ranges of data per GET request.
474    async fn read(&self, path: &str, range: impl ObjectRangeBounds) -> ObjectResult<Bytes> {
475        fail_point!("s3_read_err", |_| Err(ObjectError::internal(
476            "s3 read error"
477        )));
478
479        let val = match self.obj_store_request(path, range.clone()).send().await {
480            Ok(resp) => resp
481                .body
482                .collect()
483                .await
484                .map_err(|err| {
485                    set_error_should_retry::<GetObjectError>(self.config.clone(), err.into())
486                })?
487                .into_bytes(),
488            Err(sdk_err) => {
489                return Err(set_error_should_retry::<GetObjectError>(
490                    self.config.clone(),
491                    sdk_err.into(),
492                ));
493            }
494        };
495
496        if let Some(len) = range.len()
497            && len != val.len()
498        {
499            return Err(ObjectError::internal(format!(
500                "mismatched size: expected {}, found {} when reading {} at {:?}",
501                len,
502                val.len(),
503                path,
504                range,
505            )));
506        }
507
508        Ok(val)
509    }
510
511    async fn metadata(&self, path: &str) -> ObjectResult<ObjectMetadata> {
512        fail_point!("s3_metadata_err", |_| Err(ObjectError::internal(
513            "s3 metadata error"
514        )));
515        let resp = self
516            .client
517            .head_object()
518            .bucket(&self.bucket)
519            .key(path)
520            .send()
521            .await
522            .map_err(|err| {
523                set_error_should_retry::<HeadObjectError>(self.config.clone(), err.into())
524            })?;
525        Ok(ObjectMetadata {
526            key: path.to_owned(),
527            last_modified: resp
528                .last_modified()
529                .expect("last_modified required")
530                .as_secs_f64(),
531            total_size: resp.content_length.unwrap_or_default() as usize,
532        })
533    }
534
535    /// Returns a stream reading the object specified in `path`. If given, the stream starts at the
536    /// byte with index `start_pos` (0-based). As far as possible, the stream only loads the amount
537    /// of data into memory that is read from the stream.
538    async fn streaming_read(
539        &self,
540        path: &str,
541        range: impl ObjectRangeBounds,
542    ) -> ObjectResult<ObjectDataStream> {
543        fail_point!("s3_streaming_read_init_err", |_| Err(
544            ObjectError::internal("s3 streaming read init error")
545        ));
546
547        let resp = match self.obj_store_request(path, range.clone()).send().await {
548            Ok(resp) => resp,
549            Err(sdk_err) => {
550                return Err(set_error_should_retry::<GetObjectError>(
551                    self.config.clone(),
552                    sdk_err.into(),
553                ));
554            }
555        };
556
557        let reader = FuturesStreamCompatByteStream::new(resp.body);
558
559        Ok(Box::pin(
560            reader
561                .into_stream()
562                .map(|item| item.map_err(ObjectError::from)),
563        ))
564    }
565
566    /// Permanently deletes the whole object.
567    /// According to Amazon S3, this will simply return Ok if the object does not exist.
568    async fn delete(&self, path: &str) -> ObjectResult<()> {
569        fail_point!("s3_delete_err", |_| Err(ObjectError::internal(
570            "s3 delete error"
571        )));
572        self.client
573            .delete_object()
574            .bucket(&self.bucket)
575            .key(path)
576            .send()
577            .await
578            .map_err(|err| {
579                set_error_should_retry::<DeleteObjectError>(self.config.clone(), err.into())
580            })?;
581        Ok(())
582    }
583
584    /// Deletes the objects with the given paths permanently from the storage. If an object
585    /// specified in the request is not found, it will be considered as successfully deleted.
586    ///
587    /// Uses AWS' `DeleteObjects` API. See [AWS Docs](https://docs.aws.amazon.com/AmazonS3/latest/API/API_DeleteObjects.html) for more details.
588    async fn delete_objects(&self, paths: &[String]) -> ObjectResult<()> {
589        // AWS restricts the number of objects per request to 1000.
590        const MAX_LEN: usize = 1000;
591        let mut all_errors = Vec::new();
592
593        // If needed, split given set into subsets of size with no more than `MAX_LEN` objects.
594        for start_idx /* inclusive */ in (0..paths.len()).step_by(MAX_LEN) {
595            let end_idx /* exclusive */ = cmp::min(paths.len(), start_idx + MAX_LEN);
596            let slice = &paths[start_idx..end_idx];
597            // Create identifiers from paths.
598            let mut obj_ids = Vec::with_capacity(slice.len());
599            for path in slice {
600                obj_ids.push(ObjectIdentifier::builder().key(path).build().unwrap());
601            }
602
603            // Build and submit request to delete objects.
604            let delete_builder = Delete::builder().set_objects(Some(obj_ids));
605            let delete_output = self
606                .client
607                .delete_objects()
608                .bucket(&self.bucket)
609                .delete(delete_builder.build().unwrap()).send()
610                .await.map_err(|err| {
611                    set_error_should_retry::<DeleteObjectsError>(self.config.clone(),err.into())
612                })?;
613
614            // Check if there were errors.
615            if !delete_output.errors().is_empty() {
616                all_errors.append(&mut delete_output.errors().to_owned());
617            }
618        }
619        if !all_errors.is_empty() {
620            return Err(ObjectError::internal(format!(
621                "DeleteObjects request returned exception for some objects: {:?}",
622                all_errors
623            )));
624        }
625
626        Ok(())
627    }
628
629    async fn list(
630        &self,
631        prefix: &str,
632        start_after: Option<String>,
633        limit: Option<usize>,
634    ) -> ObjectResult<ObjectMetadataIter> {
635        Ok(Box::pin(
636            S3ObjectIter::new(
637                self.client.clone(),
638                self.bucket.clone(),
639                prefix.to_owned(),
640                self.config.clone(),
641                start_after,
642            )
643            .take(limit.unwrap_or(usize::MAX)),
644        ))
645    }
646
647    fn store_media_type(&self) -> &'static str {
648        "s3"
649    }
650}
651
652impl S3ObjectStore {
653    pub fn new_http_client(config: &ObjectStoreConfig) -> impl HttpClient + use<> {
654        let nodelay = config.s3.nodelay;
655        let pool_idle_timeout = config.s3.keepalive_ms.map(Duration::from_millis);
656
657        // Use the Smithy default (hyper 1.x) client stack. This avoids the deprecated hyper 0.14
658        // connector path (`aws-smithy-runtime/tls-rustls`).
659        let tls_provider = tls::Provider::Rustls(tls::rustls_provider::CryptoMode::AwsLc);
660        SmithyHttpClientBuilder::new().build_with_connector_fn(move |settings, _components| {
661            let mut builder = SmithyConnector::builder();
662
663            if let Some(settings) = settings {
664                builder = builder.connector_settings(settings.clone());
665            }
666
667            if let Some(nodelay) = nodelay {
668                builder = builder.enable_tcp_nodelay(nodelay);
669            }
670
671            if let Some(pool_idle_timeout) = pool_idle_timeout {
672                builder = builder.pool_idle_timeout(pool_idle_timeout);
673            }
674
675            builder.tls_provider(tls_provider.clone()).build()
676        })
677    }
678
679    /// Creates an S3 object store from environment variable.
680    ///
681    /// See [AWS Docs](https://docs.aws.amazon.com/sdk-for-rust/latest/dg/credentials.html) on how to provide credentials and region from env variable. If you are running compute-node on EC2, no configuration is required.
682    pub async fn new_with_config(
683        bucket: String,
684        metrics: Arc<ObjectStoreMetrics>,
685        config: Arc<ObjectStoreConfig>,
686    ) -> Self {
687        let sdk_config_loader = aws_config::from_env().http_client(Self::new_http_client(&config));
688
689        // Retry 3 times if we get server-side errors or throttling errors
690        let client = match std::env::var("RW_S3_ENDPOINT") {
691            Ok(endpoint) => {
692                // s3 compatible storage
693                let is_force_path_style = match std::env::var("RW_IS_FORCE_PATH_STYLE") {
694                    Ok(value) => value == "true",
695                    Err(_) => false,
696                };
697
698                let sdk_config = sdk_config_loader.load().await;
699                #[cfg(madsim)]
700                let client = Client::new(&sdk_config);
701                #[cfg(not(madsim))]
702                let client = Client::from_conf(
703                    aws_sdk_s3::config::Builder::from(&sdk_config)
704                        .endpoint_url(endpoint)
705                        .force_path_style(is_force_path_style)
706                        .identity_cache(
707                            aws_sdk_s3::config::IdentityCache::lazy()
708                                .load_timeout(Duration::from_secs(
709                                    config.s3.identity_resolution_timeout_s,
710                                ))
711                                .build(),
712                        )
713                        .stalled_stream_protection(
714                            aws_sdk_s3::config::StalledStreamProtectionConfig::disabled(),
715                        )
716                        .build(),
717                );
718                client
719            }
720            Err(_) => {
721                // s3
722                let sdk_config = sdk_config_loader.load().await;
723                #[cfg(madsim)]
724                let client = Client::new(&sdk_config);
725                #[cfg(not(madsim))]
726                let client = Client::from_conf(
727                    aws_sdk_s3::config::Builder::from(&sdk_config)
728                        .identity_cache(
729                            aws_sdk_s3::config::IdentityCache::lazy()
730                                .load_timeout(Duration::from_secs(
731                                    config.s3.identity_resolution_timeout_s,
732                                ))
733                                .build(),
734                        )
735                        .stalled_stream_protection(
736                            aws_sdk_s3::config::StalledStreamProtectionConfig::disabled(),
737                        )
738                        .build(),
739                );
740                client
741            }
742        };
743
744        Self {
745            client,
746            bucket,
747            metrics,
748            config,
749        }
750    }
751
752    /// Creates a minio client. The server should be like `minio://key:secret@address:port/bucket`.
753    pub async fn new_minio_engine(
754        server: &str,
755        metrics: Arc<ObjectStoreMetrics>,
756        object_store_config: Arc<ObjectStoreConfig>,
757    ) -> Self {
758        let server = server.strip_prefix("minio://").unwrap();
759        let (access_key_id, rest) = server.split_once(':').unwrap();
760        let (secret_access_key, mut rest) = rest.split_once('@').unwrap();
761
762        let endpoint_prefix = if let Some(rest_stripped) = rest.strip_prefix("https://") {
763            rest = rest_stripped;
764            "https://"
765        } else if let Some(rest_stripped) = rest.strip_prefix("http://") {
766            rest = rest_stripped;
767            "http://"
768        } else {
769            "http://"
770        };
771        let (address, bucket) = rest.split_once('/').unwrap();
772
773        #[cfg(madsim)]
774        let builder = aws_sdk_s3::config::Builder::new().credentials_provider(
775            Credentials::from_keys(access_key_id, secret_access_key, None),
776        );
777        #[cfg(not(madsim))]
778        let builder = aws_sdk_s3::config::Builder::from(
779            &aws_config::ConfigLoader::default()
780                // FIXME: https://github.com/awslabs/aws-sdk-rust/issues/973
781                .credentials_provider(Credentials::from_keys(
782                    access_key_id,
783                    secret_access_key,
784                    None,
785                ))
786                .load()
787                .await,
788        )
789        .force_path_style(true)
790        .identity_cache(
791            aws_sdk_s3::config::IdentityCache::lazy()
792                .load_timeout(Duration::from_secs(
793                    object_store_config.s3.identity_resolution_timeout_s,
794                ))
795                .build(),
796        )
797        .http_client(Self::new_http_client(&object_store_config))
798        .behavior_version_latest()
799        .stalled_stream_protection(aws_sdk_s3::config::StalledStreamProtectionConfig::disabled());
800        let config = builder
801            .region(Region::new("custom"))
802            .endpoint_url(format!("{}{}", endpoint_prefix, address))
803            .build();
804        let client = Client::from_conf(config);
805
806        Self {
807            client,
808            bucket: bucket.to_owned(),
809            metrics,
810            config: object_store_config,
811        }
812    }
813
814    /// Generates an HTTP GET request to download the object specified in `path`. If given,
815    /// `start_pos` and `end_pos` specify the first and last byte to download, respectively. Both
816    /// are inclusive and 0-based. For example, set `start_pos = 0` and `end_pos = 7` to download
817    /// the first 8 bytes. If neither is given, the request will download the whole object.
818    fn obj_store_request(
819        &self,
820        path: &str,
821        range: impl ObjectRangeBounds,
822    ) -> GetObjectFluentBuilder {
823        let req = self.client.get_object().bucket(&self.bucket).key(path);
824        if range.is_full() {
825            return req;
826        }
827
828        let start = range.start().map(|v| v.to_string()).unwrap_or_default();
829        let end = range.end().map(|v| (v - 1).to_string()).unwrap_or_default(); // included
830
831        req.range(format!("bytes={}-{}", start, end))
832    }
833
834    // When multipart upload is aborted, if any part uploads are in progress, those part uploads
835    // might or might not succeed. As a result, these parts will remain in the bucket and be
836    // charged for part storage. Therefore, we need to configure the bucket to purge stale
837    // parts.
838    //
839    /// Note: This configuration only works for S3. MinIO automatically enables this feature, and it
840    /// is not configurable with S3 sdk. To verify that this feature is enabled, use `mc admin
841    /// config get <alias> api`.
842    ///
843    /// Reference:
844    /// - S3
845    ///   - <https://docs.aws.amazon.com/AmazonS3/latest/API/API_AbortMultipartUpload.html>
846    ///   - <https://docs.aws.amazon.com/AmazonS3/latest/userguide/mpu-abort-incomplete-mpu-lifecycle-config.html>
847    /// - MinIO
848    ///   - <https://github.com/minio/minio/issues/15681#issuecomment-1245126561>
849    pub async fn configure_bucket_lifecycle(&self, data_directory: &str) -> bool {
850        // Check if lifecycle is already configured to avoid overriding existing configuration.
851        let bucket = self.bucket.as_str();
852        let mut configured_rules = vec![];
853        let get_config_result = self
854            .client
855            .get_bucket_lifecycle_configuration()
856            .bucket(bucket)
857            .send()
858            .await;
859        let mut is_expiration_configured = false;
860
861        if let Ok(config) = &get_config_result {
862            for rule in config.rules() {
863                if rule.expiration().is_some() {
864                    // When both of the conditions are met, it is considered that there is a risk of data deletion.
865                    //
866                    // 1. expiration status rule is enabled
867                    // 2. (a) prefix filter is not set
868                    // or (b) prefix filter is set to the data directory of RisingWave.
869                    //
870                    // P.S. 1 && (2a || 2b)
871                    is_expiration_configured |= rule.status == ExpirationStatus::Enabled // 1
872                    && match rule.filter().as_ref() {
873                        // 2a
874                        None => true,
875                        // 2b
876                        Some(filter) => {
877                            // `LifecycleRuleFilter` is a struct in newer aws-sdk-s3 versions.
878                            //
879                            // Treat an "empty filter" as applying to the entire bucket, which is
880                            // considered risky for RisingWave data deletion.
881                            let is_empty = filter.prefix().is_none()
882                                && filter.tag().is_none()
883                                && filter.and().is_none()
884                                && filter.object_size_greater_than().is_none()
885                                && filter.object_size_less_than().is_none();
886                            if is_empty {
887                                true
888                            } else {
889                                filter
890                                    .prefix()
891                                    .is_some_and(|prefix| data_directory.starts_with(prefix))
892                            }
893                        }
894                    };
895
896                    if matches!(rule.status(), ExpirationStatus::Enabled)
897                        && rule.abort_incomplete_multipart_upload().is_some()
898                    {
899                        configured_rules.push(rule);
900                    }
901                }
902            }
903        }
904
905        if !configured_rules.is_empty() {
906            tracing::info!(
907                "S3 bucket {} has already configured AbortIncompleteMultipartUpload: {:?}",
908                bucket,
909                configured_rules,
910            );
911        } else {
912            let bucket_lifecycle_rule = LifecycleRule::builder()
913                .id("abort-incomplete-multipart-upload")
914                .status(ExpirationStatus::Enabled)
915                // Empty prefix means "all objects".
916                .filter(LifecycleRuleFilter::builder().prefix("").build())
917                .abort_incomplete_multipart_upload(
918                    AbortIncompleteMultipartUpload::builder()
919                        .days_after_initiation(S3_INCOMPLETE_MULTIPART_UPLOAD_RETENTION_DAYS)
920                        .build(),
921                )
922                .build()
923                .unwrap();
924            let bucket_lifecycle_config = BucketLifecycleConfiguration::builder()
925                .rules(bucket_lifecycle_rule)
926                .build()
927                .unwrap();
928            if self
929                .client
930                .put_bucket_lifecycle_configuration()
931                .bucket(bucket)
932                .lifecycle_configuration(bucket_lifecycle_config)
933                .send()
934                .await
935                .is_ok()
936            {
937                tracing::info!(
938                    "S3 bucket {:?} is configured to automatically purge abandoned MultipartUploads after {} days",
939                    bucket,
940                    S3_INCOMPLETE_MULTIPART_UPLOAD_RETENTION_DAYS,
941                );
942            } else {
943                tracing::warn!(
944                    "Failed to configure life cycle rule for S3 bucket: {:?}. It is recommended to configure it manually to avoid unnecessary storage cost.",
945                    bucket
946                );
947            }
948        }
949        if is_expiration_configured {
950            tracing::info!(
951                "S3 bucket {} has already configured the expiration for the lifecycle.",
952                bucket,
953            );
954        }
955        is_expiration_configured
956    }
957}
958
959struct S3ObjectIter {
960    buffer: VecDeque<ObjectMetadata>,
961    client: Client,
962    bucket: String,
963    prefix: String,
964    next_continuation_token: Option<String>,
965    is_truncated: Option<bool>,
966    #[expect(clippy::type_complexity)]
967    send_future: Option<
968        BoxFuture<
969            'static,
970            Result<(Vec<ObjectMetadata>, Option<String>, Option<bool>), ObjectError>,
971        >,
972    >,
973
974    config: Arc<ObjectStoreConfig>,
975    start_after: Option<String>,
976}
977
978impl S3ObjectIter {
979    fn new(
980        client: Client,
981        bucket: String,
982        prefix: String,
983        config: Arc<ObjectStoreConfig>,
984        start_after: Option<String>,
985    ) -> Self {
986        Self {
987            buffer: VecDeque::default(),
988            client,
989            bucket,
990            prefix,
991            next_continuation_token: None,
992            is_truncated: Some(true),
993            send_future: None,
994            config,
995            start_after,
996        }
997    }
998}
999
1000impl Stream for S3ObjectIter {
1001    type Item = ObjectResult<ObjectMetadata>;
1002
1003    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
1004        if let Some(e) = self.buffer.pop_front() {
1005            return Poll::Ready(Some(Ok(e)));
1006        }
1007        if let Some(f) = self.send_future.as_mut() {
1008            return match ready!(f.poll_unpin(cx)) {
1009                Ok((more, next_continuation_token, is_truncated)) => {
1010                    self.next_continuation_token = next_continuation_token;
1011                    self.is_truncated = is_truncated;
1012                    self.buffer.extend(more);
1013                    self.send_future = None;
1014                    // only the first request may set start_after
1015                    self.start_after = None;
1016                    self.poll_next(cx)
1017                }
1018                Err(e) => {
1019                    self.send_future = None;
1020                    Poll::Ready(Some(Err(e)))
1021                }
1022            };
1023        }
1024        if !self.is_truncated.unwrap_or_default() {
1025            return Poll::Ready(None);
1026        }
1027        let mut request = self
1028            .client
1029            .list_objects_v2()
1030            .bucket(&self.bucket)
1031            .prefix(&self.prefix);
1032        #[cfg(not(madsim))]
1033        if let Some(start_after) = self.start_after.as_ref() {
1034            request = request.start_after(start_after);
1035        }
1036        if let Some(continuation_token) = self.next_continuation_token.as_ref() {
1037            request = request.continuation_token(continuation_token);
1038        }
1039        let config = self.config.clone();
1040        let f = async move {
1041            match request.send().await {
1042                Ok(r) => {
1043                    let more = r
1044                        .contents()
1045                        .iter()
1046                        .map(|obj| ObjectMetadata {
1047                            key: obj.key().expect("key required").to_owned(),
1048                            last_modified: obj
1049                                .last_modified()
1050                                .map(|l| l.as_secs_f64())
1051                                .unwrap_or(0f64),
1052                            total_size: obj.size().unwrap_or_default() as usize,
1053                        })
1054                        .collect_vec();
1055                    let is_truncated = r.is_truncated;
1056                    let next_continuation_token = r.next_continuation_token;
1057                    Ok((more, next_continuation_token, is_truncated))
1058                }
1059                Err(e) => Err(set_error_should_retry::<ListObjectsV2Error>(
1060                    config,
1061                    e.into(),
1062                )),
1063            }
1064        };
1065        self.send_future = Some(Box::pin(f));
1066        self.poll_next(cx)
1067    }
1068}
1069
1070fn set_error_should_retry<E>(config: Arc<ObjectStoreConfig>, object_err: ObjectError) -> ObjectError
1071where
1072    E: ProvideErrorMetadata + Into<BoxError> + Sync + Send + std::error::Error + 'static,
1073{
1074    let not_found = object_err.is_object_not_found_error();
1075
1076    if not_found {
1077        return object_err;
1078    }
1079
1080    let mut inner = object_err.into_inner();
1081    match inner.borrow_mut() {
1082        ObjectErrorInner::S3 {
1083            should_retry,
1084            inner,
1085        } => {
1086            let sdk_err = inner
1087                .as_ref()
1088                .downcast_ref::<SdkError<E, aws_smithy_runtime_api::http::Response<SdkBody>>>();
1089
1090            let err_should_retry = match sdk_err {
1091                Some(SdkError::DispatchFailure(e)) => {
1092                    if e.is_timeout() {
1093                        tracing::warn!(target: "http_timeout_retry", "{e:?} occurs, retry S3 get_object request.");
1094                        true
1095                    } else {
1096                        false
1097                    }
1098                }
1099
1100                Some(SdkError::ServiceError(e)) => match e.err().code() {
1101                    None => {
1102                        if config.s3.developer.retry_unknown_service_error
1103                            || config.s3.retry_unknown_service_error
1104                        {
1105                            tracing::warn!(target: "unknown_service_error", "{e:?} occurs, retry S3 get_object request.");
1106                            true
1107                        } else {
1108                            false
1109                        }
1110                    }
1111                    Some(code) => {
1112                        if config
1113                            .s3
1114                            .developer
1115                            .retryable_service_error_codes
1116                            .iter()
1117                            .any(|s| s.as_str().eq_ignore_ascii_case(code))
1118                        {
1119                            tracing::warn!(target: "retryable_service_error", "{e:?} occurs, retry S3 get_object request.");
1120                            true
1121                        } else {
1122                            false
1123                        }
1124                    }
1125                },
1126
1127                Some(SdkError::TimeoutError(_err)) => true,
1128
1129                _ => false,
1130            };
1131
1132            *should_retry = err_should_retry;
1133        }
1134
1135        _ => unreachable!(),
1136    }
1137
1138    ObjectError::from(inner)
1139}
1140
1141#[cfg(test)]
1142#[cfg(not(madsim))]
1143mod tests {
1144    use crate::object::prefix::s3::{NUM_BUCKET_PREFIXES, get_object_prefix};
1145
1146    fn get_hash_of_object(obj_id: u64) -> u32 {
1147        let crc_hash = crc32fast::hash(&obj_id.to_be_bytes());
1148        crc_hash % NUM_BUCKET_PREFIXES
1149    }
1150
1151    #[tokio::test]
1152    async fn test_get_object_prefix() {
1153        for obj_id in 0..99999 {
1154            let hash = get_hash_of_object(obj_id);
1155            let prefix = get_object_prefix(obj_id);
1156            assert_eq!(format!("{}/", hash), prefix);
1157        }
1158
1159        let obj_prefix = String::default();
1160        let path = format!("{}/{}{}.data", "hummock_001", obj_prefix, 101);
1161        assert_eq!("hummock_001/101.data", path);
1162    }
1163}