refactor(io): load resolver delegates off the lock; test real IO across installs

ResolvingFileIO::FileIOForPath held its mutex across FileIORegistry::Load and
the delegate's SetStorageCredentials, so building the first S3 client (which
can look up a bucket region) stalled every other operation on the resolver.
It now loads without the lock and caches the result only if no credential
install happened meanwhile; otherwise it loads again with the new
credentials. Delegates retired by an install are torn down outside the lock.

Tests now exercise real IO where the old ones stopped short: the multi-prefix
DeleteFiles test failed URI parsing before reaching any delegate, and the
concurrency test only built input-file wrappers. The replacements write,
delete and read through the test object store, and the stress test waits
for every worker to run before installing.

Author: plusplusjiajia

src/iceberg/arrow/s3/arrow_s3_file_io.cc

@@ -221,7 +221,7 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials {
 
   /// \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
+  /// Runs without holding `mutex_`: 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(

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 reach the network (an
+  // S3 client may look up its bucket region), 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,9 +119,14 @@ 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 {};
 }
 

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>
@@ -73,6 +74,9 @@ class ICEBERG_EXPORT ResolvingFileIO final : public FileIO,
   // 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_;
 };

src/iceberg/test/arrow_s3_file_io_test.cc

@@ -23,6 +23,7 @@
 #include <iostream>
 #include <memory>
 #include <optional>
+#include <span>
 #include <string>
 #include <string_view>
 #include <thread>
@@ -239,28 +240,6 @@ 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());
@@ -276,18 +255,28 @@ TEST_F(ArrowS3FileIOTest, OperationsSurviveConcurrentCredentialInstalls) {
   ASSERT_THAT(credentialed->SetStorageCredentials({credential("first")}), IsOk());
 
   std::atomic<bool> stop = false;
+  std::atomic<int> started = 0;
   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()) {
+      auto open = [&] {
         if (!result.value()->NewInputFile("s3://bucket/key").has_value()) {
           ++failures;
         }
+      };
+      open();
+      ++started;
+      while (!stop.load()) {
+        open();
       }
     });
   }
+  // Every worker has run and is still looping before the first install.
+  while (started.load() < 4) {
+    std::this_thread::yield();
+  }
   // 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 = {};
@@ -406,6 +395,90 @@ TEST_F(ArrowS3FileIOTest, AppliesOssCredentialInRealRoundTrip) {
   EXPECT_THAT(CheckReadWrite(*io, s3_uri, "hello oss with vended credentials"), IsOk());
 }
 
+TEST_F(ArrowS3FileIOTest, DeleteFilesReachesEveryCredentialPrefix) {
+  if (!HasIntegrationEnv()) {
+    GTEST_SKIP() << "Set ICEBERG_TEST_S3_URI to enable S3 IO test";
+  }
+
+  auto properties = PropertiesFromEnv();
+  if (!properties.contains(std::string(S3Properties::kAccessKeyId)) ||
+      !properties.contains(std::string(S3Properties::kSecretAccessKey))) {
+    GTEST_SKIP() << "Set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY to enable "
+                    "credential routing test";
+  }
+
+  // Only the prefix delegates can authenticate, so a location that falls to
+  // the default one fails the batch instead of passing unnoticed.
+  auto bad_defaults = properties;
+  for (const auto& [key, value] : BadS3Credentials()) {
+    bad_defaults.insert_or_assign(key, value);
+  }
+  ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO(std::move(bad_defaults)));
+  auto* credentialed = io->AsSupportsStorageCredentials();
+  ASSERT_NE(credentialed, nullptr);
+
+  const auto a = ObjectUri("delete_files_a/");
+  const auto b = ObjectUri("delete_files_b/");
+  ASSERT_THAT(credentialed->SetStorageCredentials({{.prefix = a, .config = properties},
+                                                   {.prefix = b, .config = properties}}),
+              IsOk());
+
+  const std::vector<std::string> paths = {a + "first", b + "only", a + "second"};
+  for (const auto& path : paths) {
+    ASSERT_THAT(io->WriteFile(path, "payload"), IsOk());
+  }
+  ASSERT_THAT(io->DeleteFiles(paths), IsOk());
+  // Deleting a missing key succeeds on S3, so check the objects are gone. The
+  // writes above authenticated on these paths, so a failed read means absent.
+  for (const auto& path : paths) {
+    EXPECT_FALSE(io->ReadFile(path, std::nullopt).has_value()) << path;
+  }
+}
+
+TEST_F(ArrowS3FileIOTest, InputFileOutlivesCredentialInstall) {
+  if (!HasIntegrationEnv()) {
+    GTEST_SKIP() << "Set ICEBERG_TEST_S3_URI to enable S3 IO test";
+  }
+
+  auto properties = PropertiesFromEnv();
+  if (!properties.contains(std::string(S3Properties::kAccessKeyId)) ||
+      !properties.contains(std::string(S3Properties::kSecretAccessKey))) {
+    GTEST_SKIP() << "Set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY to enable "
+                    "credential routing test";
+  }
+
+  auto bad_defaults = properties;
+  for (const auto& [key, value] : BadS3Credentials()) {
+    bad_defaults.insert_or_assign(key, value);
+  }
+  ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO(std::move(bad_defaults)));
+  auto* credentialed = io->AsSupportsStorageCredentials();
+  ASSERT_NE(credentialed, nullptr);
+
+  const auto prefix = ObjectUri("retained_input/");
+  const std::vector<StorageCredential> credentials = {
+      {.prefix = prefix, .config = properties}};
+  ASSERT_THAT(credentialed->SetStorageCredentials(credentials), IsOk());
+
+  const auto uri = prefix + "object";
+  constexpr std::string_view kContent = "written before the credential install";
+  ASSERT_THAT(io->WriteFile(uri, kContent), IsOk());
+  ICEBERG_UNWRAP_OR_FAIL(auto file, io->NewInputFile(uri));
+  ICEBERG_UNWRAP_OR_FAIL(auto opened_before, file->Open());
+
+  // Retires the delegate `file` came from; both handles must still read.
+  ASSERT_THAT(credentialed->SetStorageCredentials(credentials), IsOk());
+
+  ICEBERG_UNWRAP_OR_FAIL(auto opened_after, file->Open());
+  for (auto* stream : {opened_before.get(), opened_after.get()}) {
+    std::string read(kContent.size(), '\0');
+    EXPECT_THAT(stream->ReadFully(0, std::as_writable_bytes(std::span(read))), IsOk());
+    EXPECT_EQ(read, kContent);
+    EXPECT_THAT(stream->Close(), IsOk());
+  }
+  EXPECT_THAT(io->DeleteFile(uri), IsOk());
+}
+
 #if ICEBERG_S3_ENABLED
 TEST_F(ArrowS3FileIOTest, ClientRegion) {
   auto result =

src/iceberg/test/resolving_file_io_test.cc

@@ -19,8 +19,13 @@
 
 #include "iceberg/resolving_file_io.h"
 
+#include <chrono>
+#include <condition_variable>
+#include <future>
 #include <memory>
+#include <mutex>
 #include <string>
+#include <thread>
 #include <unordered_map>
 #include <vector>
 
@@ -254,4 +259,63 @@ TEST(ResolvingFileIOTest, ForwardsAllCredentialsToResolvedImplementations) {
   ASSERT_NE(last_local_io, nullptr);
 }
 
+TEST(ResolvingFileIOTest, LoadsWithoutTheLockAndDropsStaleDelegates) {
+  // The first load parks until new credentials are installed. Loading outside
+  // the lock lets that install proceed; what the parked load built then carries
+  // stale credentials and must be redone rather than cached.
+  // Static because registry factories are process-global; reset for reruns.
+  static std::mutex gate;
+  static std::condition_variable cv;
+  static bool parked;
+  static bool resume;
+  static int calls;
+  static RecordingCredentialedFileIO* last;
+  parked = resume = false;
+  calls = 0;
+  last = nullptr;
+  FileIORegistry::Register(
+      "test.file-io.slow",
+      {.create =
+           [](const FileIORegistry::Properties&) -> Result<std::unique_ptr<FileIO>> {
+         if (++calls == 1) {
+           std::unique_lock lock(gate);
+           parked = true;
+           cv.notify_all();
+           cv.wait(lock, [] { return resume; });
+         }
+         auto io = std::make_unique<RecordingCredentialedFileIO>();
+         last = io.get();
+         return io;
+       },
+       .accepts = [](std::string_view scheme) { return scheme == "slow"; }});
+
+  ResolvingFileIO io({});
+  const std::vector<StorageCredential> stale = {
+      {.prefix = "slow", .config = {{"k", "1"}}}};
+  const std::vector<StorageCredential> fresh = {
+      {.prefix = "slow", .config = {{"k", "2"}}}};
+  ASSERT_THAT(io.SetStorageCredentials(stale), IsOk());
+
+  std::thread reader([&] { (void)io.NewInputFile("slow://bucket/file"); });
+  {
+    std::unique_lock lock(gate);
+    cv.wait(lock, [] { return parked; });
+  }
+  auto install =
+      std::async(std::launch::async, [&] { return io.SetStorageCredentials(fresh); });
+  const bool install_waited =
+      install.wait_for(std::chrono::seconds(10)) != std::future_status::ready;
+  {
+    std::lock_guard lock(gate);
+    resume = true;
+  }
+  cv.notify_all();
+  reader.join();
+
+  EXPECT_FALSE(install_waited) << "a credential install waited on an in-flight load";
+  EXPECT_THAT(install.get(), IsOk());
+  ASSERT_EQ(calls, 2);
+  EXPECT_EQ(last->credentials(), fresh);
+}
+
 }  // namespace iceberg