From dcddcc3a41e4c0b1ce67d69fa7a8d981c1f691f0 Mon Sep 17 00:00:00 2001
From: zhaohehuhu <luoyedeyi@163.com>
Date: Wed, 23 Oct 2024 13:35:10 +0800
Subject: [PATCH] support retry

---
 .../org/apache/celeborn/S3MultipartUploadHandler.java     | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/multipart-uploader/src/main/java/org/apache/celeborn/S3MultipartUploadHandler.java b/multipart-uploader/src/main/java/org/apache/celeborn/S3MultipartUploadHandler.java
index 96170d2cdc3..82bc12ee604 100644
--- a/multipart-uploader/src/main/java/org/apache/celeborn/S3MultipartUploadHandler.java
+++ b/multipart-uploader/src/main/java/org/apache/celeborn/S3MultipartUploadHandler.java
@@ -23,8 +23,10 @@
 import java.util.Collections;
 import java.util.List;
 
+import com.amazonaws.ClientConfiguration;
 import com.amazonaws.auth.AWSStaticCredentialsProvider;
 import com.amazonaws.auth.BasicAWSCredentials;
+import com.amazonaws.retry.PredefinedRetryPolicies;
 import com.amazonaws.services.s3.AmazonS3;
 import com.amazonaws.services.s3.AmazonS3ClientBuilder;
 import com.amazonaws.services.s3.model.AbortMultipartUploadRequest;
@@ -58,10 +60,16 @@ public S3MultipartUploadHandler(AWSCredentials awsCredentials, String key) {
     this.awsCredentials = awsCredentials;
     BasicAWSCredentials basicAWSCredentials =
         new BasicAWSCredentials(awsCredentials.getS3AccessKey(), awsCredentials.getS3SecretKey());
+    ClientConfiguration clientConfig = new ClientConfiguration()
+            .withRetryPolicy(PredefinedRetryPolicies.getDefaultRetryPolicyWithCustomMaxRetries(5))
+            .withMaxErrorRetry(5)
+            .withConnectionTimeout(10000)
+            .withRequestTimeout(50000);
     this.s3Client =
         AmazonS3ClientBuilder.standard()
             .withCredentials(new AWSStaticCredentialsProvider(basicAWSCredentials))
             .withRegion(awsCredentials.getS3EndpointRegion())
+            .withClientConfiguration(clientConfig)
             .build();
     this.key = key;
   }