Add endpoint cache

Author: pitrou

cpp/src/arrow/filesystem/s3fs.cc

@@ -914,21 +914,21 @@ Result<std::shared_ptr<S3ClientHolder>> GetClientHolder(
 // S3 client factory: build S3Client from S3Options
 
 struct EndpointConfigKey {
-  explicit EndpointConfigKey(const S3Options& options)
-      : region(options.region),
-        scheme(options.scheme),
-        endpoint_override(options.endpoint_override),
-        force_virtual_addressing(options.force_virtual_addressing) {}
-
-  std::string region;
-  std::string scheme;
-  std::string endpoint_override;
-  bool force_virtual_addressing;
-
-  bool operator==(const EndpointConfigKey& other) {
+  explicit EndpointConfigKey(const Aws::S3::S3ClientConfiguration& config)
+      : region(config.region),
+        scheme(config.scheme),
+        endpoint_override(config.endpointOverride),
+        use_virtual_addressing(config.useVirtualAddressing) {}
+
+  Aws::String region;
+  Aws::Http::Scheme scheme;
+  Aws::String endpoint_override;
+  bool use_virtual_addressing;
+
+  bool operator==(const EndpointConfigKey& other) const noexcept {
     return region == other.region && scheme == other.scheme &&
            endpoint_override == other.endpoint_override &&
-           force_virtual_addressing == other.force_virtual_addressing;
+           use_virtual_addressing == other.use_virtual_addressing;
   }
 };
 
@@ -938,14 +938,59 @@ struct EndpointConfigKey {
 template <>
 struct std::hash<arrow::fs::EndpointConfigKey> {
   std::size_t operator()(const arrow::fs::EndpointConfigKey& key) const noexcept {
-    auto h = std::hash<std::string>{};
-    return h(key.region) ^ h(key.scheme) ^ h(key.endpoint_override);
+    auto h = std::hash<Aws::String>{};
+    return h(key.region) ^ h(key.endpoint_override);
   }
 };
 
 namespace arrow::fs {
 namespace {
 
+class EndpointProviderBuilder {
+ public:
+  struct LockedEndpointProvider {
+    std::unique_lock<std::mutex> lock;
+    std::shared_ptr<Aws::S3::S3EndpointProvider> endpoint_provider;
+  };
+
+  LockedEndpointProvider Lookup(
+      const Aws::S3::S3ClientConfiguration& config) {
+    auto key = EndpointConfigKey(config);
+    CacheValue* value;
+    {
+      std::unique_lock lock(cache_mutex_);
+      value = &cache_[std::move(key)];
+    }
+    std::unique_lock lock(value->mutex);
+    if (!value->endpoint_provider) {
+      value->endpoint_provider = std::make_shared<Aws::S3::S3EndpointProvider>();
+      value->endpoint_provider->InitBuiltInParameters(config);
+    }
+    return {std::move(lock), value->endpoint_provider};
+  }
+
+  void Reset() {
+    std::unique_lock lock(cache_mutex_);
+    cache_.clear();
+  }
+
+  static EndpointProviderBuilder* Instance() {
+    static EndpointProviderBuilder instance;
+    return &instance;
+  }
+
+ protected:
+  EndpointProviderBuilder() = default;
+
+  struct CacheValue {
+    std::mutex mutex;
+    std::shared_ptr<Aws::S3::S3EndpointProvider> endpoint_provider;
+  };
+
+  std::mutex cache_mutex_;
+  std::unordered_map<EndpointConfigKey, CacheValue> cache_;
+};
+
 class ClientBuilder {
  public:
   explicit ClientBuilder(S3Options options) : options_(std::move(options)) {}
@@ -1022,12 +1067,13 @@ class ClientBuilder {
       client_config_.maxConnections = std::max(io_context->executor()->GetCapacity(), 25);
     }
 
-    auto endpoint_provider = std::make_shared<Aws::S3::S3EndpointProvider>();
-    endpoint_provider->InitBuiltInParameters(client_config_);
-
+    // The EndPointProvider is locked because S3Client updates EndPointProvider
+    // configuration in a non-thread-safe way, even though the updates are
+    // idempotent.
+    auto locked = EndpointProviderBuilder::Instance()->Lookup(client_config_);
     auto client =
-        std::make_shared<S3Client>(credentials_provider_, nullptr, client_config_);
-    client->accessEndpointProvider() = endpoint_provider;
+        std::make_shared<S3Client>(credentials_provider_, locked.endpoint_provider,
+                                   client_config_);
 
     client->s3_retry_strategy_ = options_.retry_strategy;
     return GetClientHolder(std::move(client));
@@ -2962,6 +3008,7 @@ struct AwsInstance {
                "This could lead to a segmentation fault at exit";
       }
       GetClientFinalizer()->Finalize();
+      EndpointProviderBuilder::Instance()->Reset();
       Aws::ShutdownAPI(aws_options_);
     }
   }