From cbf2d4212a331ca0a9dc878ed961e8d305f72a89 Mon Sep 17 00:00:00 2001
From: Bingyi Sun <sunbingyi1992@gmail.com>
Date: Tue, 17 Oct 2023 18:07:58 +0800
Subject: [PATCH] [Cpp] Finalize s3 before the end (#76)

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

diff --git a/cpp/include/milvus-storage/storage/space.h b/cpp/include/milvus-storage/storage/space.h
index c3d6ff3..0bfb7f1 100644
--- a/cpp/include/milvus-storage/storage/space.h
+++ b/cpp/include/milvus-storage/storage/space.h
@@ -42,6 +42,8 @@ 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 9cc0517..61a914a 100644
--- a/cpp/src/common/fs_util.cpp
+++ b/cpp/src/common/fs_util.cpp
@@ -23,10 +23,10 @@ Result<std::shared_ptr<arrow::fs::FileSystem>> BuildFileSystem(const std::string
   // }
 
   if (schema == "s3") {
+    RETURN_ARROW_NOT_OK(arrow::fs::InitializeS3(arrow::fs::S3GlobalOptions{}));
     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));
 
-    RETURN_ARROW_NOT_OK(arrow::fs::InitializeS3(arrow::fs::S3GlobalOptions{}));
     return std::shared_ptr<arrow::fs::FileSystem>(fs);
   }
 
diff --git a/cpp/src/storage/space.cpp b/cpp/src/storage/space.cpp
index 409bc1d..7a8e45b 100644
--- a/cpp/src/storage/space.cpp
+++ b/cpp/src/storage/space.cpp
@@ -1,5 +1,6 @@
 
 #include <arrow/filesystem/filesystem.h>
+#include <arrow/filesystem/s3fs.h>
 #include <arrow/filesystem/type_fwd.h>
 #include <arrow/status.h>
 #include <algorithm>
@@ -12,6 +13,7 @@
 
 #include "arrow/array/builder_primitive.h"
 #include "common/fs_util.h"
+#include "common/log.h"
 #include "common/macro.h"
 #include "file/delete_fragment.h"
 #include "filter/constant_filter.h"
@@ -307,4 +309,13 @@ 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