diff --git a/include/cucascade/io/datasource_factory.hpp b/include/cucascade/io/datasource_factory.hpp index a73ec2a..03ff4d8 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,27 @@ class io_context_registry { */ void register_ioctx(io_context_type type, scheme_checker_type checker, factory_type factory); + /** + * @brief Atomically hand @p old_type's registration to @p new_type — + * bootstrap-only arbitration (e.g. an engine swapping the s3:// + * claimant for an alternative transport before any routing). + * + * Unlike an unregister+register pair, there is no observable no-claimant + * gap. + * + * @throws std::invalid_argument when @p old_type is not registered, when + * @p new_type is already registered, or when @p checker / + * @p factory is null. Strong guarantee: the registry is + * unchanged on any throw. + * @throws std::logic_error once @c lookup_path has run — the registry + * latches its first lookup; arbitration is legal strictly before + * routing begins (bootstrap is single-threaded by contract). + */ + 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 +127,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/object_store_listing.hpp b/include/cucascade/io/object_store_listing.hpp new file mode 100644 index 0000000..a89d275 --- /dev/null +++ b/include/cucascade/io/object_store_listing.hpp @@ -0,0 +1,57 @@ +/* + * 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 + +namespace cucascade::io { + +/// Listing capability of an object-store backend. A glob / LIST layer +/// depends on this interface, not a concrete ioctx type. Listing is prefix +/// resolution on the control plane; it is independent of how the data plane +/// reads the resolved keys. +class object_store_listing { + public: + virtual ~object_store_listing() = default; + + /// 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. + virtual 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) = 0; + + /// The backend's configured matched cap for glob resolution. + [[nodiscard]] virtual std::size_t list_max_matches() const = 0; +}; + +} // namespace cucascade::io diff --git a/include/cucascade/io/rest/rest_ioctx.hpp b/include/cucascade/io/rest/rest_ioctx.hpp index 9701570..559688a 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 @@ -46,7 +47,7 @@ namespace cucascade::io::rest { * via a blocking HEAD before constructing the @c rest_io_object — the static * reactor factory cannot do this since it needs the authorizer + a round-trip. */ -class rest_ioctx : public templated_ioctx { +class rest_ioctx : public templated_ioctx, public object_store_listing { public: /// Build a pool of @p n_reactors reactors, all sharing @p ctx (one context per /// pool: it carries the per-reactor @c config, the presigning authorizer, and @@ -74,7 +75,7 @@ class rest_ioctx : public templated_ioctx { std::string_view prefix, std::size_t page_size, std::function const& sink, - std::optional max_scanned = std::nullopt); + std::optional max_scanned = std::nullopt) override; /// Whole-listing convenience over @c list_objects_paged: every object under /// @p prefix, in document order, with sizes. Throws (never truncates) when @@ -90,7 +91,7 @@ class rest_ioctx : public templated_ioctx { /// glob layer one level up can bound its match set without a reactor handle. /// Falls back to the built-in default when the pool is empty (never in /// practice). - [[nodiscard]] std::size_t list_max_matches() const; + [[nodiscard]] std::size_t list_max_matches() const override; protected: /// Backend hook invoked by @c ioctx::open_io_object: parse @p path diff --git a/include/cucascade/io/s3rdma/s3rdma_ioctx.hpp b/include/cucascade/io/s3rdma/s3rdma_ioctx.hpp new file mode 100644 index 0000000..42e9b42 --- /dev/null +++ b/include/cucascade/io/s3rdma/s3rdma_ioctx.hpp @@ -0,0 +1,113 @@ +/* + * 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 S3-over-RDMA object-store ioctx — PLACEHOLDER. + * + * The backend (an alternative `s3://` data plane that lands object reads + * directly in device memory over RDMA, with LIST/HEAD staying on the HTTP + * control plane) is under active development downstream, in Sirius, against + * the extension points this framework already carries: it will identify as + * @c io_context_type::s3rdma, take over the `s3://` claim via + * @c io_context_registry::replace_ioctx, expose listing through + * @c object_store_listing, and probe context health through + * @c templated_ioctx::on_device_dispatch_failure. + * + * This declaration reserves the backend's name and surface; it will be + * replaced wholesale by the full implementation when the backend is + * contributed upstream. Until then the class is deliberately + * non-constructible — the constructor is deleted (pinned by the + * static_assert below) and the overrides are declared but not defined — + * so it claims no paths and cannot disturb routing. + */ +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..514f80a 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,19 @@ class templated_ioctx : public ioctx { } } + protected: + /// Backend policy point, called from the catch of every device-plane + /// dispatch (@c device_read_async_io / @c host_to_device_read_async_io) + /// BEFORE the failure is softened into an errored future — exactly once + /// per failed dispatch, never on a success path and never on the + /// empty-reactor-pool return (no exception occurred there). A backend + /// whose contract makes a poisoned CUDA context process-fatal probes + /// context health here. Must not throw (the dispatch wrappers are + /// noexcept) and must not re-enter this ioctx. The default keeps the + /// plain error-future behavior for every other backend. + 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 +431,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..9b4e8dd 100644 --- a/src/io/CMakeLists.txt +++ b/src/io/CMakeLists.txt @@ -23,6 +23,7 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/rest/curl_handle.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/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..5562997 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_listing.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_listing.cpp b/test/io/rest/test_object_store_listing.cpp new file mode 100644 index 0000000..d8a8285 --- /dev/null +++ b/test/io/rest/test_object_store_listing.cpp @@ -0,0 +1,488 @@ +/* + * 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 +#include + +namespace { + +using cucascade::io::object_store_listing; +using cucascade::io::rest::authorized_request; +using cucascade::io::rest::config; +using cucascade::io::rest::object_ref; +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; +} + +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(); + } + + [[nodiscard]] object_store_listing& listing() const { return *ioctx; } + + scripted_list_server server; + std::shared_ptr authorizer; + std::shared_ptr ioctx; +}; + +} // namespace + +TEST_CASE("rest listing is reachable through the object store interface", "[rest][listing]") +{ + static_assert(std::derived_from); + + 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.listing().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/"); +} + +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.listing().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.listing().list_objects_paged("bucket", "", 0, consume); + fixture.listing().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.listing().list_objects_paged( + "bucket", + "prefix/", + 1000, + [&](list_objects_v2_page const&) { + ++pages_seen; + return true; + }, + 1), + Catch::Matchers::ContainsSubstring("scanned more than 1 objects")); + CHECK(pages_seen == 0); + CHECK(fixture.server.request_count() == 1); +} + +TEST_CASE("listing exposes the configured match cap through the interface", "[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.listing().list_max_matches() == configured_cap); + CHECK(fixture.server.request_count() == 0); +} 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"); +}