diff --git a/src/iceberg/arrow/s3/arrow_s3_file_io.cc b/src/iceberg/arrow/s3/arrow_s3_file_io.cc index 7c2799f7d..95f277c8b 100644 --- a/src/iceberg/arrow/s3/arrow_s3_file_io.cc +++ b/src/iceberg/arrow/s3/arrow_s3_file_io.cc @@ -17,9 +17,12 @@ * under the License. */ +#include #include #include +#include #include +#include #include #include #include @@ -190,7 +193,7 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials { public: ArrowS3FileIO(std::shared_ptr<::arrow::fs::FileSystem> arrow_fs, std::unordered_map default_properties) - : default_file_io_(std::move(arrow_fs)), + : default_file_io_(std::make_shared(std::move(arrow_fs))), default_properties_(std::move(default_properties)) {} Result> NewInputFile(std::string file_location) override; @@ -207,27 +210,67 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials { Status SetStorageCredentials( const std::vector& storage_credentials) override; - const std::vector& credentials() const override { + std::vector 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 FileIOForPath(std::string_view location); + + using DelegatesByPrefix = + std::vector>>; + + /// \brief Longest-prefix match against one consistent view of the delegates. + static std::shared_ptr MatchDelegate( + const std::shared_ptr& 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 BuildDelegates( + const std::vector& 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& storage_credentials, + DelegatesByPrefix& delegates); + + std::shared_ptr default_file_io_; std::unordered_map default_properties_; + // Guards everything below; shared because reads happen per file operation. + mutable std::shared_mutex mutex_; std::vector storage_credentials_; - std::vector>> - file_io_by_prefix_; + DelegatesByPrefix file_io_by_prefix_; }; Status ArrowS3FileIO::SetStorageCredentials( const std::vector& storage_credentials) { - std::vector>> - 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::BuildDelegates( + const std::vector& 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()); @@ -242,11 +285,10 @@ 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(std::move(fs))); + delegates.emplace_back(CanonicalizeS3Scheme(credential.prefix), + std::make_shared(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( @@ -254,50 +296,80 @@ Status ArrowS3FileIO::SetStorageCredentials( "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& storage_credentials, DelegatesByPrefix& delegates) { + file_io_by_prefix_.swap(delegates); + storage_credentials_.swap(storage_credentials); +} + +std::shared_ptr ArrowS3FileIO::MatchDelegate( + const std::shared_ptr& 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 ArrowS3FileIO::FileIOForPath( + std::string_view location) { + std::shared_lock lock(mutex_); + return MatchDelegate(default_file_io_, file_io_by_prefix_, location); } Result> 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> 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> 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& file_locations) { - std::unordered_map> 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 fallback; + DelegatesByPrefix by_prefix; + { + std::shared_lock lock(mutex_); + fallback = default_file_io_; + by_prefix = file_io_by_prefix_; + } + std::vector, std::vector>> + 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{file_location}); + } else { + it->second.push_back(file_location); + } } for (auto& [file_io, locations] : locations_by_io) { ICEBERG_RETURN_UNEXPECTED(file_io->DeleteFiles(locations)); diff --git a/src/iceberg/file_io.h b/src/iceberg/file_io.h index 3ea4afa49..e22e5cf21 100644 --- a/src/iceberg/file_io.h +++ b/src/iceberg/file_io.h @@ -193,8 +193,11 @@ class ICEBERG_EXPORT SupportsStorageCredentials { virtual Status SetStorageCredentials( const std::vector& storage_credentials) = 0; - /// \brief Return currently installed storage credentials. - virtual const std::vector& 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 credentials() const = 0; }; } // namespace iceberg diff --git a/src/iceberg/resolving_file_io.cc b/src/iceberg/resolving_file_io.cc index 8a4e81138..72b7e1576 100644 --- a/src/iceberg/resolving_file_io.cc +++ b/src/iceberg/resolving_file_io.cc @@ -109,7 +109,8 @@ Status ResolvingFileIO::SetStorageCredentials( return {}; } -const std::vector& ResolvingFileIO::credentials() const { +std::vector ResolvingFileIO::credentials() const { + std::shared_lock lock(mutex_); return storage_credentials_; } diff --git a/src/iceberg/resolving_file_io.h b/src/iceberg/resolving_file_io.h index 837afd60e..d2f1254dd 100644 --- a/src/iceberg/resolving_file_io.h +++ b/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& storage_credentials) override; - const std::vector& credentials() const override; + std::vector credentials() const override; SupportsStorageCredentials* AsSupportsStorageCredentials() override { return this; } @@ -67,8 +70,8 @@ class ICEBERG_EXPORT ResolvingFileIO final : public FileIO, Result> FileIOForPath(std::string_view location); std::unordered_map properties_; - // Guards lazy resolution and credential refresh. - std::shared_mutex mutex_; + // Guards lazy resolution and credential state. + mutable std::shared_mutex mutex_; std::vector storage_credentials_; std::unordered_map, StringHash, StringEqual> io_by_name_; diff --git a/src/iceberg/test/arrow_s3_file_io_test.cc b/src/iceberg/test/arrow_s3_file_io_test.cc index 40719827a..238a90909 100644 --- a/src/iceberg/test/arrow_s3_file_io_test.cc +++ b/src/iceberg/test/arrow_s3_file_io_test.cc @@ -18,12 +18,14 @@ */ #include +#include #include #include #include #include #include #include +#include #include #include #include @@ -226,6 +228,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 stop = false; + std::atomic failures = 0; + std::vector 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"}}); diff --git a/src/iceberg/test/resolving_file_io_test.cc b/src/iceberg/test/resolving_file_io_test.cc index 788b59beb..1f01cc517 100644 --- a/src/iceberg/test/resolving_file_io_test.cc +++ b/src/iceberg/test/resolving_file_io_test.cc @@ -60,9 +60,7 @@ class RecordingCredentialedFileIO : public RecordingFileIO, return {}; } - const std::vector& credentials() const override { - return credentials_; - } + std::vector credentials() const override { return credentials_; } SupportsStorageCredentials* AsSupportsStorageCredentials() override { return this; } diff --git a/src/iceberg/test/rest_file_io_test.cc b/src/iceberg/test/rest_file_io_test.cc index dde3238a8..2fce9dacf 100644 --- a/src/iceberg/test/rest_file_io_test.cc +++ b/src/iceberg/test/rest_file_io_test.cc @@ -61,7 +61,7 @@ class MockCredentialedFileIO : public MockFileIO, public SupportsStorageCredenti return {}; } - const std::vector& credentials() const override { + std::vector credentials() const override { return captured_storage_credentials; }