Ensure aws library initialization is called by GetAwsAllocator()

Also rework aws global variables into a singleton AwsApi class, and update
the lazy initialization of the class members.

PiperOrigin-RevId: 749910414
Change-Id: Ia3a0f5cd658cc1a69a59b5a034a0918b6abd4437

Author: laramiel

tensorstore/internal/aws/BUILD

@@ -20,11 +20,12 @@ tensorstore_cc_library(
         "@aws_c_common",
         "@aws_c_http",
         "@aws_c_io",
-        "@com_google_absl//absl/base",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/base:log_severity",
+        "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/debugging:leak_check",
         "@com_google_absl//absl/log:absl_log",
+        "@com_google_absl//absl/synchronization",
     ],
 )
 

tensorstore/internal/aws/aws_api.cc

@@ -23,10 +23,12 @@
 #include <string_view>
 
 #include "absl/base/attributes.h"
-#include "absl/base/call_once.h"
 #include "absl/base/log_severity.h"
+#include "absl/base/no_destructor.h"
+#include "absl/base/thread_annotations.h"
 #include "absl/debugging/leak_check.h"
 #include "absl/log/absl_log.h"
+#include "absl/synchronization/mutex.h"
 #include <aws/auth/auth.h>
 #include <aws/cal/cal.h>
 #include <aws/common/allocator.h>
@@ -145,74 +147,104 @@ aws_logger s_absl_logger{
 };
 
 // AWS apis rely on global initialization; do that here.
-ABSL_CONST_INIT absl::once_flag g_init;
-
-aws_event_loop_group *g_event_loop_group = nullptr;
-aws_host_resolver *g_resolver = nullptr;
-aws_client_bootstrap *g_client_bootstrap = nullptr;
-aws_tls_ctx *g_tls_ctx = nullptr;
-
-void InitAwsLibraries() {
-  absl::LeakCheckDisabler disabler;
-
-  auto *allocator = GetAwsAllocator();
-
-  /* Initialize AWS libraries.*/
-  aws_common_library_init(allocator);
-
-  s_absl_logger.allocator = allocator;
-  aws_logger_set(&s_absl_logger);
-
-  aws_cal_library_init(allocator);
-  aws_io_library_init(allocator);
-  aws_http_library_init(allocator);
-  aws_auth_library_init(allocator);
-
-  /* event loop */
-  g_event_loop_group = aws_event_loop_group_new_default(allocator, 0, nullptr);
-
-  /* resolver */
-  aws_host_resolver_default_options resolver_options;
-  AWS_ZERO_STRUCT(resolver_options);
-  resolver_options.el_group = g_event_loop_group;
-  resolver_options.max_entries = 32;  // defaults to 8?
-  g_resolver = aws_host_resolver_new_default(allocator, &resolver_options);
-
-  /* client bootstrap */
-  aws_client_bootstrap_options bootstrap_options;
-  AWS_ZERO_STRUCT(bootstrap_options);
-  bootstrap_options.event_loop_group = g_event_loop_group;
-  bootstrap_options.host_resolver = g_resolver;
-  g_client_bootstrap = aws_client_bootstrap_new(allocator, &bootstrap_options);
-  if (g_client_bootstrap == nullptr) {
-    ABSL_LOG(FATAL) << "ERROR initializing client bootstrap: "
-                    << aws_error_debug_str(aws_last_error());
+class AwsApi {
+ public:
+  AwsApi() : allocator_(aws_default_allocator()) {
+    absl::LeakCheckDisabler disabler;
+
+    /* Initialize AWS libraries.*/
+    aws_common_library_init(allocator_);
+
+    s_absl_logger.allocator = allocator_;
+    aws_logger_set(&s_absl_logger);
+
+    aws_cal_library_init(allocator_);
+    aws_io_library_init(allocator_);
+    aws_http_library_init(allocator_);
+    aws_auth_library_init(allocator_);
+  }
+
+  aws_allocator *allocator() { return allocator_; }
+
+  aws_client_bootstrap *client_bootstrap() ABSL_LOCKS_EXCLUDED(mutex_) {
+    absl::MutexLock l(&mutex_);
+    init_client_bootstrap();
+    return client_bootstrap_;
+  }
+
+  aws_tls_ctx *tls_ctx() ABSL_LOCKS_EXCLUDED(mutex_) {
+    absl::MutexLock l(&mutex_);
+    init_tls_ctx();
+    return tls_ctx_;
+  }
+
+ private:
+  void init_event_loop_group() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) {
+    if (event_loop_group_ != nullptr) return;
+    event_loop_group_ =
+        aws_event_loop_group_new_default(allocator_, 0, nullptr);
   }
 
-  AwsTlsCtx tls_ctx = AwsTlsCtxBuilder(allocator).Build();
-  if (tls_ctx == nullptr) {
-    ABSL_LOG(FATAL) << "ERROR initializing TLS context: "
-                    << aws_error_debug_str(aws_last_error());
+  void init_resolver() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) {
+    if (resolver_ != nullptr) return;
+    init_event_loop_group();
+
+    aws_host_resolver_default_options resolver_options;
+    AWS_ZERO_STRUCT(resolver_options);
+    resolver_options.el_group = event_loop_group_;
+    resolver_options.max_entries = 32;  // defaults to 8?
+    resolver_ = aws_host_resolver_new_default(allocator_, &resolver_options);
+  }
+
+  void init_client_bootstrap() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) {
+    if (client_bootstrap_ != nullptr) return;
+    init_event_loop_group();
+    init_resolver();
+
+    aws_client_bootstrap_options bootstrap_options;
+    AWS_ZERO_STRUCT(bootstrap_options);
+    bootstrap_options.event_loop_group = event_loop_group_;
+    bootstrap_options.host_resolver = resolver_;
+    client_bootstrap_ =
+        aws_client_bootstrap_new(allocator_, &bootstrap_options);
+    if (client_bootstrap_ == nullptr) {
+      ABSL_LOG(FATAL) << "ERROR initializing client bootstrap: "
+                      << aws_error_debug_str(aws_last_error());
+    }
+  }
+
+  void init_tls_ctx() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) {
+    if (tls_ctx_ != nullptr) return;
+    auto my_tls_ctx = AwsTlsCtxBuilder(allocator_).Build();
+    if (my_tls_ctx == nullptr) {
+      ABSL_LOG(FATAL) << "ERROR initializing TLS context: "
+                      << aws_error_debug_str(aws_last_error());
+    }
+    tls_ctx_ = my_tls_ctx.release();
   }
-  g_tls_ctx = tls_ctx.release();
+
+  absl::Mutex mutex_;
+  aws_allocator *allocator_ = nullptr;
+  aws_event_loop_group *event_loop_group_ ABSL_GUARDED_BY(mutex_) = nullptr;
+  aws_host_resolver *resolver_ ABSL_GUARDED_BY(mutex_) = nullptr;
+  aws_client_bootstrap *client_bootstrap_ ABSL_GUARDED_BY(mutex_) = nullptr;
+  aws_tls_ctx *tls_ctx_ ABSL_GUARDED_BY(mutex_) = nullptr;
+};
+
+AwsApi &GetAwsApi() {
+  static absl::NoDestructor<AwsApi> aws_api;
+  return *aws_api;
 }
 
 }  // namespace
 
-aws_allocator *GetAwsAllocator() {
-  // The default allocator is used for all AWS API objects.
-  return aws_default_allocator();
-}
+aws_allocator *GetAwsAllocator() { return GetAwsApi().allocator(); }
 
 aws_client_bootstrap *GetAwsClientBootstrap() {
-  absl::call_once(g_init, InitAwsLibraries);
-  return g_client_bootstrap;
+  return GetAwsApi().client_bootstrap();
 }
 
-aws_tls_ctx *GetAwsTlsCtx() {
-  absl::call_once(g_init, InitAwsLibraries);
-  return g_tls_ctx;
-}
+aws_tls_ctx *GetAwsTlsCtx() { return GetAwsApi().tls_ctx(); }
 
 }  // namespace internal_aws
 }  // namespace tensorstore

tensorstore/kvstore/s3/BUILD

@@ -30,6 +30,7 @@ tensorstore_cc_library(
         "//tensorstore/internal:intrusive_ptr",
         "//tensorstore/internal:source_location",
         "//tensorstore/internal:uri_utils",
+        "//tensorstore/internal/aws:aws_api",
         "//tensorstore/internal/aws:aws_credentials",
         "//tensorstore/internal/cache_key",
         "//tensorstore/internal/digest:sha256",
@@ -380,6 +381,7 @@ tensorstore_cc_library(
         "//tensorstore:json_serialization_options_base",
         "//tensorstore/internal:type_traits",
         "//tensorstore/internal:uri_utils",
+        "//tensorstore/internal/aws:aws_api",
         "//tensorstore/internal/aws:aws_credentials",
         "//tensorstore/internal/aws/credentials:common",
         "//tensorstore/internal/json_binding",

tensorstore/kvstore/s3/aws_credentials_spec.cc

@@ -26,6 +26,7 @@
 #include "absl/strings/str_cat.h"
 #include <aws/common/error.h>
 #include <nlohmann/json_fwd.hpp>
+#include "tensorstore/internal/aws/aws_api.h"
 #include "tensorstore/internal/aws/aws_credentials.h"
 #include "tensorstore/internal/aws/credentials/common.h"
 #include "tensorstore/internal/json_binding/bindable.h"
@@ -166,6 +167,7 @@ Result<AwsCredentialsProvider> MakeAwsCredentialsProvider(const Spec& spec) {
     }
   };
 
+  (void)internal_aws::GetAwsAllocator();  // Force AWS library initialization.
   AwsCredentialsProvider credentials_provider =
       std::visit(MakeCredentialsVisitor{}, spec.config);
 

tensorstore/kvstore/s3/s3_key_value_store.cc

@@ -37,6 +37,7 @@
 #include "absl/time/clock.h"
 #include "absl/time/time.h"
 #include "tensorstore/context.h"
+#include "tensorstore/internal/aws/aws_api.h"
 #include "tensorstore/internal/aws/aws_credentials.h"
 #include "tensorstore/internal/data_copy_concurrency_resource.h"
 #include "tensorstore/internal/digest/sha256.h"
@@ -1342,6 +1343,7 @@ Future<const S3EndpointRegion> S3KeyValueStore::MaybeResolveRegion() {
 }
 
 Future<kvstore::DriverPtr> S3KeyValueStoreSpec::DoOpen() const {
+  (void)internal_aws::GetAwsAllocator();  // Force AWS library initialization.
   // TODO: The transport should support the AWS_CA_BUNDLE environment
   // variable.
   TENSORSTORE_ASSIGN_OR_RETURN(