From 37dce8a66054157c1b5f8574f9bef47c8e1c6e4e Mon Sep 17 00:00:00 2001
From: bigsheeper <yihao.dai@zilliz.com>
Date: Mon, 31 Oct 2022 17:09:33 +0800
Subject: [PATCH] Make s3 client creation thread safe, prevent crashing
 (#20200)

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
---
 .../core/src/storage/MinioChunkManager.cpp    | 25 ++++++++++++++-----
 internal/core/src/storage/MinioChunkManager.h |  5 ++++
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/internal/core/src/storage/MinioChunkManager.cpp b/internal/core/src/storage/MinioChunkManager.cpp
index 467bab05b4..27176f2721 100644
--- a/internal/core/src/storage/MinioChunkManager.cpp
+++ b/internal/core/src/storage/MinioChunkManager.cpp
@@ -43,6 +43,7 @@
 namespace milvus::storage {
 
 std::atomic<size_t> MinioChunkManager::init_count_(0);
+std::mutex MinioChunkManager::client_mutex_;
 
 /**
  * @brief convert std::string to Aws::String
@@ -67,12 +68,27 @@ ConvertFromAwsString(const Aws::String& aws_str) {
     return std::string(aws_str.c_str(), aws_str.size());
 }
 
-MinioChunkManager::MinioChunkManager(const StorageConfig& storage_config)
-    : default_bucket_name_(storage_config.bucket_name) {
+void
+MinioChunkManager::InitSDKAPI() {
+    std::scoped_lock lock{client_mutex_};
     const size_t initCount = init_count_++;
     if (initCount == 0) {
         Aws::InitAPI(sdk_options_);
     }
+}
+
+void
+MinioChunkManager::ShutdownSDKAPI() {
+    std::scoped_lock lock{client_mutex_};
+    const size_t initCount = --init_count_;
+    if (initCount == 0) {
+        Aws::ShutdownAPI(sdk_options_);
+    }
+}
+
+MinioChunkManager::MinioChunkManager(const StorageConfig& storage_config)
+    : default_bucket_name_(storage_config.bucket_name) {
+    InitSDKAPI();
     Aws::Client::ClientConfiguration config;
     config.endpointOverride = ConvertToAwsString(storage_config.address);
 
@@ -117,10 +133,7 @@ MinioChunkManager::MinioChunkManager(const StorageConfig& storage_config)
 }
 
 MinioChunkManager::~MinioChunkManager() {
-    const size_t initCount = --init_count_;
-    if (initCount == 0) {
-        Aws::ShutdownAPI(sdk_options_);
-    }
+    ShutdownSDKAPI();
     client_.reset();
 }
 
diff --git a/internal/core/src/storage/MinioChunkManager.h b/internal/core/src/storage/MinioChunkManager.h
index e1261d9349..580b1bba76 100644
--- a/internal/core/src/storage/MinioChunkManager.h
+++ b/internal/core/src/storage/MinioChunkManager.h
@@ -112,10 +112,15 @@ class MinioChunkManager : public RemoteChunkManager {
     GetObjectBuffer(const std::string& bucket_name, const std::string& object_name, void* buf, uint64_t size);
     std::vector<std::string>
     ListObjects(const char* bucket_name, const char* prefix = NULL);
+    void
+    InitSDKAPI();
+    void
+    ShutdownSDKAPI();
 
  private:
     const Aws::SDKOptions sdk_options_;
     static std::atomic<size_t> init_count_;
+    static std::mutex client_mutex_;
     std::shared_ptr<Aws::S3::S3Client> client_;
     std::string default_bucket_name_;
 };