From b24df3eaa9426633ff4e3f7d44602914f4db2677 Mon Sep 17 00:00:00 2001
From: Lei Zhang <zhanglei@selectdb.com>
Date: Mon, 11 Nov 2024 11:27:36 +0800
Subject: [PATCH] [fix](cloud) Fix aws sdk default retry policy not taking
 effect

* AWS SDK default retry policy check response xml `code` to determine
  retry or not, but other object storage such as `OSS`, `COS` not compatible
  with AWS S3, so adding a custome retry policy for solving the problem
---
 be/src/util/s3_util.cpp           |  5 +++--
 common/cpp/obj_retry_strategy.cpp | 11 ++++++++---
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/be/src/util/s3_util.cpp b/be/src/util/s3_util.cpp
index 18058469ee4566..f87389b6b3e119 100644
--- a/be/src/util/s3_util.cpp
+++ b/be/src/util/s3_util.cpp
@@ -42,6 +42,7 @@
 #include "common/config.h"
 #include "common/logging.h"
 #include "common/status.h"
+#include "cpp/obj_retry_strategy.h"
 #include "cpp/sync_point.h"
 #ifdef USE_AZURE
 #include "io/fs/azure_obj_storage_client.h"
@@ -307,8 +308,8 @@ std::shared_ptr<io::ObjStorageClient> S3ClientFactory::_create_s3_client(
         aws_config.scheme = Aws::Http::Scheme::HTTP;
     }
 
-    aws_config.retryStrategy =
-            std::make_shared<Aws::Client::DefaultRetryStrategy>(config::max_s3_client_retry);
+    aws_config.retryStrategy = std::make_shared<S3CustomRetryStrategy>(
+            config::max_s3_client_retry /*scaleFactor = 25*/);
     std::shared_ptr<Aws::S3::S3Client> new_client;
     if (!s3_conf.ak.empty() && !s3_conf.sk.empty()) {
         Aws::Auth::AWSCredentials aws_cred(s3_conf.ak, s3_conf.sk);
diff --git a/common/cpp/obj_retry_strategy.cpp b/common/cpp/obj_retry_strategy.cpp
index 4daa31dc588493..0226d1af28f038 100644
--- a/common/cpp/obj_retry_strategy.cpp
+++ b/common/cpp/obj_retry_strategy.cpp
@@ -17,6 +17,7 @@
 
 #include "obj_retry_strategy.h"
 
+#include <aws/core/http/HttpResponse.h>
 #include <bvar/reducer.h>
 
 namespace doris {
@@ -29,12 +30,16 @@ S3CustomRetryStrategy::~S3CustomRetryStrategy() = default;
 
 bool S3CustomRetryStrategy::ShouldRetry(const Aws::Client::AWSError<Aws::Client::CoreErrors>& error,
                                         long attemptedRetries) const {
-    if (attemptedRetries < m_maxRetries &&
-        error.GetResponseCode() == Aws::Http::HttpResponseCode::TOO_MANY_REQUESTS) {
+    if (attemptedRetries >= m_maxRetries) {
+        return false;
+    }
+
+    if (Aws::Http::IsRetryableHttpResponseCode(error.GetResponseCode()) || error.ShouldRetry()) {
         s3_too_many_request_retry_cnt << 1;
         return true;
     }
-    return Aws::Client::DefaultRetryStrategy::ShouldRetry(error, attemptedRetries);
+
+    return false;
 }
 #ifdef USE_AZURE
 AzureRetryRecordPolicy::AzureRetryRecordPolicy(int retry_cnt) : retry_cnt(retry_cnt) {}