refactor(io): hand out S3 delegates by shared_ptr and build them off the lock

Author: plusplusjiajia

src/iceberg/arrow/s3/arrow_s3_file_io.cc

@@ -17,9 +17,12 @@
  * under the License.
  */
 
+#include <algorithm>
 #include <cstdlib>
 #include <memory>
+#include <mutex>
 #include <optional>
+#include <shared_mutex>
 #include <string>
 #include <string_view>
 #include <unordered_map>
@@ -179,7 +182,7 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials {
  public:
   ArrowS3FileIO(std::shared_ptr<::arrow::fs::FileSystem> arrow_fs,
                 std::unordered_map<std::string, std::string> default_properties)
-      : default_file_io_(std::move(arrow_fs)),
+      : default_file_io_(std::make_shared<ArrowFileSystemFileIO>(std::move(arrow_fs))),
         default_properties_(std::move(default_properties)) {}
 
   Result<std::unique_ptr<InputFile>> NewInputFile(std::string file_location) override;
@@ -196,27 +199,67 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials {
   Status SetStorageCredentials(
       const std::vector<StorageCredential>& storage_credentials) override;
 
-  const std::vector<StorageCredential>& credentials() const override {
+  std::vector<StorageCredential> credentials() const override {
+    std::shared_lock lock(mutex_);
     return storage_credentials_;
   }
 
   SupportsStorageCredentials* AsSupportsStorageCredentials() override { return this; }
 
  private:
-  ArrowFileSystemFileIO& FileIOForPath(std::string_view location);
-
-  ArrowFileSystemFileIO default_file_io_;
+  /// \brief Delegate serving `location`, pinned by the caller against a
+  /// concurrent credential install.
+  std::shared_ptr<ArrowFileSystemFileIO> FileIOForPath(std::string_view location);
+
+  using DelegatesByPrefix =
+      std::vector<std::pair<std::string, std::shared_ptr<ArrowFileSystemFileIO>>>;
+
+  /// \brief Longest-prefix match against one consistent view of the delegates.
+  static std::shared_ptr<ArrowFileSystemFileIO> MatchDelegate(
+      const std::shared_ptr<ArrowFileSystemFileIO>& fallback,
+      const DelegatesByPrefix& by_prefix, std::string_view location);
+
+  /// \brief Build a delegate for each credential this FileIO can serve.
+  ///
+  /// Lock-free on purpose: building an S3 client can reach out to discover a
+  /// bucket region, which would stall every concurrent operation. Reads no
+  /// mutable member state.
+  Result<DelegatesByPrefix> BuildDelegates(
+      const std::vector<StorageCredential>& storage_credentials) const;
+
+  /// \brief Swap in credentials and delegates, handing back the retired ones.
+  ///
+  /// Callers must hold `mutex_` exclusively and let the returned generation
+  /// destruct only after releasing it: tearing down an S3 client can block on
+  /// in-flight requests, which would stall every operation.
+  void InstallCredentials(std::vector<StorageCredential>& storage_credentials,
+                          DelegatesByPrefix& delegates);
+
+  std::shared_ptr<ArrowFileSystemFileIO> default_file_io_;
   std::unordered_map<std::string, std::string> default_properties_;
+  // Guards everything below; shared because reads happen per file operation.
+  mutable std::shared_mutex mutex_;
   std::vector<StorageCredential> storage_credentials_;
-  std::vector<std::pair<std::string, std::unique_ptr<ArrowFileSystemFileIO>>>
-      file_io_by_prefix_;
+  DelegatesByPrefix file_io_by_prefix_;
 };
 
 Status ArrowS3FileIO::SetStorageCredentials(
     const std::vector<StorageCredential>& storage_credentials) {
-  std::vector<std::pair<std::string, std::unique_ptr<ArrowFileSystemFileIO>>>
-      file_io_by_prefix;
-  file_io_by_prefix.reserve(storage_credentials.size());
+  ICEBERG_ASSIGN_OR_RAISE(auto delegates, BuildDelegates(storage_credentials));
+  auto credentials = storage_credentials;
+  {
+    std::unique_lock lock(mutex_);
+    InstallCredentials(credentials, delegates);
+  }
+  // `credentials` and `delegates` now hold the retired generation and destruct
+  // here, outside the lock.
+  return {};
+}
+
+Result<ArrowS3FileIO::DelegatesByPrefix> ArrowS3FileIO::BuildDelegates(
+    const std::vector<StorageCredential>& storage_credentials) const {
+  DelegatesByPrefix delegates;
+  delegates.reserve(storage_credentials.size());
   // TODO(gangwu): Refresh vended credentials via credentials.uri before tokens expire.
   for (const auto& credential : storage_credentials) {
     ICEBERG_RETURN_UNEXPECTED(credential.Validate());
@@ -231,62 +274,91 @@ Status ArrowS3FileIO::SetStorageCredentials(
       properties[key] = value;
     }
     ICEBERG_ASSIGN_OR_RAISE(auto fs, BuildArrowS3FileSystem(properties));
-    file_io_by_prefix.emplace_back(
-        CanonicalizeS3Scheme(credential.prefix),
-        std::make_unique<ArrowFileSystemFileIO>(std::move(fs)));
+    delegates.emplace_back(CanonicalizeS3Scheme(credential.prefix),
+                           std::make_shared<ArrowFileSystemFileIO>(std::move(fs)));
   }
-  if (file_io_by_prefix.empty() && !storage_credentials.empty()) {
+  if (delegates.empty() && !storage_credentials.empty()) {
     // Silent skipping of every vended credential is hard to diagnose: S3 access
     // would proceed with the default credentials and fail only at IO time.
     ICEBERG_LOG_WARN(
         "None of the {} vended storage credential(s) has an S3-compatible prefix; "
         "S3 access will use the default credentials",
         storage_credentials.size());
   }
-  file_io_by_prefix_ = std::move(file_io_by_prefix);
-  storage_credentials_ = storage_credentials;
-  return {};
+  return delegates;
 }
 
-ArrowFileSystemFileIO& ArrowS3FileIO::FileIOForPath(std::string_view location) {
-  if (file_io_by_prefix_.empty()) {
-    return default_file_io_;
+void ArrowS3FileIO::InstallCredentials(
+    std::vector<StorageCredential>& storage_credentials, DelegatesByPrefix& delegates) {
+  file_io_by_prefix_.swap(delegates);
+  storage_credentials_.swap(storage_credentials);
+}
+
+std::shared_ptr<ArrowFileSystemFileIO> ArrowS3FileIO::MatchDelegate(
+    const std::shared_ptr<ArrowFileSystemFileIO>& fallback,
+    const DelegatesByPrefix& by_prefix, std::string_view location) {
+  if (by_prefix.empty()) {
+    return fallback;
   }
   const std::string canonical = CanonicalizeS3Scheme(location);
-  ArrowFileSystemFileIO* best = &default_file_io_;
+  auto best = fallback;
   size_t best_len = 0;
-  for (const auto& [prefix, file_io] : file_io_by_prefix_) {
+  for (const auto& [prefix, file_io] : by_prefix) {
     if (prefix.size() > best_len && canonical.starts_with(prefix)) {
-      best = file_io.get();
+      best = file_io;
       best_len = prefix.size();
     }
   }
-  return *best;
+  return best;
+}
+
+std::shared_ptr<ArrowFileSystemFileIO> ArrowS3FileIO::FileIOForPath(
+    std::string_view location) {
+  std::shared_lock lock(mutex_);
+  return MatchDelegate(default_file_io_, file_io_by_prefix_, location);
 }
 
 Result<std::unique_ptr<InputFile>> ArrowS3FileIO::NewInputFile(
     std::string file_location) {
-  return FileIOForPath(file_location).NewInputFile(std::move(file_location));
+  return FileIOForPath(file_location)->NewInputFile(std::move(file_location));
 }
 
 Result<std::unique_ptr<InputFile>> ArrowS3FileIO::NewInputFile(std::string file_location,
                                                                size_t length) {
-  return FileIOForPath(file_location).NewInputFile(std::move(file_location), length);
+  return FileIOForPath(file_location)->NewInputFile(std::move(file_location), length);
 }
 
 Result<std::unique_ptr<OutputFile>> ArrowS3FileIO::NewOutputFile(
     std::string file_location) {
-  return FileIOForPath(file_location).NewOutputFile(std::move(file_location));
+  return FileIOForPath(file_location)->NewOutputFile(std::move(file_location));
 }
 
 Status ArrowS3FileIO::DeleteFile(const std::string& file_location) {
-  return FileIOForPath(file_location).DeleteFile(file_location);
+  return FileIOForPath(file_location)->DeleteFile(file_location);
 }
 
 Status ArrowS3FileIO::DeleteFiles(const std::vector<std::string>& file_locations) {
-  std::unordered_map<ArrowFileSystemFileIO*, std::vector<std::string>> locations_by_io;
+  // One snapshot so the whole batch matches the same delegate generation; only
+  // ever a handful of delegates, so a linear scan beats hashing.
+  std::shared_ptr<ArrowFileSystemFileIO> fallback;
+  DelegatesByPrefix by_prefix;
+  {
+    std::shared_lock lock(mutex_);
+    fallback = default_file_io_;
+    by_prefix = file_io_by_prefix_;
+  }
+  std::vector<std::pair<std::shared_ptr<ArrowFileSystemFileIO>, std::vector<std::string>>>
+      locations_by_io;
   for (const auto& file_location : file_locations) {
-    locations_by_io[&FileIOForPath(file_location)].push_back(file_location);
+    auto file_io = MatchDelegate(fallback, by_prefix, file_location);
+    auto it = std::ranges::find_if(
+        locations_by_io, [&](const auto& entry) { return entry.first == file_io; });
+    if (it == locations_by_io.end()) {
+      locations_by_io.emplace_back(std::move(file_io),
+                                   std::vector<std::string>{file_location});
+    } else {
+      it->second.push_back(file_location);
+    }
   }
   for (auto& [file_io, locations] : locations_by_io) {
     ICEBERG_RETURN_UNEXPECTED(file_io->DeleteFiles(locations));

src/iceberg/file_io.h

@@ -193,8 +193,11 @@ class ICEBERG_EXPORT SupportsStorageCredentials {
   virtual Status SetStorageCredentials(
       const std::vector<StorageCredential>& storage_credentials) = 0;
 
-  /// \brief Return currently installed storage credentials.
-  virtual const std::vector<StorageCredential>& credentials() const = 0;
+  /// \brief Return the storage credentials this FileIO holds.
+  ///
+  /// By value because a concurrent install may replace them. An implementation
+  /// that delegates may report what was installed on it.
+  virtual std::vector<StorageCredential> credentials() const = 0;
 };
 
 }  // namespace iceberg

src/iceberg/resolving_file_io.cc

@@ -109,7 +109,8 @@ Status ResolvingFileIO::SetStorageCredentials(
   return {};
 }
 
-const std::vector<StorageCredential>& ResolvingFileIO::credentials() const {
+std::vector<StorageCredential> ResolvingFileIO::credentials() const {
+  std::shared_lock lock(mutex_);
   return storage_credentials_;
 }
 

src/iceberg/resolving_file_io.h

@@ -38,6 +38,9 @@
 namespace iceberg {
 
 /// \brief FileIO that resolves and caches implementations by registry name.
+///
+/// Vended credentials are forwarded to every resolved implementation that
+/// supports them; each applies what it understands.
 class ICEBERG_EXPORT ResolvingFileIO final : public FileIO,
                                              public SupportsStorageCredentials {
  public:
@@ -58,7 +61,7 @@ class ICEBERG_EXPORT ResolvingFileIO final : public FileIO,
   Status SetStorageCredentials(
       const std::vector<StorageCredential>& storage_credentials) override;
 
-  const std::vector<StorageCredential>& credentials() const override;
+  std::vector<StorageCredential> credentials() const override;
 
   SupportsStorageCredentials* AsSupportsStorageCredentials() override { return this; }
 
@@ -67,8 +70,8 @@ class ICEBERG_EXPORT ResolvingFileIO final : public FileIO,
   Result<std::shared_ptr<FileIO>> FileIOForPath(std::string_view location);
 
   std::unordered_map<std::string, std::string> properties_;
-  // Guards lazy resolution and credential refresh.
-  std::shared_mutex mutex_;
+  // Guards lazy resolution and credential state.
+  mutable std::shared_mutex mutex_;
   std::vector<StorageCredential> storage_credentials_;
   std::unordered_map<std::string, std::shared_ptr<FileIO>, StringHash, StringEqual>
       io_by_name_;

src/iceberg/test/arrow_s3_file_io_test.cc

@@ -18,12 +18,14 @@
  */
 
 #include <algorithm>
+#include <atomic>
 #include <cstdlib>
 #include <iostream>
 #include <memory>
 #include <optional>
 #include <string>
 #include <string_view>
+#include <thread>
 #include <unordered_map>
 #include <utility>
 #include <vector>
@@ -237,6 +239,69 @@ TEST_F(ArrowS3FileIOTest, WarnsWhenNoCredentialApplies) {
   EXPECT_TRUE(HasWarning(*logger));
 }
 
+TEST_F(ArrowS3FileIOTest, DeleteFilesDispatchesAcrossCredentialPrefixes) {
+  auto result = MakeS3FileIO({});
+  ASSERT_THAT(result, IsOk());
+  auto* credentialed = result.value()->AsSupportsStorageCredentials();
+  ASSERT_NE(credentialed, nullptr);
+
+  auto credential = [](std::string_view prefix, std::string_view access_key) {
+    return StorageCredential{
+        .prefix = std::string(prefix),
+        .config = {{std::string(S3Properties::kAccessKeyId), std::string(access_key)},
+                   {std::string(S3Properties::kSecretAccessKey), "secret"}}};
+  };
+  ASSERT_THAT(credentialed->SetStorageCredentials({credential("s3://bucket-a", "key-a"),
+                                                   credential("s3://bucket-b", "key-b")}),
+              IsOk());
+
+  auto status = result.value()->DeleteFiles({"s3://bucket-a/%ZZ.parquet",
+                                             "s3://bucket-a/second.parquet",
+                                             "s3://bucket-b/other.parquet"});
+  EXPECT_THAT(status, HasErrorMessage("Cannot parse URI"));
+}
+
+TEST_F(ArrowS3FileIOTest, OperationsSurviveConcurrentCredentialInstalls) {
+  auto result = MakeS3FileIO({});
+  ASSERT_THAT(result, IsOk());
+  auto* credentialed = result.value()->AsSupportsStorageCredentials();
+  ASSERT_NE(credentialed, nullptr);
+
+  auto credential = [](std::string_view access_key) {
+    return StorageCredential{
+        .prefix = "s3://bucket",
+        .config = {{std::string(S3Properties::kAccessKeyId), std::string(access_key)},
+                   {std::string(S3Properties::kSecretAccessKey), "secret"}}};
+  };
+  ASSERT_THAT(credentialed->SetStorageCredentials({credential("first")}), IsOk());
+
+  std::atomic<bool> stop = false;
+  std::atomic<int> failures = 0;
+  std::vector<std::thread> operations;
+  operations.reserve(4);
+  for (int i = 0; i < 4; ++i) {
+    operations.emplace_back([&] {
+      while (!stop.load()) {
+        if (!result.value()->NewInputFile("s3://bucket/key").has_value()) {
+          ++failures;
+        }
+      }
+    });
+  }
+  // No assertions until the threads are joined: a fatal assertion here would
+  // destroy joinable threads and terminate the binary, masking the failure.
+  Status install_status = {};
+  for (int round = 0; round < 3 && install_status.has_value(); ++round) {
+    install_status = credentialed->SetStorageCredentials({credential("replacement")});
+  }
+  stop = true;
+  for (auto& operation : operations) {
+    operation.join();
+  }
+  ASSERT_THAT(install_status, IsOk());
+  EXPECT_EQ(failures, 0);
+}
+
 TEST_F(ArrowS3FileIOTest, RejectsIncompleteStaticCredentials) {
   auto result =
       MakeS3FileIO({{std::string(S3Properties::kAccessKeyId), "access-key-only"}});

src/iceberg/test/resolving_file_io_test.cc

@@ -60,9 +60,7 @@ class RecordingCredentialedFileIO : public RecordingFileIO,
     return {};
   }
 
-  const std::vector<StorageCredential>& credentials() const override {
-    return credentials_;
-  }
+  std::vector<StorageCredential> credentials() const override { return credentials_; }
 
   SupportsStorageCredentials* AsSupportsStorageCredentials() override { return this; }
 

src/iceberg/test/rest_file_io_test.cc

@@ -61,7 +61,7 @@ class MockCredentialedFileIO : public MockFileIO, public SupportsStorageCredenti
     return {};
   }
 
-  const std::vector<StorageCredential>& credentials() const override {
+  std::vector<StorageCredential> credentials() const override {
     return captured_storage_credentials;
   }