expose explicit AWS SDK lifecycle for SigV4

Author: plusplusjiajia

src/iceberg/catalog/rest/auth/aws_sdk.h

@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#pragma once
+
+#include "iceberg/catalog/rest/iceberg_rest_export.h"
+#include "iceberg/result.h"
+
+/// \file iceberg/catalog/rest/auth/aws_sdk.h
+/// \brief Process-wide AWS SDK lifecycle for SigV4 authentication.
+///
+/// Applications using SigV4 should call InitializeAwsSdk() at startup and
+/// FinalizeAwsSdk() before exit. If never called, the SDK is lazily
+/// initialized on first SigV4 use and leaked at process exit. FinalizeAwsSdk()
+/// is intended for process-shutdown sequencing, not concurrent teardown.
+
+namespace iceberg::rest::auth {
+
+/// \brief Initialize the AWS SDK. Idempotent.
+ICEBERG_REST_EXPORT Status InitializeAwsSdk();
+
+/// \brief Shut down the AWS SDK. Refuses if any SigV4 sessions are alive.
+ICEBERG_REST_EXPORT Status FinalizeAwsSdk();
+
+ICEBERG_REST_EXPORT bool IsAwsSdkInitialized();
+ICEBERG_REST_EXPORT bool IsAwsSdkFinalized();
+
+}  // namespace iceberg::rest::auth

src/iceberg/catalog/rest/auth/sigv4_auth_manager.cc

@@ -18,10 +18,13 @@
  */
 
 #include "iceberg/catalog/rest/auth/auth_manager_internal.h"
+#include "iceberg/catalog/rest/auth/aws_sdk.h"
 #include "iceberg/catalog/rest/auth/sigv4_auth_manager_internal.h"
 
 #ifdef ICEBERG_SIGV4
 
+#  include <atomic>
+#  include <mutex>
 #  include <sstream>
 
 #  include <aws/core/Aws.h>
@@ -41,23 +44,17 @@ namespace iceberg::rest::auth {
 
 namespace {
 
-/// \brief Ensures the AWS SDK is initialized exactly once per process.
-///
-/// Aws::InitAPI / ShutdownAPI must bracket the process lifetime, which a
-/// library cannot enforce, so we never call ShutdownAPI (leak by design).
-class AwsSdkGuard {
- public:
-  static void EnsureInitialized() {
-    static AwsSdkGuard instance;
-    (void)instance;
-  }
+enum class LifecycleState : uint8_t { kUninitialized, kInitialized, kFinalized };
 
- private:
-  AwsSdkGuard() {
-    Aws::SDKOptions options;
-    Aws::InitAPI(options);
-  }
-};
+std::atomic<LifecycleState> g_state{LifecycleState::kUninitialized};
+std::mutex g_lifecycle_mutex;
+Aws::SDKOptions g_sdk_options;
+std::atomic<size_t> g_active_session_count{0};
+
+Status EnsureSdkInitialized() {
+  if (g_state.load() == LifecycleState::kInitialized) return {};
+  return InitializeAwsSdk();
+}
 
 Aws::Http::HttpMethod ToAwsMethod(HttpMethod method) {
   switch (method) {
@@ -114,9 +111,14 @@ SigV4AuthSession::SigV4AuthSession(
       signing_name_(std::move(signing_name)),
       credentials_provider_(std::move(credentials_provider)),
       signer_(std::make_unique<RestSigV4Signer>(
-          credentials_provider_, signing_name_.c_str(), signing_region_.c_str())) {}
+          credentials_provider_, signing_name_.c_str(), signing_region_.c_str())) {
+  // Counted so FinalizeAwsSdk() refuses to ShutdownAPI while sessions exist.
+  g_active_session_count.fetch_add(1, std::memory_order_relaxed);
+}
 
-SigV4AuthSession::~SigV4AuthSession() = default;
+SigV4AuthSession::~SigV4AuthSession() {
+  g_active_session_count.fetch_sub(1, std::memory_order_relaxed);
+}
 
 Result<HttpRequest> SigV4AuthSession::Authenticate(const HttpRequest& request) {
   ICEBERG_ASSIGN_OR_RAISE(auto delegate_request, delegate_->Authenticate(request));
@@ -189,7 +191,7 @@ SigV4AuthManager::~SigV4AuthManager() = default;
 Result<std::shared_ptr<AuthSession>> SigV4AuthManager::InitSession(
     HttpClient& init_client,
     const std::unordered_map<std::string, std::string>& properties) {
-  AwsSdkGuard::EnsureInitialized();
+  ICEBERG_RETURN_UNEXPECTED(EnsureSdkInitialized());
   ICEBERG_ASSIGN_OR_RAISE(auto delegate_session,
                           delegate_->InitSession(init_client, properties));
   return WrapSession(std::move(delegate_session), properties);
@@ -198,7 +200,7 @@ Result<std::shared_ptr<AuthSession>> SigV4AuthManager::InitSession(
 Result<std::shared_ptr<AuthSession>> SigV4AuthManager::CatalogSession(
     HttpClient& shared_client,
     const std::unordered_map<std::string, std::string>& properties) {
-  AwsSdkGuard::EnsureInitialized();
+  ICEBERG_RETURN_UNEXPECTED(EnsureSdkInitialized());
   catalog_properties_ = properties;
   ICEBERG_ASSIGN_OR_RAISE(auto delegate_session,
                           delegate_->CatalogSession(shared_client, properties));
@@ -323,6 +325,35 @@ Result<std::unique_ptr<AuthManager>> MakeSigV4AuthManager(
   return std::make_unique<SigV4AuthManager>(std::move(delegate));
 }
 
+Status InitializeAwsSdk() {
+  std::lock_guard<std::mutex> lock(g_lifecycle_mutex);
+  auto state = g_state.load();
+  if (state == LifecycleState::kInitialized) return {};
+  if (state == LifecycleState::kFinalized) {
+    return InvalidArgument("AWS SDK has already been finalized; cannot reinitialize");
+  }
+  Aws::InitAPI(g_sdk_options);
+  g_state.store(LifecycleState::kInitialized);
+  return {};
+}
+
+Status FinalizeAwsSdk() {
+  std::lock_guard<std::mutex> lock(g_lifecycle_mutex);
+  if (g_state.load() != LifecycleState::kInitialized) return {};
+  auto live = g_active_session_count.load();
+  if (live != 0) {
+    return Invalid(
+        "Cannot finalize AWS SDK while {} SigV4 auth session(s) are still alive", live);
+  }
+  Aws::ShutdownAPI(g_sdk_options);
+  g_state.store(LifecycleState::kFinalized);
+  return {};
+}
+
+bool IsAwsSdkInitialized() { return g_state.load() == LifecycleState::kInitialized; }
+
+bool IsAwsSdkFinalized() { return g_state.load() == LifecycleState::kFinalized; }
+
 }  // namespace iceberg::rest::auth
 
 #else  // !ICEBERG_SIGV4
@@ -336,6 +367,17 @@ Result<std::unique_ptr<AuthManager>> MakeSigV4AuthManager(
       "SigV4 authentication is not built; configure with -DICEBERG_SIGV4=ON");
 }
 
+Status InitializeAwsSdk() {
+  return NotSupported(
+      "SigV4 authentication is not built; configure with -DICEBERG_SIGV4=ON");
+}
+
+Status FinalizeAwsSdk() { return {}; }
+
+bool IsAwsSdkInitialized() { return false; }
+
+bool IsAwsSdkFinalized() { return false; }
+
 }  // namespace iceberg::rest::auth
 
 #endif  // ICEBERG_SIGV4

src/iceberg/catalog/rest/meson.build

@@ -96,6 +96,7 @@ iceberg_rest_auth_headers = [
     'auth/auth_managers.h',
     'auth/auth_properties.h',
     'auth/auth_session.h',
+    'auth/aws_sdk.h',
     'auth/oauth2_util.h',
 ]
 install_headers(iceberg_rest_auth_headers, subdir: 'iceberg/catalog/rest/auth')

src/iceberg/test/sigv4_auth_test.cc

@@ -22,13 +22,13 @@
 #  include <string>
 #  include <unordered_map>
 
-#  include <aws/core/Aws.h>
 #  include <aws/core/auth/AWSCredentialsProvider.h>
 #  include <gtest/gtest.h>
 
 #  include "iceberg/catalog/rest/auth/auth_managers.h"
 #  include "iceberg/catalog/rest/auth/auth_properties.h"
 #  include "iceberg/catalog/rest/auth/auth_session.h"
+#  include "iceberg/catalog/rest/auth/aws_sdk.h"
 #  include "iceberg/catalog/rest/auth/sigv4_auth_manager_internal.h"
 #  include "iceberg/catalog/rest/http_client.h"
 #  include "iceberg/table_identifier.h"
@@ -38,14 +38,7 @@ namespace iceberg::rest::auth {
 
 class SigV4AuthTest : public ::testing::Test {
  protected:
-  static void SetUpTestSuite() {
-    static bool initialized = false;
-    if (!initialized) {
-      Aws::SDKOptions options;
-      Aws::InitAPI(options);
-      initialized = true;
-    }
-  }
+  static void SetUpTestSuite() { ASSERT_THAT(InitializeAwsSdk(), IsOk()); }
 
   HttpClient client_{{}};
 
@@ -61,6 +54,23 @@ class SigV4AuthTest : public ::testing::Test {
   }
 };
 
+TEST_F(SigV4AuthTest, LifecycleInitializeIsIdempotent) {
+  EXPECT_THAT(InitializeAwsSdk(), IsOk());
+  EXPECT_TRUE(IsAwsSdkInitialized());
+  EXPECT_FALSE(IsAwsSdkFinalized());
+}
+
+TEST_F(SigV4AuthTest, LifecycleFinalizeRefusesWhileSessionsAlive) {
+  auto properties = MakeSigV4Properties();
+  auto manager_result = AuthManagers::Load("test-catalog", properties);
+  ASSERT_THAT(manager_result, IsOk());
+  auto session_result = manager_result.value()->CatalogSession(client_, properties);
+  ASSERT_THAT(session_result, IsOk());
+
+  EXPECT_THAT(FinalizeAwsSdk(), IsError(ErrorKind::kInvalid));
+  EXPECT_TRUE(IsAwsSdkInitialized());
+}
+
 TEST_F(SigV4AuthTest, LoadSigV4AuthManager) {
   auto properties = MakeSigV4Properties();
   auto manager_result = AuthManagers::Load("test-catalog", properties);