diff --git a/include/cucascade/io/datasource_factory.hpp b/include/cucascade/io/datasource_factory.hpp index a73ec2a..f172719 100644 --- a/include/cucascade/io/datasource_factory.hpp +++ b/include/cucascade/io/datasource_factory.hpp @@ -21,6 +21,7 @@ #include #include +#include #include #include #include @@ -80,6 +81,39 @@ class io_context_registry { */ void register_ioctx(io_context_type type, scheme_checker_type checker, factory_type factory); + /** + * @brief Replaces one backend registration before path routing begins. + * + * A URI scheme may have more than one backend implementation, but only one + * of them should claim that scheme explicitly. For example, the registry + * initially assigns `s3://` paths to @c io_context_type::restful; when + * S3-over-RDMA is selected, the caller replaces that registration with + * @c io_context_type::s3rdma. + * + * The checker and factory for @p new_type are installed and the @p old_type + * entry is removed while holding the registry's exclusive lock, so no lookup + * observes an intermediate routing state. + * + * This operation changes registration metadata only. It does not construct, + * shut down, or migrate any ioctx instance. + * + * @param old_type Registered backend to remove. + * @param new_type Backend type to register. + * @param checker Predicate used to claim paths for @p new_type. + * @param factory Factory used by @c make_ioctx for @p new_type. + * + * @throws std::invalid_argument if @p old_type is absent, @p new_type is + * already registered, or @p checker or @p factory is empty. + * @throws std::logic_error if @c lookup_path has already been called. + * + * Provides the strong exception guarantee: the registry is unchanged if the + * replacement fails. Intended for single-threaded bootstrap. + */ + void replace_ioctx(io_context_type old_type, + io_context_type new_type, + scheme_checker_type checker, + factory_type factory); + /// Resolve the backend for a full @p path (not a bare scheme — the checkers /// parse the URI / stat the filesystem themselves). Explicit backends /// (uring / restful) take precedence over the kvikio catch-all, so `s3://` @@ -105,6 +139,9 @@ class io_context_registry { cucascade::memory::memory_reservation_manager& _reservation_manager; mutable std::shared_mutex _mtx; std::unordered_map _entries; + /// Set by the first @c lookup_path; @c replace_ioctx refuses afterwards + /// (bootstrap-only — see its contract). + mutable std::atomic _lookup_latched{false}; }; // --------------------------------------------------------------------------- diff --git a/include/cucascade/io/io_context.hpp b/include/cucascade/io/io_context.hpp index 97d87b5..d6c5d4d 100644 --- a/include/cucascade/io/io_context.hpp +++ b/include/cucascade/io/io_context.hpp @@ -37,7 +37,7 @@ namespace cucascade::io { -enum class io_context_type { uring, restful, kvikio }; +enum class io_context_type { uring, restful, kvikio, s3rdma }; /// Hint passed to @c open_io_object so a backend can tailor how it resolves an /// object's metadata. @c generic resolves the size however is cheapest for the diff --git a/include/cucascade/io/rest/object_store_lister.hpp b/include/cucascade/io/rest/object_store_lister.hpp new file mode 100644 index 0000000..5359165 --- /dev/null +++ b/include/cucascade/io/rest/object_store_lister.hpp @@ -0,0 +1,103 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +// The page/entry types live under rest/s3 but are S3-PROTOCOL shapes +// (ListObjectsV2 responses), not REST-transport shapes — any backend that +// lists an S3-compatible store speaks them, whatever its data plane. +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace cucascade::io::rest { + +// --------------------------------------------------------------------------- +// object_store_lister +// --------------------------------------------------------------------------- + +/** + * @brief Backend-ioctx-owned control-plane component for object-store listing. + * + * Each S3-capable ioctx composes one instance and injects its own page-fetch + * operation. The lister owns the ListObjectsV2 protocol logic — canonical + * query construction, page-size clamping, response parsing, continuation-token + * advancement, the anti-loop guards, and the @c max_scanned / @c max_matches + * caps. Transport, retries, signing, and admission policy stay behind the + * injected fetch. + */ +class object_store_lister { + public: + /// Fetches one raw ListObjectsV2 response body for @p bucket. + /// @p canonical_query is the fully encoded query string in SigV4 canonical + /// order; @p prefix repeats the (unencoded) prefix inside it for fetchers + /// that address by prefix. + using page_fetch_fn = std::function; + + /// @p error_context prefixes the messages of exceptions the lister itself + /// throws; exceptions from the fetch, the parser, or the sink propagate + /// unwrapped. + object_store_lister(page_fetch_fn fetch, + std::size_t max_scanned, + std::size_t max_matches, + std::string error_context) + : _fetch(std::move(fetch)), + _max_scanned(max_scanned), + _max_matches(max_matches), + _error_context(std::move(error_context)) + { + } + + /// Stream ListObjectsV2 pages under @p prefix to @p sink, one call per + /// page. @p sink returns false to stop early. @p page_size is clamped + /// to [1,1000] (0 and >1000 mean 1000). Throws (never truncates) on a + /// truncated page without a continuation token, and once more than + /// @p max_scanned entries have been scanned across pages. + void list_objects_paged(std::string_view bucket, + std::string_view prefix, + std::size_t page_size, + std::function const& sink, + std::optional max_scanned = std::nullopt); + + /// Whole-listing convenience over @c list_objects_paged: every object under + /// @p prefix, in document order, with sizes. Throws (never truncates) when + /// the accumulated entries would exceed @p max_keys. + [[nodiscard]] std::vector list_objects( + std::string_view bucket, + std::string_view prefix, + std::size_t page_size = 1000, + std::optional max_keys = std::nullopt); + + /// The configured matched cap for glob resolution. + [[nodiscard]] std::size_t list_max_matches() const noexcept { return _max_matches; } + + private: + page_fetch_fn _fetch; + std::size_t _max_scanned; + std::size_t _max_matches; + std::string _error_context; +}; + +} // namespace cucascade::io::rest diff --git a/include/cucascade/io/rest/rest_ioctx.hpp b/include/cucascade/io/rest/rest_ioctx.hpp index 9701570..69e1cc9 100644 --- a/include/cucascade/io/rest/rest_ioctx.hpp +++ b/include/cucascade/io/rest/rest_ioctx.hpp @@ -18,6 +18,7 @@ #pragma once +#include #include #include #include @@ -114,6 +115,8 @@ class rest_ioctx : public templated_ioctx { /// footer reads are served locally. Falls back to a plain HEAD (no stash) /// when the response is unusable. std::shared_ptr create_footer_probe_object(std::string path); + + object_store_lister _lister; }; } // namespace cucascade::io::rest diff --git a/include/cucascade/io/s3rdma/s3rdma_ioctx.hpp b/include/cucascade/io/s3rdma/s3rdma_ioctx.hpp new file mode 100644 index 0000000..34b14a3 --- /dev/null +++ b/include/cucascade/io/s3rdma/s3rdma_ioctx.hpp @@ -0,0 +1,112 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cucascade::io::s3rdma { + +// --------------------------------------------------------------------------- +// s3rdma_ioctx (placeholder) +// --------------------------------------------------------------------------- + +/** + * @brief Non-constructible placeholder for the S3-over-RDMA I/O context. + * + * The planned backend uses RDMA for S3 range-read payloads. Data lands in + * registered GPU memory and is copied on device to the caller's destination. + * LIST and HEAD use the configured HTTP(S) control plane; glob resolution + * uses LIST results. + * + * Integration uses @c io_context_type::s3rdma and + * @c io_context_registry::replace_ioctx, composes @c rest::object_store_lister + * with its own page fetch for LIST, and overrides + * @c templated_ioctx::on_device_dispatch_failure. + * + * This declaration has no implementation or factory registration. Its deleted + * constructor prevents accidental use. + * + * @see https://github.com/sirius-db/sirius/blob/dev/experimental/s3-rdma-transport-design.md + */ +class s3rdma_ioctx : public ioctx { + public: + s3rdma_ioctx() = delete; + + [[nodiscard]] io_context_type type() const noexcept override; + + void shutdown() noexcept override; + + [[nodiscard]] bool supports(std::string_view path) const noexcept override; + + [[nodiscard]] bool supports_device_read() const noexcept override; + [[nodiscard]] bool supports_host_to_device_read() const noexcept override; + [[nodiscard]] bool supports_vector_host_read() const noexcept override; + [[nodiscard]] cache::prefetching_stage preferred_prefetching_stage() const noexcept override; + + [[nodiscard]] std::vector align_and_coalesce( + std::span ranges, + std::optional alignment = std::nullopt) const noexcept override; + + size_t host_read_io(const io_object& obj, size_t offset, size_t size, uint8_t* dst) override; + + exec::semi_future host_read_async_io(const io_object& obj, + size_t offset, + size_t size, + uint8_t* dst) noexcept override; + + exec::semi_future device_read_async_io(const io_object& obj, + size_t offset, + size_t size, + uint8_t* dst, + rmm::cuda_stream_view stream) noexcept override; + + exec::semi_future host_to_device_read_async_io( + const io_object& obj, + std::span slices, + size_t offset, + size_t size, + uint8_t* dst, + rmm::cuda_stream_view stream) noexcept override; + + exec::semi_future host_read_ranges_async_io( + const io_object& obj, std::span segments) noexcept override; + + protected: + std::shared_ptr create_io_object(std::string path) override; +}; + +static_assert(!std::is_default_constructible_v, + "s3rdma_ioctx is a placeholder and must stay non-constructible " + "until the backend implementation lands"); + +} // namespace cucascade::io::s3rdma diff --git a/include/cucascade/io/templated_ioctx.hpp b/include/cucascade/io/templated_ioctx.hpp index e28976d..4547e02 100644 --- a/include/cucascade/io/templated_ioctx.hpp +++ b/include/cucascade/io/templated_ioctx.hpp @@ -380,6 +380,7 @@ class templated_ioctx : public ioctx { }); return semi; } catch (...) { + on_device_dispatch_failure(); return exec::make_semi_future(std::current_exception()); } } else { @@ -388,6 +389,26 @@ class templated_ioctx : public ioctx { } } + protected: + /** + * @brief Applies backend policy after synchronous device dispatch fails. + * + * Called from the exception handlers in device_read_async_io() and + * host_to_device_read_async_io(), before the exception is returned through + * an errored future. + * + * An S3-over-RDMA backend overrides this hook to check for a sticky CUDA + * context error. Returning such an error as an ordinary request failure + * could allow registered GPU memory to be reused or released before RDMA + * writes and CUDA work are known to be quiescent. In that case the backend + * must invoke its fatal policy instead of returning. + * + * The default implementation does nothing. An override must not throw or + * re-enter this ioctx. + */ + virtual void on_device_dispatch_failure() noexcept {} + + public: exec::semi_future host_to_device_read_async_io( const io_object& obj, std::span slices, @@ -417,6 +438,7 @@ class templated_ioctx : public ioctx { }); return semi; } catch (...) { + on_device_dispatch_failure(); return exec::make_semi_future(std::current_exception()); } } else { diff --git a/src/io/CMakeLists.txt b/src/io/CMakeLists.txt index d645721..f8f2d2a 100644 --- a/src/io/CMakeLists.txt +++ b/src/io/CMakeLists.txt @@ -21,8 +21,10 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/io_context.cpp ${CMAKE_CURRENT_SOURCE_DIR}/datasource_factory.cpp ${CMAKE_CURRENT_SOURCE_DIR}/rest/curl_handle.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/rest/object_store_lister.cpp ${CMAKE_CURRENT_SOURCE_DIR}/rest/rest_ioctx.cpp ${CMAKE_CURRENT_SOURCE_DIR}/rest/rest_reactor.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/s3rdma/s3rdma_ioctx.cpp ${CMAKE_CURRENT_SOURCE_DIR}/uring/uring_ioctx.cpp ${CMAKE_CURRENT_SOURCE_DIR}/uring/uring_reactor.cpp ${CMAKE_CURRENT_SOURCE_DIR}/kvikio/kvikio_context.cpp diff --git a/src/io/datasource_factory.cpp b/src/io/datasource_factory.cpp index edefff4..b5c52ed 100644 --- a/src/io/datasource_factory.cpp +++ b/src/io/datasource_factory.cpp @@ -185,10 +185,37 @@ void io_context_registry::register_ioctx(io_context_type type, _entries[type] = {type, std::move(checker), std::move(factory)}; } +void io_context_registry::replace_ioctx(io_context_type old_type, + io_context_type new_type, + scheme_checker_type checker, + factory_type factory) +{ + if (!checker) { + throw std::invalid_argument("datasource_registry: replace_ioctx: null scheme checker"); + } + if (!factory) { throw std::invalid_argument("datasource_registry: replace_ioctx: null factory"); } + std::lock_guard lk{_mtx}; + if (_lookup_latched.load(std::memory_order_acquire)) { + throw std::logic_error( + "datasource_registry: replace_ioctx after the first lookup_path (bootstrap-only)"); + } + if (!_entries.contains(old_type)) { + throw std::invalid_argument("datasource_registry: replace_ioctx: old type not registered"); + } + if (_entries.contains(new_type)) { + throw std::invalid_argument("datasource_registry: replace_ioctx: new type already registered"); + } + // Strong guarantee: the emplace is the only throwing step and precedes the + // erase; erase by KEY, not by a pre-emplace iterator (emplace may rehash). + _entries.emplace(new_type, entry{new_type, std::move(checker), std::move(factory)}); + _entries.erase(old_type); +} + std::optional io_context_registry::lookup_path( std::string_view path) const noexcept { std::shared_lock lk{_mtx}; + _lookup_latched.store(true, std::memory_order_release); // kvikio's checker matches everything; _entries iterates in unspecified order, // so defer the catch-all and let an explicit backend (uring/restful) win. std::optional fallback; diff --git a/src/io/rest/object_store_lister.cpp b/src/io/rest/object_store_lister.cpp new file mode 100644 index 0000000..991c4c5 --- /dev/null +++ b/src/io/rest/object_store_lister.cpp @@ -0,0 +1,107 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +#include +#include + +namespace cucascade::io::rest { + +void object_store_lister::list_objects_paged( + std::string_view bucket, + std::string_view prefix, + std::size_t page_size, + std::function const& sink, + std::optional max_scanned) +{ + std::size_t const clamped = (page_size == 0 || page_size > 1000) ? 1000 : page_size; + std::size_t const scanned_cap = max_scanned.value_or(_max_scanned); + + std::size_t scanned = 0; + std::string token; + bool truncated = false; + do { + // SigV4 canonical order = byte order of the encoded keys; for these params + // that is continuation-token < list-type < max-keys < prefix. + std::string query; + if (!token.empty()) { + query += "continuation-token="; + query += s3::uri_encode(token, /*encode_slash=*/true); + query += '&'; + } + query += "list-type=2&max-keys="; + query += std::to_string(clamped); + query += "&prefix="; + query += s3::uri_encode(prefix, /*encode_slash=*/true); + + auto const page = s3::parse_list_objects_v2(_fetch(bucket, prefix, query)); + + scanned += page.entries.size(); + if (scanned > scanned_cap) { + throw std::runtime_error(_error_context + ": scanned more than " + + std::to_string(scanned_cap) + " objects under s3://" + + std::string(bucket) + "/" + std::string(prefix) + + " — narrow the glob prefix"); + } + if (page.is_truncated && page.next_continuation_token.empty()) { + throw std::runtime_error(_error_context + + ": truncated ListObjectsV2 page without a continuation token for " + "s3://" + + std::string(bucket) + "/" + std::string(prefix)); + } + // A truncated page must contain entries and advance the token. Together + // with scanned_cap these bound pagination for non-conforming backends that + // would otherwise loop on empty or non-advancing pages. + if (page.is_truncated && page.entries.empty()) { + throw std::runtime_error(_error_context + + ": truncated ListObjectsV2 page with no entries for s3://" + + std::string(bucket) + "/" + std::string(prefix)); + } + if (page.is_truncated && page.next_continuation_token == token) { + throw std::runtime_error(_error_context + + ": ListObjectsV2 continuation token did not advance for s3://" + + std::string(bucket) + "/" + std::string(prefix)); + } + truncated = page.is_truncated; + token = page.next_continuation_token; + if (!sink(page)) { return; } + } while (truncated); +} + +std::vector object_store_lister::list_objects(std::string_view bucket, + std::string_view prefix, + std::size_t page_size, + std::optional max_keys) +{ + std::size_t const keys_cap = max_keys.value_or(_max_matches); + std::vector out; + list_objects_paged(bucket, prefix, page_size, [&](s3::list_objects_v2_page const& page) { + if (out.size() + page.entries.size() > keys_cap) { + throw std::runtime_error(_error_context + ": more than " + std::to_string(keys_cap) + + " objects under s3://" + std::string(bucket) + "/" + + std::string(prefix) + " — narrow the glob prefix"); + } + out.insert(out.end(), page.entries.begin(), page.entries.end()); + return true; + }); + return out; +} + +} // namespace cucascade::io::rest diff --git a/src/io/rest/rest_ioctx.cpp b/src/io/rest/rest_ioctx.cpp index 6fc6fb5..9e57d05 100644 --- a/src/io/rest/rest_ioctx.cpp +++ b/src/io/rest/rest_ioctx.cpp @@ -17,7 +17,6 @@ */ #include -#include #include #include @@ -30,9 +29,23 @@ namespace cucascade::io::rest { rest_ioctx::rest_ioctx(std::size_t n_reactors, std::shared_ptr ctx) - : templated_ioctx(n_reactors, [ctx = std::move(ctx), i = 0]() mutable { - return std::make_unique(ctx, "rest-" + std::to_string(i++)); - }) + : templated_ioctx(n_reactors, + [ctx = std::move(ctx), i = 0]() mutable { + return std::make_unique( + ctx, "rest-" + std::to_string(i++)); + }), + _lister( + [this](std::string_view bucket, std::string_view prefix, std::string_view canonical_query) { + if (_reactors.empty()) { + throw std::runtime_error("rest_ioctx::list_objects: no reactors"); + } + return _reactors.front()->list_page(bucket, prefix, canonical_query); + }, + _reactors.empty() ? s3::default_max_scanned_objects + : _reactors.front()->get_config().list_max_scanned, + _reactors.empty() ? s3::default_max_list_objects + : _reactors.front()->get_config().list_max_matches, + "rest_ioctx::list_objects") { } @@ -72,60 +85,7 @@ void rest_ioctx::list_objects_paged( std::optional max_scanned) { if (_reactors.empty()) { throw std::runtime_error("rest_ioctx::list_objects: no reactors"); } - std::size_t const clamped = (page_size == 0 || page_size > 1000) ? 1000 : page_size; - std::size_t const scanned_cap = - max_scanned.value_or(_reactors.front()->get_config().list_max_scanned); - - std::size_t scanned = 0; - std::string token; - bool truncated = false; - do { - // SigV4 canonical order = byte order of the encoded keys; for these params - // that is continuation-token < list-type < max-keys < prefix. - std::string query; - if (!token.empty()) { - query += "continuation-token="; - query += s3::uri_encode(token, /*encode_slash=*/true); - query += '&'; - } - query += "list-type=2&max-keys="; - query += std::to_string(clamped); - query += "&prefix="; - query += s3::uri_encode(prefix, /*encode_slash=*/true); - - auto const page = - s3::parse_list_objects_v2(_reactors.front()->list_page(bucket, prefix, query)); - - scanned += page.entries.size(); - if (scanned > scanned_cap) { - throw std::runtime_error("rest_ioctx::list_objects: scanned more than " + - std::to_string(scanned_cap) + " objects under s3://" + - std::string(bucket) + "/" + std::string(prefix) + - " — narrow the glob prefix"); - } - if (page.is_truncated && page.next_continuation_token.empty()) { - throw std::runtime_error( - "rest_ioctx::list_objects: truncated ListObjectsV2 page without a continuation token for " - "s3://" + - std::string(bucket) + "/" + std::string(prefix)); - } - // A truncated page must contain entries and advance the token. Together - // with scanned_cap these bound pagination for non-conforming backends that - // would otherwise loop on empty or non-advancing pages. - if (page.is_truncated && page.entries.empty()) { - throw std::runtime_error( - "rest_ioctx::list_objects: truncated ListObjectsV2 page with no entries for s3://" + - std::string(bucket) + "/" + std::string(prefix)); - } - if (page.is_truncated && page.next_continuation_token == token) { - throw std::runtime_error( - "rest_ioctx::list_objects: ListObjectsV2 continuation token did not advance for s3://" + - std::string(bucket) + "/" + std::string(prefix)); - } - truncated = page.is_truncated; - token = page.next_continuation_token; - if (!sink(page)) { return; } - } while (truncated); + _lister.list_objects_paged(bucket, prefix, page_size, sink, max_scanned); } std::vector rest_ioctx::list_objects(std::string_view bucket, @@ -133,25 +93,10 @@ std::vector rest_ioctx::list_objects(std::string_view bucket, std::size_t page_size, std::optional max_keys) { - std::size_t const keys_cap = max_keys.value_or(list_max_matches()); - std::vector out; - list_objects_paged(bucket, prefix, page_size, [&](s3::list_objects_v2_page const& page) { - if (out.size() + page.entries.size() > keys_cap) { - throw std::runtime_error("rest_ioctx::list_objects: more than " + std::to_string(keys_cap) + - " objects under s3://" + std::string(bucket) + "/" + - std::string(prefix) + " — narrow the glob prefix"); - } - out.insert(out.end(), page.entries.begin(), page.entries.end()); - return true; - }); - return out; + return _lister.list_objects(bucket, prefix, page_size, max_keys); } -std::size_t rest_ioctx::list_max_matches() const -{ - return _reactors.empty() ? s3::default_max_list_objects - : _reactors.front()->get_config().list_max_matches; -} +std::size_t rest_ioctx::list_max_matches() const { return _lister.list_max_matches(); } std::shared_ptr rest_ioctx::create_io_object(std::string path) { diff --git a/src/io/s3rdma/s3rdma_ioctx.cpp b/src/io/s3rdma/s3rdma_ioctx.cpp new file mode 100644 index 0000000..8afc36a --- /dev/null +++ b/src/io/s3rdma/s3rdma_ioctx.cpp @@ -0,0 +1,25 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Placeholder translation unit: compiles the s3rdma_ioctx declaration so CI +// verifies the header, and is replaced wholesale together with it when the +// S3-over-RDMA backend is contributed upstream (see the header's class doc). +// The class is deliberately non-constructible until then — no definitions +// here. + +#include diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index dd613aa..b57e63f 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -61,9 +61,12 @@ if(NOT CUCASCADE_TOPOLOGY_ONLY AND CUCASCADE_BUILD_IO) # IO test executable - links the cudf-free cucascade-io datasource layer. add_executable( cucascade_io_tests + io/test_datasource_registry.cpp + io/test_dispatch_failure_hook.cpp io/test_uri_parser.cpp io/cache/test_metadata_store.cpp io/kvikio/test_kvikio_config.cpp + io/rest/test_object_store_lister.cpp io/rest/test_rest_perf_snapshot.cpp io/rest/test_rest_validation_tag.cpp io/rest/test_shared_byte_span.cpp diff --git a/test/io/rest/test_object_store_lister.cpp b/test/io/rest/test_object_store_lister.cpp new file mode 100644 index 0000000..4da6f31 --- /dev/null +++ b/test/io/rest/test_object_store_lister.cpp @@ -0,0 +1,611 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using cucascade::io::rest::authorized_request; +using cucascade::io::rest::config; +using cucascade::io::rest::object_ref; +using cucascade::io::rest::object_store_lister; +using cucascade::io::rest::request_authorizer; +using cucascade::io::rest::request_method; +using cucascade::io::rest::rest_ioctx; +using cucascade::io::rest::rest_reactor; +using cucascade::io::rest::s3::list_objects_v2_page; +using namespace std::chrono_literals; + +struct listed_object { + std::string key; + std::uint64_t size; +}; + +struct scripted_page { + std::string request_token; + std::vector objects; + bool truncated{false}; + std::string next_token; +}; + +struct observed_query { + std::string max_keys; + std::string continuation_token; + std::string prefix; +}; + +class scripted_list_server { + public: + explicit scripted_list_server(std::vector pages) : _pages(std::move(pages)) + { + if (_pages.empty()) { throw std::invalid_argument("scripted LIST server needs a page"); } + + _listen_fd = ::socket(AF_INET, SOCK_STREAM, 0); + if (_listen_fd < 0) { throw std::runtime_error("socket failed: " + errno_message()); } + + int one = 1; + if (::setsockopt(_listen_fd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one)) != 0) { + close_listener(); + throw std::runtime_error("setsockopt failed: " + errno_message()); + } + + sockaddr_in address{}; + address.sin_family = AF_INET; + address.sin_addr.s_addr = htonl(INADDR_LOOPBACK); + address.sin_port = 0; + if (::bind(_listen_fd, reinterpret_cast(&address), sizeof(address)) != 0) { + close_listener(); + throw std::runtime_error("bind failed: " + errno_message()); + } + if (::listen(_listen_fd, 16) != 0) { + close_listener(); + throw std::runtime_error("listen failed: " + errno_message()); + } + + socklen_t length = sizeof(address); + if (::getsockname(_listen_fd, reinterpret_cast(&address), &length) != 0) { + close_listener(); + throw std::runtime_error("getsockname failed: " + errno_message()); + } + _port = ntohs(address.sin_port); + int const listen_fd = _listen_fd; + _thread = std::thread([this, listen_fd] { accept_loop(listen_fd); }); + } + + ~scripted_list_server() + { + _stop.store(true, std::memory_order_relaxed); + if (_listen_fd >= 0) { (void)::shutdown(_listen_fd, SHUT_RDWR); } + if (_thread.joinable()) { _thread.join(); } + close_listener(); + } + + scripted_list_server(scripted_list_server const&) = delete; + scripted_list_server& operator=(scripted_list_server const&) = delete; + + [[nodiscard]] std::string endpoint() const { return "http://127.0.0.1:" + std::to_string(_port); } + + [[nodiscard]] std::size_t request_count() const noexcept + { + return _request_count.load(std::memory_order_relaxed); + } + + [[nodiscard]] std::vector observations() const + { + std::scoped_lock lock{_observations_mutex}; + return _observations; + } + + private: + static std::string errno_message() { return std::strerror(errno); } + + void close_listener() noexcept + { + if (_listen_fd < 0) { return; } + (void)::close(_listen_fd); + _listen_fd = -1; + } + + void accept_loop(int listen_fd) + { + while (!_stop.load(std::memory_order_relaxed)) { + sockaddr_in client{}; + socklen_t length = sizeof(client); + int const fd = ::accept(listen_fd, reinterpret_cast(&client), &length); + if (fd < 0) { + if (_stop.load(std::memory_order_relaxed)) { return; } + continue; + } + handle_client(fd); + (void)::close(fd); + } + } + + void handle_client(int fd) + { + timeval timeout{}; + timeout.tv_sec = 5; + (void)::setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout)); + + std::string request; + std::array buffer{}; + while (request.find("\r\n\r\n") == std::string::npos && request.size() < 64 * 1024) { + ssize_t const received = ::recv(fd, buffer.data(), buffer.size(), 0); + if (received <= 0) { return; } + request.append(buffer.data(), static_cast(received)); + } + + std::string const target = request_target(request); + if (request.rfind("GET ", 0) != 0 || target.find("list-type=2") == std::string::npos) { + send_all(fd, + "HTTP/1.1 405 Method Not Allowed\r\nContent-Length: 0\r\nConnection: " + "close\r\n\r\n"); + return; + } + + observed_query observation{.max_keys = query_value(target, "max-keys"), + .continuation_token = query_value(target, "continuation-token"), + .prefix = query_value(target, "prefix")}; + { + std::scoped_lock lock{_observations_mutex}; + _observations.push_back(observation); + } + _request_count.fetch_add(1, std::memory_order_relaxed); + + auto const* page = find_page(observation.continuation_token); + if (page == nullptr) { + send_all(fd, "HTTP/1.1 400 Bad Request\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"); + return; + } + + std::string const body = page_xml(*page); + send_all(fd, + "HTTP/1.1 200 OK\r\nContent-Type: application/xml\r\nContent-Length: " + + std::to_string(body.size()) + "\r\nConnection: close\r\n\r\n" + body); + } + + [[nodiscard]] scripted_page const* find_page(std::string_view request_token) const noexcept + { + for (auto const& page : _pages) { + if (page.request_token == request_token) { return &page; } + } + return nullptr; + } + + static std::string request_target(std::string const& request) + { + auto const first_space = request.find(' '); + if (first_space == std::string::npos) { return {}; } + auto const second_space = request.find(' ', first_space + 1); + if (second_space == std::string::npos) { return {}; } + return request.substr(first_space + 1, second_space - first_space - 1); + } + + static int hex_value(char c) noexcept + { + if (c >= '0' && c <= '9') { return c - '0'; } + if (c >= 'a' && c <= 'f') { return c - 'a' + 10; } + if (c >= 'A' && c <= 'F') { return c - 'A' + 10; } + return -1; + } + + static std::string url_decode(std::string_view encoded) + { + std::string decoded; + decoded.reserve(encoded.size()); + for (std::size_t i = 0; i < encoded.size(); ++i) { + if (encoded[i] == '%' && i + 2 < encoded.size()) { + int const high = hex_value(encoded[i + 1]); + int const low = hex_value(encoded[i + 2]); + if (high >= 0 && low >= 0) { + decoded.push_back(static_cast((high << 4) | low)); + i += 2; + continue; + } + } + decoded.push_back(encoded[i] == '+' ? ' ' : encoded[i]); + } + return decoded; + } + + static std::string query_value(std::string_view target, std::string_view wanted_key) + { + auto const question = target.find('?'); + if (question == std::string_view::npos) { return {}; } + std::string_view query = target.substr(question + 1); + while (!query.empty()) { + auto const ampersand = query.find('&'); + auto const part = query.substr(0, ampersand); + auto const equals = part.find('='); + if (equals != std::string_view::npos && part.substr(0, equals) == wanted_key) { + return url_decode(part.substr(equals + 1)); + } + if (ampersand == std::string_view::npos) { break; } + query.remove_prefix(ampersand + 1); + } + return {}; + } + + static std::string xml_escape(std::string_view value) + { + std::string escaped; + for (char c : value) { + switch (c) { + case '&': escaped += "&"; break; + case '<': escaped += "<"; break; + case '>': escaped += ">"; break; + case '\"': escaped += """; break; + case '\'': escaped += "'"; break; + default: escaped.push_back(c); break; + } + } + return escaped; + } + + static std::string page_xml(scripted_page const& page) + { + std::string body = + "" + ""; + body += page.truncated ? "true" : "false"; + body += ""; + if (!page.next_token.empty()) { + body += "" + xml_escape(page.next_token) + ""; + } + for (auto const& object : page.objects) { + body += "" + xml_escape(object.key) + "" + + std::to_string(object.size) + ""; + } + body += ""; + return body; + } + + static void send_all(int fd, std::string_view response) + { + std::size_t sent = 0; + while (sent < response.size()) { + ssize_t const written = + ::send(fd, response.data() + sent, response.size() - sent, MSG_NOSIGNAL); + if (written <= 0) { return; } + sent += static_cast(written); + } + } + + int _listen_fd{-1}; + std::uint16_t _port{0}; + std::vector _pages; + std::atomic _stop{false}; + std::atomic _request_count{0}; + mutable std::mutex _observations_mutex; + std::vector _observations; + std::thread _thread; +}; + +class loopback_list_authorizer final : public request_authorizer { + public: + explicit loopback_list_authorizer(std::string endpoint) : _endpoint(std::move(endpoint)) {} + + authorized_request authorize(object_ref const& object, + request_method, + std::chrono::seconds) override + { + return {_endpoint + "/" + object.bucket + "/" + object.key, {}}; + } + + authorized_request authorize_list(std::string_view bucket, + std::string_view canonical_query, + std::chrono::seconds) override + { + return {_endpoint + "/" + std::string{bucket} + "?" + std::string{canonical_query}, {}}; + } + + private: + std::string _endpoint; +}; + +config listing_config(std::size_t list_max_matches = 100'000) +{ + config cfg{}; + cfg.request_timeout_s = 5; + cfg.tls_verify = false; + cfg.max_connections = 1; + cfg.max_retry_attempts = 1; + cfg.max_auth_retry_attempts = 1; + cfg.retry_backoff_base = 1ms; + cfg.retry_jitter = 0ms; + cfg.honor_retry_after = false; + cfg.list_max_matches = list_max_matches; + return cfg; +} + +std::string direct_page_xml(std::vector objects, + bool truncated, + std::optional next_token = std::nullopt) +{ + std::string body = ""; + body += truncated ? "true" : "false"; + body += ""; + if (next_token.has_value()) { + body += "" + std::string{*next_token} + ""; + } + for (auto const& object : objects) { + body += "" + object.key + "" + std::to_string(object.size) + + ""; + } + body += ""; + return body; +} + +class stub_page_fetch { + public: + explicit stub_page_fetch(std::vector responses) : _responses(std::move(responses)) {} + + std::string fetch(std::string_view, std::string_view, std::string_view) + { + if (_next == _responses.size()) { + throw std::runtime_error("stub page fetch exhausted its responses"); + } + return _responses[_next++]; + } + + private: + std::vector _responses; + std::size_t _next{0}; +}; + +object_store_lister make_direct_lister(std::shared_ptr fetch, + std::size_t max_scanned = 100, + std::size_t max_matches = 100) +{ + return object_store_lister{ + [fetch = std::move(fetch)]( + std::string_view bucket, std::string_view prefix, std::string_view canonical_query) { + return fetch->fetch(bucket, prefix, canonical_query); + }, + max_scanned, + max_matches, + "test_lister::list_objects"}; +} + +struct listing_fixture { + explicit listing_fixture(std::vector pages, std::size_t list_max_matches = 100'000) + : server(std::move(pages)), + authorizer(std::make_shared(server.endpoint())) + { + auto context = std::make_shared( + listing_config(list_max_matches), authorizer, nullptr); + ioctx = std::make_shared(1, std::move(context)); + ioctx->start(); + } + + scripted_list_server server; + std::shared_ptr authorizer; + std::shared_ptr ioctx; +}; + +} // namespace + +TEST_CASE("rest ioctx delegates paged listing to its composed lister", "[rest][listing]") +{ + listing_fixture fixture{ + {scripted_page{.request_token = "", + .objects = {{"prefix/a.parquet", 11}, {"prefix/b.parquet", 22}}, + .truncated = true, + .next_token = "page/2"}, + scripted_page{.request_token = "page/2", + .objects = {{"prefix/c.parquet", 33}}, + .truncated = false, + .next_token = ""}}}; + std::vector delivered; + + fixture.ioctx->list_objects_paged("bucket", "prefix/", 2, [&](list_objects_v2_page const& page) { + delivered.push_back(page); + return true; + }); + + REQUIRE(delivered.size() == 2); + REQUIRE(delivered[0].entries.size() == 2); + REQUIRE(delivered[1].entries.size() == 1); + CHECK(delivered[0].entries[0].key == "prefix/a.parquet"); + CHECK(delivered[0].entries[0].size == 11); + CHECK(delivered[0].entries[1].key == "prefix/b.parquet"); + CHECK(delivered[1].entries[0].key == "prefix/c.parquet"); + CHECK(delivered[1].entries[0].size == 33); + + auto const observations = fixture.server.observations(); + REQUIRE(observations.size() == 2); + CHECK(observations[0].max_keys == "2"); + CHECK(observations[0].prefix == "prefix/"); + CHECK(observations[0].continuation_token.empty()); + CHECK(observations[1].continuation_token == "page/2"); + CHECK(observations[1].prefix == "prefix/"); + + SECTION("whole-list delegation preserves order and enforces max_keys") + { + auto const objects = fixture.ioctx->list_objects("bucket", "prefix/", 2); + + REQUIRE(objects.size() == 3); + CHECK(objects[0].key == "prefix/a.parquet"); + CHECK(objects[0].size == 11); + CHECK(objects[1].key == "prefix/b.parquet"); + CHECK(objects[1].size == 22); + CHECK(objects[2].key == "prefix/c.parquet"); + CHECK(objects[2].size == 33); + + CHECK_THROWS_WITH(fixture.ioctx->list_objects("bucket", "prefix/", 2, 2), + Catch::Matchers::ContainsSubstring("rest_ioctx::list_objects:") && + Catch::Matchers::ContainsSubstring("more than 2 objects")); + } +} + +TEST_CASE("a listing sink can stop before the next page request", "[rest][listing]") +{ + listing_fixture fixture{ + {scripted_page{ + .request_token = "", .objects = {{"prefix/a", 1}}, .truncated = true, .next_token = "next"}, + scripted_page{.request_token = "next", + .objects = {{"prefix/b", 2}}, + .truncated = false, + .next_token = ""}}}; + std::size_t pages_seen = 0; + + fixture.ioctx->list_objects_paged("bucket", "prefix/", 1, [&](list_objects_v2_page const&) { + ++pages_seen; + return false; + }); + + CHECK(pages_seen == 1); + CHECK(fixture.server.request_count() == 1); +} + +TEST_CASE("listing page size is clamped on the wire", "[rest][listing]") +{ + listing_fixture fixture{{scripted_page{ + .request_token = "", .objects = {{"key", 1}}, .truncated = false, .next_token = ""}}}; + auto const consume = [](list_objects_v2_page const&) { return true; }; + + fixture.ioctx->list_objects_paged("bucket", "", 0, consume); + fixture.ioctx->list_objects_paged("bucket", "", 1001, consume); + + auto const observations = fixture.server.observations(); + REQUIRE(observations.size() == 2); + CHECK(observations[0].max_keys == "1000"); + CHECK(observations[1].max_keys == "1000"); +} + +TEST_CASE("listing throws when the scanned object cap is exceeded", "[rest][listing]") +{ + listing_fixture fixture{{scripted_page{.request_token = "", + .objects = {{"prefix/a", 1}, {"prefix/b", 2}}, + .truncated = false, + .next_token = ""}}}; + std::size_t pages_seen = 0; + + CHECK_THROWS_WITH(fixture.ioctx->list_objects_paged( + "bucket", + "prefix/", + 1000, + [&](list_objects_v2_page const&) { + ++pages_seen; + return true; + }, + 1), + Catch::Matchers::ContainsSubstring("rest_ioctx::list_objects:") && + Catch::Matchers::ContainsSubstring("scanned more than 1 objects")); + CHECK(pages_seen == 0); + CHECK(fixture.server.request_count() == 1); +} + +TEST_CASE("rest ioctx exposes the configured listing match cap", "[rest][listing]") +{ + constexpr std::size_t configured_cap = 37; + listing_fixture fixture{ + {scripted_page{.request_token = "", .objects = {}, .truncated = false, .next_token = ""}}, + configured_cap}; + + CHECK(fixture.ioctx->list_max_matches() == configured_cap); + CHECK(fixture.server.request_count() == 0); +} + +TEST_CASE("object store lister rejects unsafe pagination and result growth", "[rest][listing]") +{ + auto const consume = [](list_objects_v2_page const&) { return true; }; + + SECTION("a truncated page requires a non-empty continuation token") + { + // A missing element is rejected by the parser; an empty element reaches the lister guard. + auto fetch = std::make_shared(std::vector{direct_page_xml( + {{"prefix/a", 1}}, true, std::optional{std::string_view{}})}); + auto lister = make_direct_lister(std::move(fetch)); + + CHECK_THROWS_WITH(lister.list_objects_paged("bucket", "prefix/", 1000, consume), + Catch::Matchers::ContainsSubstring("test_lister::list_objects:") && + Catch::Matchers::ContainsSubstring("without a continuation token")); + } + + SECTION("a truncated page cannot be empty") + { + auto fetch = std::make_shared( + std::vector{direct_page_xml({}, true, std::string_view{"next"})}); + auto lister = make_direct_lister(std::move(fetch)); + + CHECK_THROWS_WITH(lister.list_objects_paged("bucket", "prefix/", 1000, consume), + Catch::Matchers::ContainsSubstring("test_lister::list_objects:") && + Catch::Matchers::ContainsSubstring("with no entries")); + } + + SECTION("a continuation token must advance") + { + auto fetch = std::make_shared( + std::vector{direct_page_xml({{"prefix/a", 1}}, true, std::string_view{"next"}), + direct_page_xml({{"prefix/b", 2}}, true, std::string_view{"next"})}); + auto lister = make_direct_lister(std::move(fetch)); + + CHECK_THROWS_WITH(lister.list_objects_paged("bucket", "prefix/", 1000, consume), + Catch::Matchers::ContainsSubstring("test_lister::list_objects:") && + Catch::Matchers::ContainsSubstring("continuation token did not advance")); + } + + SECTION("the scanned object cap is enforced") + { + auto fetch = std::make_shared( + std::vector{direct_page_xml({{"prefix/a", 1}, {"prefix/b", 2}}, false)}); + auto lister = make_direct_lister(std::move(fetch), 1); + + CHECK_THROWS_WITH(lister.list_objects_paged("bucket", "prefix/", 1000, consume), + Catch::Matchers::ContainsSubstring("test_lister::list_objects:") && + Catch::Matchers::ContainsSubstring("scanned more than 1 objects")); + } + + SECTION("the whole-list match cap is enforced") + { + auto fetch = std::make_shared( + std::vector{direct_page_xml({{"prefix/a", 1}, {"prefix/b", 2}}, false)}); + auto lister = make_direct_lister(std::move(fetch), 100, 1); + + CHECK_THROWS_WITH(lister.list_objects("bucket", "prefix/"), + Catch::Matchers::ContainsSubstring("test_lister::list_objects:") && + Catch::Matchers::ContainsSubstring("more than 1 objects")); + } +} diff --git a/test/io/test_datasource_registry.cpp b/test/io/test_datasource_registry.cpp new file mode 100644 index 0000000..5464993 --- /dev/null +++ b/test/io/test_datasource_registry.cpp @@ -0,0 +1,149 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include + +#include + +#include +#include +#include +#include + +namespace { + +using cucascade::io::io_config; +using cucascade::io::io_context_registry; +using cucascade::io::io_context_type; +using cucascade::io::ioctx; +using cucascade::memory::disk_memory_space_config; +using cucascade::memory::memory_reservation_manager; +using cucascade::memory::memory_space_config; + +class registry_fixture { + public: + registry_fixture() + : manager(std::vector{disk_memory_space_config{ + .disk_id = 0, .memory_capacity = 1UL << 20, .mount_paths = "/tmp"}}), + registry(io_config{}, manager) + { + } + + memory_reservation_manager manager; + io_context_registry registry; +}; + +bool s3_checker(std::string_view path) { return path.starts_with("s3://"); } + +bool rdma_checker(std::string_view path) { return path.starts_with("rdma://"); } + +std::shared_ptr null_factory(io_config const&) { return nullptr; } + +} // namespace + +TEST_CASE("replace hands the s3 claimant to the new backend", "[io][registry]") +{ + std::size_t old_factory_calls = 0; + registry_fixture fixture; + + fixture.registry.register_ioctx(io_context_type::restful, + &s3_checker, + [&old_factory_calls](io_config const&) -> std::shared_ptr { + ++old_factory_calls; + return nullptr; + }); + + fixture.registry.replace_ioctx( + io_context_type::restful, io_context_type::s3rdma, &s3_checker, &null_factory); + + CHECK(fixture.registry.lookup_path("s3://bucket/key") == io_context_type::s3rdma); + CHECK(fixture.registry.make_ioctx(io_context_type::restful) == nullptr); + CHECK(old_factory_calls == 0); + CHECK(fixture.registry.lookup_path("/proc/self/exe") == io_context_type::uring); +} + +TEST_CASE("replace rejects a missing old backend without changing routing", "[io][registry]") +{ + registry_fixture fixture; + + CHECK_THROWS_AS(fixture.registry.replace_ioctx( + io_context_type::s3rdma, io_context_type::s3rdma, &s3_checker, &null_factory), + std::invalid_argument); + CHECK(fixture.registry.lookup_path("s3://bucket/key") == io_context_type::restful); +} + +TEST_CASE("replace rejects an already registered new backend", "[io][registry]") +{ + registry_fixture fixture; + + CHECK_THROWS_AS(fixture.registry.replace_ioctx( + io_context_type::restful, io_context_type::kvikio, &s3_checker, &null_factory), + std::invalid_argument); + CHECK(fixture.registry.lookup_path("s3://bucket/key") == io_context_type::restful); +} + +TEST_CASE("replace rejects a null checker without changing routing", "[io][registry]") +{ + registry_fixture fixture; + + CHECK_THROWS_AS(fixture.registry.replace_ioctx(io_context_type::restful, + io_context_type::s3rdma, + io_context_registry::scheme_checker_type{}, + &null_factory), + std::invalid_argument); + CHECK(fixture.registry.lookup_path("s3://bucket/key") == io_context_type::restful); +} + +TEST_CASE("replace rejects a null factory without changing routing", "[io][registry]") +{ + registry_fixture fixture; + + CHECK_THROWS_AS(fixture.registry.replace_ioctx(io_context_type::restful, + io_context_type::s3rdma, + &s3_checker, + io_context_registry::factory_type{}), + std::invalid_argument); + CHECK(fixture.registry.lookup_path("s3://bucket/key") == io_context_type::restful); +} + +TEST_CASE("replace is forbidden after the first path lookup", "[io][registry]") +{ + registry_fixture fixture; + + REQUIRE(fixture.registry.lookup_path("unmatched-before-bootstrap") == io_context_type::kvikio); + CHECK_THROWS_AS(fixture.registry.replace_ioctx( + io_context_type::restful, io_context_type::s3rdma, &s3_checker, &null_factory), + std::logic_error); + CHECK(fixture.registry.lookup_path("s3://bucket/key") == io_context_type::restful); +} + +TEST_CASE("register remains legal after path lookup", "[io][registry]") +{ + registry_fixture fixture; + + REQUIRE(fixture.registry.lookup_path("s3://bucket/key") == io_context_type::restful); + CHECK_NOTHROW( + fixture.registry.register_ioctx(io_context_type::s3rdma, &rdma_checker, &null_factory)); + CHECK(fixture.registry.lookup_path("rdma://bucket/key") == io_context_type::s3rdma); +} + +TEST_CASE("s3 rdma has a distinct context type", "[io][registry]") +{ + CHECK(io_context_type::s3rdma != io_context_type::restful); +} diff --git a/test/io/test_dispatch_failure_hook.cpp b/test/io/test_dispatch_failure_hook.cpp new file mode 100644 index 0000000..6a38aa4 --- /dev/null +++ b/test/io/test_dispatch_failure_hook.cpp @@ -0,0 +1,365 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +struct dispatch_controls { + bool throw_device_prep{false}; + bool throw_staged_prep{false}; + bool throw_enqueue{false}; +}; + +class stub_io_object final : public cucascade::io::io_object { + public: + explicit stub_io_object(std::shared_ptr controls, + std::string path = "stub://object", + std::size_t size = 64) + : _controls(std::move(controls)), _path(std::move(path)), _size(size) + { + } + + [[nodiscard]] std::shared_ptr const& controls() const noexcept + { + return _controls; + } + + [[nodiscard]] const std::string& raw_file_cache_id() const noexcept override { return _path; } + [[nodiscard]] const std::string& object_path() const noexcept override { return _path; } + [[nodiscard]] std::size_t size() const noexcept override { return _size; } + + private: + std::shared_ptr _controls; + std::string _path; + std::size_t _size; +}; + +class stub_request { + public: + stub_request(std::size_t bytes, std::shared_ptr controls) + : _state(std::make_shared(bytes, std::move(controls))) + { + } + + [[nodiscard]] cucascade::exec::semi_future get_future() noexcept + { + return _state->promise.get_semi_future(); + } + + static std::vector> splits(std::unique_ptr request, + std::size_t n_splits) noexcept + { + std::vector> result; + if (request != nullptr && n_splits != 0) { result.push_back(std::move(request)); } + return result; + } + + [[nodiscard]] dispatch_controls const& controls() const noexcept { return *_state->controls; } + + void complete() { _state->promise.set_value(std::size_t{_state->bytes}); } + + private: + struct state { + state(std::size_t bytes, std::shared_ptr controls) + : bytes(bytes), controls(std::move(controls)) + { + } + + std::size_t bytes; + std::shared_ptr controls; + cucascade::exec::promise promise; + }; + + std::shared_ptr _state; +}; + +struct stub_reactor_config {}; + +class stub_reactor { + public: + using io_object_type = stub_io_object; + using request_type = stub_request; + using request_type_ptr = std::unique_ptr; + using reactor_config_type = stub_reactor_config; + + [[nodiscard]] const reactor_config_type& get_config() const noexcept { return _config; } + + static request_type_ptr prep_host_rx_request(const reactor_config_type&, + const io_object_type& file, + cucascade::io::io_object_segment segment) + { + return std::make_unique(segment.size, file.controls()); + } + + static request_type_ptr prep_device_rx_request(const reactor_config_type&, + const io_object_type& file, + std::uint8_t*, + std::size_t, + std::size_t size, + rmm::cuda_stream_view, + int) + { + if (file.controls()->throw_device_prep) { throw std::runtime_error("device prep failure"); } + return std::make_unique(size, file.controls()); + } + + static request_type_ptr prep_host_to_device_rx_request( + const reactor_config_type&, + const io_object_type& file, + std::span, + std::uint8_t*, + std::size_t, + std::size_t size, + rmm::cuda_stream_view, + int) + { + if (file.controls()->throw_staged_prep) { + throw std::runtime_error("host-to-device prep failure"); + } + return std::make_unique(size, file.controls()); + } + + void enqueue(request_type_ptr request) + { + if (request->controls().throw_enqueue) { throw std::runtime_error("enqueue failure"); } + request->complete(); + } + + std::size_t host_read(const io_object_type&, std::size_t, std::size_t size, std::uint8_t*) + { + return size; + } + + void start() {} + void shutdown() {} + void interrupt() {} + + static std::unique_ptr create_io_object(std::string path) + { + return std::make_unique(std::make_shared(), std::move(path)); + } + + [[nodiscard]] static bool supports(std::string_view) { return true; } + + [[nodiscard]] static constexpr cucascade::io::cache::prefetching_stage + preferred_prefetching_stage() noexcept + { + return cucascade::io::cache::prefetching_stage::none; + } + + private: + reactor_config_type _config; +}; + +static_assert(cucascade::io::io_reactor_c); +static_assert(cucascade::io::reactor_has_device_rx); +static_assert(cucascade::io::reactor_has_host_to_device_rx); + +std::vector> make_reactors() +{ + std::vector> reactors; + reactors.push_back(std::make_unique()); + return reactors; +} + +class hooked_ioctx final : public cucascade::io::templated_ioctx { + public: + explicit hooked_ioctx(bool empty_selection = false) + : templated_ioctx(make_reactors()), _empty_selection(empty_selection) + { + } + + [[nodiscard]] cucascade::io::io_context_type type() const noexcept override + { + return cucascade::io::io_context_type::s3rdma; + } + + [[nodiscard]] std::size_t hook_calls() const noexcept { return _hook_calls; } + + std::vector next_reactor(const stub_io_object& object, + std::size_t n_chunks, + io_op_type operation, + int device_id = -1) noexcept override + { + if (_empty_selection) { return {}; } + return templated_ioctx::next_reactor(object, n_chunks, operation, device_id); + } + + protected: + void on_device_dispatch_failure() noexcept override { ++_hook_calls; } + + private: + bool _empty_selection; + std::size_t _hook_calls{0}; +}; + +class plain_ioctx final : public cucascade::io::templated_ioctx { + public: + plain_ioctx() : templated_ioctx(make_reactors()) {} + + [[nodiscard]] cucascade::io::io_context_type type() const noexcept override + { + return cucascade::io::io_context_type::kvikio; + } +}; + +std::shared_ptr make_object(std::shared_ptr controls) +{ + return std::make_shared(std::move(controls)); +} + +void check_error(cucascade::exec::semi_future future, std::string_view message) +{ + CHECK_THROWS_WITH(std::move(future).get(), + Catch::Matchers::ContainsSubstring(std::string{message})); +} + +} // namespace + +TEST_CASE("device prep failure fires the dispatch hook once", "[io][hook]") +{ + auto controls = std::make_shared(); + controls->throw_device_prep = true; + auto object = make_object(controls); + hooked_ioctx ioctx; + + auto future = + ioctx.device_read_async_io(*object, 0, object->size(), nullptr, rmm::cuda_stream_default); + + check_error(std::move(future), "device prep failure"); + CHECK(ioctx.hook_calls() == 1); +} + +TEST_CASE("device enqueue failure fires the dispatch hook once", "[io][hook]") +{ + auto controls = std::make_shared(); + controls->throw_enqueue = true; + auto object = make_object(controls); + hooked_ioctx ioctx; + + auto future = + ioctx.device_read_async_io(*object, 0, object->size(), nullptr, rmm::cuda_stream_default); + + check_error(std::move(future), "enqueue failure"); + CHECK(ioctx.hook_calls() == 1); +} + +TEST_CASE("host to device prep failure fires the dispatch hook once", "[io][hook]") +{ + auto controls = std::make_shared(); + controls->throw_staged_prep = true; + auto object = make_object(controls); + std::array bounce{}; + hooked_ioctx ioctx; + + auto future = ioctx.host_to_device_read_async_io( + *object, bounce, 0, object->size(), nullptr, rmm::cuda_stream_default); + + check_error(std::move(future), "host-to-device prep failure"); + CHECK(ioctx.hook_calls() == 1); +} + +TEST_CASE("host to device enqueue failure fires the dispatch hook once", "[io][hook]") +{ + auto controls = std::make_shared(); + controls->throw_enqueue = true; + auto object = make_object(controls); + std::array bounce{}; + hooked_ioctx ioctx; + + auto future = ioctx.host_to_device_read_async_io( + *object, bounce, 0, object->size(), nullptr, rmm::cuda_stream_default); + + check_error(std::move(future), "enqueue failure"); + CHECK(ioctx.hook_calls() == 1); +} + +TEST_CASE("empty reactor selection returns errors without firing the hook", "[io][hook]") +{ + auto controls = std::make_shared(); + auto object = make_object(controls); + hooked_ioctx ioctx{true}; + + SECTION("device read") + { + auto future = + ioctx.device_read_async_io(*object, 0, object->size(), nullptr, rmm::cuda_stream_default); + check_error(std::move(future), "device_read_async_io: no available reactors"); + CHECK(ioctx.hook_calls() == 0); + } + + SECTION("host to device read") + { + std::array bounce{}; + auto future = ioctx.host_to_device_read_async_io( + *object, bounce, 0, object->size(), nullptr, rmm::cuda_stream_default); + check_error(std::move(future), "host_to_device_read_async_io: no available reactors"); + CHECK(ioctx.hook_calls() == 0); + } +} + +TEST_CASE("successful device dispatches do not fire the hook", "[io][hook]") +{ + auto controls = std::make_shared(); + auto object = make_object(controls); + hooked_ioctx ioctx; + + SECTION("device read") + { + auto future = + ioctx.device_read_async_io(*object, 0, object->size(), nullptr, rmm::cuda_stream_default); + CHECK(std::move(future).get() == object->size()); + CHECK(ioctx.hook_calls() == 0); + } + + SECTION("host to device read") + { + std::array bounce{}; + auto future = ioctx.host_to_device_read_async_io( + *object, bounce, 0, object->size(), nullptr, rmm::cuda_stream_default); + CHECK(std::move(future).get() == object->size()); + CHECK(ioctx.hook_calls() == 0); + } +} + +TEST_CASE("the default dispatch failure hook preserves error futures", "[io][hook]") +{ + auto controls = std::make_shared(); + controls->throw_device_prep = true; + auto object = make_object(controls); + plain_ioctx ioctx; + + auto future = + ioctx.device_read_async_io(*object, 0, object->size(), nullptr, rmm::cuda_stream_default); + + check_error(std::move(future), "device prep failure"); +}