From 5551d9691323e4a8bcc6c074ac872028c0827e1d Mon Sep 17 00:00:00 2001
From: sunby <sunbingyi1992@gmail.com>
Date: Tue, 17 Oct 2023 18:28:26 +0800
Subject: [PATCH] [Cpp] Call FinalizeS3 at exit

Signed-off-by: sunby <sunbingyi1992@gmail.com>
---
 cpp/include/milvus-storage/storage/space.h |  2 --
 cpp/src/common/fs_util.cpp                 | 12 +++++++++++-
 cpp/src/storage/space.cpp                  |  9 ---------
 3 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/cpp/include/milvus-storage/storage/space.h b/cpp/include/milvus-storage/storage/space.h
index 0bfb7f1..c3d6ff3 100644
--- a/cpp/include/milvus-storage/storage/space.h
+++ b/cpp/include/milvus-storage/storage/space.h
@@ -42,8 +42,6 @@ class Space {
 
   std::vector<Blob> StatisticsBlobs();
 
-  ~Space();
-
   private:
   Status Init();
 
diff --git a/cpp/src/common/fs_util.cpp b/cpp/src/common/fs_util.cpp
index 61a914a..04bf06f 100644
--- a/cpp/src/common/fs_util.cpp
+++ b/cpp/src/common/fs_util.cpp
@@ -3,6 +3,8 @@
 #include <arrow/filesystem/hdfs.h>
 #include <arrow/filesystem/s3fs.h>
 #include <arrow/util/uri.h>
+#include <cstdlib>
+#include "common/log.h"
 #include "common/macro.h"
 namespace milvus_storage {
 
@@ -23,7 +25,15 @@ Result<std::shared_ptr<arrow::fs::FileSystem>> BuildFileSystem(const std::string
   // }
 
   if (schema == "s3") {
-    RETURN_ARROW_NOT_OK(arrow::fs::InitializeS3(arrow::fs::S3GlobalOptions{}));
+    if (!arrow::fs::IsS3Initialized()) {
+      RETURN_ARROW_NOT_OK(arrow::fs::EnsureS3Initialized());
+      std::atexit([]() {
+        auto status = arrow::fs::EnsureS3Finalized();
+        if (!status.ok()) {
+          LOG_STORAGE_WARNING_ << "Failed to finalize S3: " << status.message();
+        }
+      });
+    }
     ASSIGN_OR_RETURN_ARROW_NOT_OK(auto option, arrow::fs::S3Options::FromUri(uri_parser));
     ASSIGN_OR_RETURN_ARROW_NOT_OK(auto fs, arrow::fs::S3FileSystem::Make(option));
 
diff --git a/cpp/src/storage/space.cpp b/cpp/src/storage/space.cpp
index 7a8e45b..7368fce 100644
--- a/cpp/src/storage/space.cpp
+++ b/cpp/src/storage/space.cpp
@@ -309,13 +309,4 @@ Result<std::shared_ptr<arrow::RecordBatchReader>> Space::ScanData() {
   return RecordReader::MakeScanDataReader(manifest_, fs_);
 }
 
-Space::~Space() {
-  if (fs_->type_name() == "s3") {
-    auto status = arrow::fs::FinalizeS3();
-    if (!status.ok()) {
-      LOG_STORAGE_WARNING_ << "FinalizeS3 failed: " << status.message();
-    }
-  }
-}
-
 }  // namespace milvus_storage