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

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.
+  ///
+  /// Runs without holding `mutex_`: building an S3 client can block (without
+  /// static keys the AWS SDK may wait on the EC2 metadata service), 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

@@ -40,28 +40,44 @@ Result<std::shared_ptr<FileIO>> ResolvingFileIO::FileIOForPath(
   const auto scheme = StringUtils::ToLower(LocationUtil::ParseScheme(location));
   ICEBERG_ASSIGN_OR_RAISE(const auto name, FileIORegistry::Resolve(scheme));
 
-  {
-    std::shared_lock lock(mutex_);
-    if (const auto cached = io_by_name_.find(name); cached != io_by_name_.end()) {
-      return cached->second;
-    }
-  }
-
-  std::unique_lock lock(mutex_);
-  auto it = io_by_name_.find(name);
-  if (it == io_by_name_.end()) {
-    ICEBERG_ASSIGN_OR_RAISE(auto io, FileIORegistry::Load(name, properties_));
-    // Forward all credentials; each implementation applies the prefixes it
-    // understands.
-    if (!storage_credentials_.empty()) {
+  // Loads without holding `mutex_`: building a client can block (an S3 client
+  // without static keys may wait on the EC2 metadata service), which would
+  // stall every other operation. Forwards all credentials; each implementation
+  // applies the prefixes it understands.
+  auto load = [&](const std::vector<StorageCredential>& credentials)
+      -> Result<std::shared_ptr<FileIO>> {
+    ICEBERG_ASSIGN_OR_RAISE(std::shared_ptr<FileIO> io,
+                            FileIORegistry::Load(name, properties_));
+    if (!credentials.empty()) {
       if (auto* credentialed = io->AsSupportsStorageCredentials()) {
-        ICEBERG_RETURN_UNEXPECTED(
-            credentialed->SetStorageCredentials(storage_credentials_));
+        ICEBERG_RETURN_UNEXPECTED(credentialed->SetStorageCredentials(credentials));
       }
     }
-    it = io_by_name_.emplace(std::string(name), std::move(io)).first;
+    return io;
+  };
+
+  while (true) {
+    uint64_t generation = 0;
+    std::vector<StorageCredential> credentials;
+    {
+      std::shared_lock lock(mutex_);
+      if (const auto cached = io_by_name_.find(name); cached != io_by_name_.end()) {
+        return cached->second;
+      }
+      generation = credential_generation_;
+      credentials = storage_credentials_;
+    }
+    // Declared before the lock, so a delegate that is not cached is torn down
+    // only after the lock is released.
+    auto loaded = load(credentials);
+    std::unique_lock lock(mutex_);
+    if (generation != credential_generation_) {
+      continue;  // Credentials were replaced mid-load; load again with them.
+    }
+    ICEBERG_RETURN_UNEXPECTED(loaded);
+    // A concurrent first access may have cached one already; that one wins.
+    return io_by_name_.try_emplace(name, *loaded).first->second;
   }
-  return it->second;
 }
 
 Result<std::unique_ptr<InputFile>> ResolvingFileIO::NewInputFile(
@@ -103,13 +119,19 @@ Status ResolvingFileIO::SetStorageCredentials(
     const std::vector<StorageCredential>& storage_credentials) {
   // Rebuild delegates lazily with the new credentials. Updating live delegates
   // instead would leave the resolver inconsistent if one of them rejected them.
-  std::unique_lock lock(mutex_);
-  storage_credentials_ = storage_credentials;
-  io_by_name_.clear();
+  // Retired outside the lock: tearing down a delegate can block.
+  decltype(io_by_name_) retired;
+  {
+    std::unique_lock lock(mutex_);
+    storage_credentials_ = storage_credentials;
+    ++credential_generation_;
+    retired.swap(io_by_name_);
+  }
   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

@@ -22,6 +22,7 @@
 /// \file iceberg/resolving_file_io.h
 /// \brief FileIO that resolves the concrete implementation per file-path scheme.
 
+#include <cstdint>
 #include <memory>
 #include <shared_mutex>
 #include <string>
@@ -38,6 +39,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 +62,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,9 +71,12 @@ 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_;
+  // Bumped by every credential install, so a delegate loaded from an older set
+  // never reaches the cache.
+  uint64_t credential_generation_ = 0;
   std::unordered_map<std::string, std::shared_ptr<FileIO>, StringHash, StringEqual>
       io_by_name_;
 };