From 4edecd9ff4d7fafb833a68a0a86f2c0c95b46106 Mon Sep 17 00:00:00 2001 From: Felipe Aramburu Date: Fri, 14 Aug 2026 17:54:51 -0500 Subject: [PATCH 1/5] fix(data): lock and shared_ptr-ify data_repository_manager::get_repository MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit get_repository() was the only method in data_repository_manager that did not take _mutex, and it returned a reference INTO the map. A concurrent add_new_repository or clear_all_repositories therefore raced both the lookup and the returned reference — with concurrent queries creating and clearing repositories this is a routine use-after-free (silent in release builds). - Store repositories as shared_ptr (add_new_repository still accepts unique_ptr; the manager remains the one logical owner). - New locked accessor get_repository_shared() returns a shared_ptr copy under _mutex, so a repository obtained just before a concurrent clear_all_repositories remains valid for the caller's use. - get_repository() is kept temporarily for source compatibility (deprecated; forwards to get_repository_shared). - Regression tests: a deterministic obtained-before-clear lifetime test and a multi-thread get-vs-add/clear hammer. Co-Authored-By: Claude Fable 5 --- .../data/data_repository_manager.hpp | 46 ++++-- test/data/test_data_repository_manager.cpp | 154 ++++++++++++++---- 2 files changed, 158 insertions(+), 42 deletions(-) diff --git a/include/cucascade/data/data_repository_manager.hpp b/include/cucascade/data/data_repository_manager.hpp index 0436256..00d6c29 100644 --- a/include/cucascade/data/data_repository_manager.hpp +++ b/include/cucascade/data/data_repository_manager.hpp @@ -110,12 +110,15 @@ class data_repository_manager { std::string_view port_id, std::unique_ptr repository) { - std::unique_ptr old_repository; + // Stored as shared_ptr so accessors can hand out lifetime-safe references + // under _mutex (see get_repository_shared); callers keep passing + // unique_ptr because each repository still has exactly one logical owner. + std::shared_ptr shared_repository{std::move(repository)}; { std::lock_guard lock(_mutex); auto it = _repositories.find({operator_id, std::string(port_id)}); if (it != _repositories.end()) { throw std::runtime_error("Repository already exists"); } - _repositories[{operator_id, std::string(port_id)}] = std::move(repository); + _repositories[{operator_id, std::string(port_id)}] = std::move(shared_repository); } } @@ -139,24 +142,42 @@ class data_repository_manager { } /** - * @brief Get direct access to a repository for advanced operations. + * @brief Get lifetime-safe access to a repository for advanced operations. * - * Provides direct access to the underlying repository implementation, allowing - * for repository-specific operations that aren't covered by the common interface. + * Looks the repository up under _mutex and returns a shared_ptr copy, so the + * returned repository stays valid even if a concurrent add_new_repository or + * clear_all_repositories mutates the map after this call returns. (The old + * variant returned a reference into the map without taking _mutex — a + * concurrent mutation raced both the lookup and the returned reference.) * * @param operator_id The unique ID of the operator whose repository is requested * @param port_id The port identifier for the repository - * @return std::unique_ptr& Reference to the repository + * @return std::shared_ptr Shared ownership of the repository * * @throws std::out_of_range If no repository exists for the specified operator/port - * @note Thread-safe for read access, but modifications should use the repository's own thread - * safety + * @note Thread-safe — the lookup holds the manager mutex; the repository's own + * thread safety covers subsequent operations on it */ - std::unique_ptr& get_repository(size_t operator_id, std::string_view port_id) + std::shared_ptr get_repository_shared(size_t operator_id, + std::string_view port_id) { + std::lock_guard lock(_mutex); return _repositories.at({operator_id, std::string(port_id)}); } + /** + * @brief DEPRECATED — use get_repository_shared(). + * + * Kept temporarily for source compatibility (callers used `.get()` / `->`, + * which keep compiling against the returned shared_ptr). Now forwards to the + * locked, lifetime-safe accessor; the historical signature returned a bare + * reference into the map without taking _mutex. + */ + std::shared_ptr get_repository(size_t operator_id, std::string_view port_id) + { + return get_repository_shared(operator_id, port_id); + } + /** * @brief Generate a globally unique data batch identifier. * @@ -244,8 +265,11 @@ class data_repository_manager { std::mutex _mutex; ///< Mutex for thread-safe access std::atomic _next_data_batch_id = 0; ///< Atomic counter for generating unique data batch identifiers - std::map> - _repositories; ///< Map of operator ID to data_repository + /// Map of operator ID/port to data_repository. Held by shared_ptr so + /// get_repository_shared can hand out references that survive a concurrent + /// clear_all_repositories / add_new_repository (the manager remains the one + /// logical owner; accessors only extend lifetime across their use). + std::map> _repositories; }; using shared_data_repository_manager = data_repository_manager; diff --git a/test/data/test_data_repository_manager.cpp b/test/data/test_data_repository_manager.cpp index 44de8a7..e1baf01 100644 --- a/test/data/test_data_repository_manager.cpp +++ b/test/data/test_data_repository_manager.cpp @@ -47,7 +47,7 @@ TEST_CASE("data_repository_manager Construction", "[data_repository_manager]") // Manager should be empty initially // Accessing non-existent repository should throw - REQUIRE_THROWS_AS(manager.get_repository(0, "default"), std::out_of_range); + REQUIRE_THROWS_AS(manager.get_repository_shared(0, "default"), std::out_of_range); } // Test adding a single repository @@ -60,7 +60,7 @@ TEST_CASE("data_repository_manager Add Single Repository", "[data_repository_man manager.add_new_repository(operator_id, "default", std::move(repository)); // Repository should be accessible - auto& repo = manager.get_repository(operator_id, "default"); + auto repo = manager.get_repository_shared(operator_id, "default"); REQUIRE(repo != nullptr); } @@ -79,7 +79,7 @@ TEST_CASE("data_repository_manager Add Multiple Repositories", "[data_repository // All repositories should be accessible for (size_t i = 0; i < num_operators; ++i) { - auto& repo = manager.get_repository(i, "default"); + auto repo = manager.get_repository_shared(i, "default"); REQUIRE(repo != nullptr); } } @@ -133,7 +133,7 @@ TEST_CASE("data_repository_manager Add Data Batch Single Operator", "[data_repos manager.add_data_batch(batch, operator_ports); // Repository should have the batch - auto& repo = manager.get_repository(operator_id, "default"); + auto repo = manager.get_repository_shared(operator_id, "default"); auto pulled_batch = repo->pop_next_data_batch(); REQUIRE(pulled_batch != nullptr); REQUIRE(pulled_batch->get_batch_id() == batch_id); @@ -163,7 +163,7 @@ TEST_CASE("data_repository_manager Add Data Batch Multiple Operators", "[data_re // All repositories should have the batch (same shared_ptr) for (size_t id : operator_ids) { - auto& repo = manager.get_repository(id, "default"); + auto repo = manager.get_repository_shared(id, "default"); auto pulled = repo->pop_next_data_batch(); REQUIRE(pulled != nullptr); REQUIRE(pulled->get_batch_id() == batch_id); @@ -241,8 +241,8 @@ TEST_CASE("data_repository_manager Thread-Safe Add Batch", "[data_repository_man } // Repository should have all batches - auto& repo = manager.get_repository(operator_id, "default"); - int count = 0; + auto repo = manager.get_repository_shared(operator_id, "default"); + int count = 0; while (true) { auto batch = repo->pop_next_data_batch(); if (!batch) break; @@ -300,8 +300,8 @@ TEST_CASE("data_repository_manager Full Workflow", "[data_repository_manager]") // Verify: Operator 0 should have 2 batches (batch 0 and 1) { - auto& repo = manager.get_repository(0, "default"); - int count = 0; + auto repo = manager.get_repository_shared(0, "default"); + int count = 0; while (true) { auto batch = repo->pop_next_data_batch(); if (!batch) break; @@ -312,8 +312,8 @@ TEST_CASE("data_repository_manager Full Workflow", "[data_repository_manager]") // Verify: Operator 1 should have 2 batches (batch 0 and 2) { - auto& repo = manager.get_repository(1, "default"); - int count = 0; + auto repo = manager.get_repository_shared(1, "default"); + int count = 0; while (true) { auto batch = repo->pop_next_data_batch(); if (!batch) break; @@ -324,8 +324,8 @@ TEST_CASE("data_repository_manager Full Workflow", "[data_repository_manager]") // Verify: Operator 2 should have 2 batches (batch 0 and 2) { - auto& repo = manager.get_repository(2, "default"); - int count = 0; + auto repo = manager.get_repository_shared(2, "default"); + int count = 0; while (true) { auto batch = repo->pop_next_data_batch(); if (!batch) break; @@ -349,7 +349,7 @@ TEST_CASE("data_repository_manager Large Number of Operators", "[data_repository // All operators should be accessible for (int i = 0; i < num_operators; ++i) { - auto& repo = manager.get_repository(i, "default"); + auto repo = manager.get_repository_shared(i, "default"); REQUIRE(repo != nullptr); } } @@ -375,8 +375,8 @@ TEST_CASE("data_repository_manager Large Number of Batches", "[data_repository_m } // Repository should have all batches - auto& repo = manager.get_repository(operator_id, "default"); - int count = 0; + auto repo = manager.get_repository_shared(operator_id, "default"); + int count = 0; while (true) { auto batch = repo->pop_next_data_batch(); if (!batch) break; @@ -412,7 +412,7 @@ TEST_CASE("data_repository_manager Thread-Safe Add Repository", "[data_repositor // All repositories should be accessible for (int i = 0; i < num_threads; ++i) { - auto& repo = manager.get_repository(i, "default"); + auto repo = manager.get_repository_shared(i, "default"); REQUIRE(repo != nullptr); } } @@ -454,7 +454,7 @@ TEST_CASE("data_repository_manager Thread-Safe Mixed Operations", "[data_reposit // Occasionally pull and store a batch (to test concurrent pull operations) if (j % 10 == 0) { - auto& repo = manager.get_repository(operator_id, "default"); + auto repo = manager.get_repository_shared(operator_id, "default"); auto pulled = repo->pop_next_data_batch(); if (pulled) { std::lock_guard lock(pull_mutex); @@ -524,7 +524,7 @@ TEST_CASE("data_repository_manager Concurrent Add and Pull", "[data_repository_m for (int i = 0; i < num_puller_threads; ++i) { threads.emplace_back([&, i]() { size_t operator_id = i % num_operators; - auto& repo = manager.get_repository(operator_id, "default"); + auto repo = manager.get_repository_shared(operator_id, "default"); // Keep pulling while adders are working while (keep_adding.load()) { @@ -567,7 +567,7 @@ TEST_CASE("data_repository_manager Concurrent Add and Pull", "[data_repository_m // All repositories should be empty for (int i = 0; i < num_operators; ++i) { - auto& repo = manager.get_repository(i, "default"); + auto repo = manager.get_repository_shared(i, "default"); auto batch = repo->pop_next_data_batch(); REQUIRE(batch == nullptr); } @@ -592,7 +592,7 @@ TEST_CASE("data_repository_manager High Contention Add Pull", "[data_repository_ // Launch threads doing both add and pull operations for (int i = 0; i < num_threads; ++i) { threads.emplace_back([&]() { - auto& repo = manager.get_repository(operator_id, "default"); + auto repo = manager.get_repository_shared(operator_id, "default"); std::vector> operator_ports = {{operator_id, "default"}}; for (int j = 0; j < operations_per_thread; ++j) { @@ -619,7 +619,7 @@ TEST_CASE("data_repository_manager High Contention Add Pull", "[data_repository_ REQUIRE(total_added == num_threads * operations_per_thread); // Clean up remaining batches - auto& repo = manager.get_repository(operator_id, "default"); + auto repo = manager.get_repository_shared(operator_id, "default"); while (true) { auto batch = repo->pop_next_data_batch(); if (!batch) break; @@ -663,7 +663,7 @@ TEST_CASE("data_repository_manager Concurrent Add Multiple Operators Per Batch", for (int i = 0; i < num_operators; ++i) { threads.emplace_back([&, i]() { - auto& repo = manager.get_repository(i, "default"); + auto repo = manager.get_repository_shared(i, "default"); // Pull all batches from this operator while (true) { @@ -685,7 +685,7 @@ TEST_CASE("data_repository_manager Concurrent Add Multiple Operators Per Batch", // All repositories should be empty for (int i = 0; i < num_operators; ++i) { - auto& repo = manager.get_repository(i, "default"); + auto repo = manager.get_repository_shared(i, "default"); auto batch = repo->pop_next_data_batch(); REQUIRE(batch == nullptr); } @@ -703,7 +703,7 @@ TEST_CASE("data_repository_manager Operator ID Zero", "[data_repository_manager] // Operator ID 0 should work like any other ID manager.add_new_repository(0, "default", std::make_unique()); - auto& repo = manager.get_repository(0, "default"); + auto repo = manager.get_repository_shared(0, "default"); REQUIRE(repo != nullptr); } @@ -721,7 +721,7 @@ TEST_CASE("data_repository_manager Large Operator IDs", "[data_repository_manage // All should be accessible for (size_t id : large_ids) { - auto& repo = manager.get_repository(id, "default"); + auto repo = manager.get_repository_shared(id, "default"); REQUIRE(repo != nullptr); } } @@ -747,8 +747,8 @@ TEST_CASE("data_repository_manager Batches With Different Sizes", "[data_reposit } // All batches should be accessible - auto& repo = manager.get_repository(operator_id, "default"); - int count = 0; + auto repo = manager.get_repository_shared(operator_id, "default"); + int count = 0; while (true) { auto batch = repo->pop_next_data_batch(); if (!batch) break; @@ -778,8 +778,8 @@ TEST_CASE("data_repository_manager Batches With Different Tiers", "[data_reposit } // All batches should be accessible - auto& repo = manager.get_repository(operator_id, "default"); - int count = 0; + auto repo = manager.get_repository_shared(operator_id, "default"); + int count = 0; while (true) { auto batch = repo->pop_next_data_batch(); if (!batch) break; @@ -798,7 +798,7 @@ TEST_CASE("data_repository_manager Rapid Add Pull Cycles", "[data_repository_man manager.add_new_repository(operator_id, "default", std::make_unique()); std::vector> operator_ports = {{operator_id, "default"}}; - auto& repo = manager.get_repository(operator_id, "default"); + auto repo = manager.get_repository_shared(operator_id, "default"); // Perform many cycles of add and pull for (int cycle = 0; cycle < 100; ++cycle) { @@ -816,3 +816,95 @@ TEST_CASE("data_repository_manager Rapid Add Pull Cycles", "[data_repository_man auto empty = repo->pop_next_data_batch(); REQUIRE(empty == nullptr); } + +// ============================================================================= +// Regression tests: get_repository_shared is locked and lifetime-safe (the old +// get_repository returned a reference into the map without taking _mutex, so a +// concurrent add_new_repository or clear_all_repositories raced both the +// lookup and the returned reference). +// ============================================================================= + +// A repository obtained just before clear_all_repositories must remain usable: +// the shared_ptr copy keeps it alive past the map erase. +TEST_CASE("data_repository_manager Repository Survives Clear All", "[data_repository_manager]") +{ + data_repository_manager manager; + size_t operator_id = 1; + manager.add_new_repository(operator_id, "default", std::make_unique()); + + auto data = std::make_unique(memory::Tier::GPU, 1024); + auto batch = data_batch::make(manager.get_next_data_batch_id(), std::move(data)); + manager.add_data_batch(batch, {{operator_id, "default"}}); + + auto repo = manager.get_repository_shared(operator_id, "default"); + REQUIRE(repo != nullptr); + + manager.clear_all_repositories(); + REQUIRE_THROWS_AS(manager.get_repository_shared(operator_id, "default"), std::out_of_range); + + // The detached repository is still fully usable through the shared_ptr. + REQUIRE(repo->total_size() == 1); + auto pulled = repo->pop_next_data_batch(); + REQUIRE(pulled != nullptr); + REQUIRE(repo->pop_next_data_batch() == nullptr); +} + +// Hammer the locked accessor against concurrent add/clear churn. Pre-fix this +// was a use-after-free of the map node (silent in release builds); post-fix +// every successfully returned repository must stay dereferenceable. +TEST_CASE("data_repository_manager Concurrent Get Vs Add And Clear", "[data_repository_manager]") +{ + data_repository_manager manager; + constexpr size_t num_operators = 8; + constexpr int num_iterations = 2000; + constexpr int num_readers = 4; + + std::atomic stop{false}; + std::atomic successful_gets{0}; + std::atomic failed_derefs{0}; + + // Writer: churn the map — register every operator, then clear them all. + std::thread writer([&] { + for (int iter = 0; iter < num_iterations; ++iter) { + for (size_t op = 0; op < num_operators; ++op) { + manager.add_new_repository(op, "default", std::make_unique()); + } + manager.clear_all_repositories(); + } + stop.store(true); + }); + + // Readers: race lookups against the churn and dereference every hit. + // (Catch2 assertion macros are not thread-safe — workers only count.) + std::vector readers; + for (int r = 0; r < num_readers; ++r) { + readers.emplace_back([&] { + while (!stop.load()) { + for (size_t op = 0; op < num_operators; ++op) { + try { + auto repo = manager.get_repository_shared(op, "default"); + if (!repo) { + failed_derefs.fetch_add(1); + continue; + } + // Dereference: pre-fix this touched a freed map node when the + // writer's clear ran between lookup and use. + (void)repo->total_size(); + successful_gets.fetch_add(1); + } catch (const std::out_of_range&) { + // The operator was between clear and re-add — expected. + } + } + } + }); + } + + writer.join(); + for (auto& t : readers) { + t.join(); + } + + REQUIRE(failed_derefs.load() == 0); + // The interleave actually exercised the racy window. + REQUIRE(successful_gets.load() > 0); +} From cc7badcb734149fbc141e45f2a8f0e3a55d3108a Mon Sep 17 00:00:00 2001 From: Felipe Aramburu Date: Sat, 15 Aug 2026 01:19:43 -0500 Subject: [PATCH 2/5] =?UTF-8?q?chore(data):=20H-group=20hygiene=20?= =?UTF-8?q?=E2=80=94=20drop=20the=20dead=20manager=20APIs,=20explain=20the?= =?UTF-8?q?=20aliases?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Register H10/H11 (sirius concurrency register), verified by grep across cucascade and the sirius tree before deletion: - data_repository_manager::get_repository() — the unlocked accessor's deprecated forwarding shim that B9 kept for source compatibility. Zero callers remain (sirius and cucascade both use get_repository_shared), so the migration window is over. The docs example that still showed it now shows get_repository_shared. - data_repository_manager::for_each_repository() — zero callers; every consumer uses the get_repositories() snapshot instead. - The include only for_each_repository needed. H10 judgement: shared_data_repository / shared_data_repository_manager stay, as compatibility aliases with comments saying exactly that. The names are no longer flatly wrong — data_repository stores shared_ptr (one batch sits in several repositories on fan-out) and since the shared_ptr map change the manager hands out lifetime-safe shared_ptr copies — and the aliases are load-bearing across the sirius tree (public signatures in batch_stream, gpu_pipeline_task, repository_wiring, batch_telemetry, the registry). The misleading part was that they read like distinct types with stronger sharing semantics; the new comments close that off and steer new code to the plain names. Co-Authored-By: Claude Fable 5 --- docs/data-management.md | 5 +-- include/cucascade/data/data_repository.hpp | 4 +++ .../data/data_repository_manager.hpp | 36 +++---------------- 3 files changed, 12 insertions(+), 33 deletions(-) diff --git a/docs/data-management.md b/docs/data-management.md index 48419a0..a04c330 100644 --- a/docs/data-management.md +++ b/docs/data-management.md @@ -448,8 +448,9 @@ manager.add_new_repository(0, "output", std::make_unique()); manager.add_new_repository(1, "input", std::make_unique()); manager.add_new_repository(1, "output", std::make_unique()); -// Access a specific repository -auto& repo = manager.get_repository(1, "input"); +// Access a specific repository (lifetime-safe shared_ptr, looked up under the +// manager mutex) +auto repo = manager.get_repository_shared(1, "input"); ``` ### Batch ID Generation diff --git a/include/cucascade/data/data_repository.hpp b/include/cucascade/data/data_repository.hpp index 6428533..a1949d6 100644 --- a/include/cucascade/data/data_repository.hpp +++ b/include/cucascade/data/data_repository.hpp @@ -316,6 +316,10 @@ class data_repository { _data_batches; ///< Container for data batch pointers (partitioned) }; +/// Compatibility alias, NOT a distinct type: the historical class that stored +/// `shared_ptr` was merged into data_repository (which now always +/// does — one batch can sit in several repositories on fan-out). Kept so old +/// call sites keep compiling; prefer `data_repository` in new code. using shared_data_repository = data_repository; } // namespace cucascade diff --git a/include/cucascade/data/data_repository_manager.hpp b/include/cucascade/data/data_repository_manager.hpp index 00d6c29..51447eb 100644 --- a/include/cucascade/data/data_repository_manager.hpp +++ b/include/cucascade/data/data_repository_manager.hpp @@ -21,7 +21,6 @@ #include #include -#include #include #include #include @@ -165,19 +164,6 @@ class data_repository_manager { return _repositories.at({operator_id, std::string(port_id)}); } - /** - * @brief DEPRECATED — use get_repository_shared(). - * - * Kept temporarily for source compatibility (callers used `.get()` / `->`, - * which keep compiling against the returned shared_ptr). Now forwards to the - * locked, lifetime-safe accessor; the historical signature returned a bare - * reference into the map without taking _mutex. - */ - std::shared_ptr get_repository(size_t operator_id, std::string_view port_id) - { - return get_repository_shared(operator_id, port_id); - } - /** * @brief Generate a globally unique data batch identifier. * @@ -223,23 +209,6 @@ class data_repository_manager { return leaked; } - /** - * @brief Iterate over all repositories, calling the visitor for each one. - * - * The visitor receives a raw pointer to each repository. The visitor must not - * remove or add repositories during iteration. - * - * @param visitor Callback invoked for each repository - * @note Thread-safe — holds the manager mutex for the duration of iteration. - */ - void for_each_repository(std::function visitor) - { - std::lock_guard lock(_mutex); - for (auto& [key, repo] : _repositories) { - if (repo) { visitor(repo.get()); } - } - } - /** * @brief Get a snapshot of all current repository pointers. * @@ -272,6 +241,11 @@ class data_repository_manager { std::map> _repositories; }; +/// Compatibility alias, NOT a distinct type: kept so call sites written +/// against the pre-merge class keep compiling. Since the map moved to +/// shared_ptr storage the "shared" in the name is loosely true (accessors hand +/// out lifetime-safe shared_ptr copies), but the manager remains each +/// repository's one logical owner. Prefer `data_repository_manager` in new code. using shared_data_repository_manager = data_repository_manager; } // namespace cucascade From 20c8cf27d37b8d6fad772bc4fc58cc613b5e50e7 Mon Sep 17 00:00:00 2001 From: Felipe Aramburu Date: Sat, 15 Aug 2026 04:04:28 -0500 Subject: [PATCH 3/5] fix(memory): FIFO ticket handoff for exclusive_stream_pool BLOCK checkout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All BLOCK-policy waiters used to share one CV: a released stream went to whichever waiter won the wake-up race, and a caller that released and immediately re-acquired always beat a parked waiter (no CV round trip), so checkout could starve under contention (Sirius register F9). Each BLOCK caller now draws a ticket under the pool lock and a released stream is handed to the lowest ticket — strict arrival order, so checkout is starvation-free. GROW callers never wait and no longer take a pooled stream a parked waiter is owed; they mint a fresh stream instead. API shapes unchanged. Co-Authored-By: Claude Fable 5 --- include/cucascade/memory/stream_pool.hpp | 13 +++++++++++ src/memory/stream_pool.cpp | 28 +++++++++++++++++++----- 2 files changed, 36 insertions(+), 5 deletions(-) diff --git a/include/cucascade/memory/stream_pool.hpp b/include/cucascade/memory/stream_pool.hpp index 8494474..7246100 100644 --- a/include/cucascade/memory/stream_pool.hpp +++ b/include/cucascade/memory/stream_pool.hpp @@ -22,6 +22,7 @@ #include #include +#include #include #include @@ -86,6 +87,12 @@ class exclusive_stream_pool { * * This function is thread safe with respect to other calls to the same function. * + * BLOCK-policy checkout is starvation-free: callers that find the pool empty take a FIFO + * ticket, and a released stream is handed to the longest-waiting caller. A caller that + * releases a stream and immediately re-acquires therefore queues behind every already-parked + * waiter instead of racing them for the freed stream. GROW-policy callers never wait and + * never take a pooled stream a parked waiter is owed — they mint a fresh stream instead. + * * @return rmm::cuda_stream_view */ borrowed_stream acquire_stream( @@ -103,6 +110,12 @@ class exclusive_stream_pool { // Streams are acquired from the front and returned to the back so the pool cycles through // all streams round-robin, rather than repeatedly reusing the most-recently-returned one. std::deque _streams; + // FIFO ticket handoff for BLOCK-policy checkout. Every BLOCK caller draws a ticket under + // _mutex; only the caller whose ticket equals _grant_ticket may take a stream, so waiters are + // served strictly in arrival order (one shared CV would otherwise wake an arbitrary waiter + // and let it win the race — a busy caller's re-acquire could starve a parked one forever). + std::uint64_t _next_ticket{0}; + std::uint64_t _grant_ticket{0}; }; } // namespace memory diff --git a/src/memory/stream_pool.cpp b/src/memory/stream_pool.cpp index 9b07450..ae5a74a 100644 --- a/src/memory/stream_pool.cpp +++ b/src/memory/stream_pool.cpp @@ -76,20 +76,35 @@ exclusive_stream_pool::exclusive_stream_pool(rmm::cuda_device_id device_id, borrowed_stream exclusive_stream_pool::acquire_stream(stream_acquire_policy policy) noexcept { std::unique_lock lock(_mutex); - if (_streams.empty()) { - if (policy == stream_acquire_policy::GROW) { + if (policy == stream_acquire_policy::GROW) { + // GROW never waits — and never takes a pooled stream a parked BLOCK waiter is owed + // (_grant_ticket != _next_ticket means at least one waiter is still unserved). Minting a + // fresh stream costs the same either way: the pool ends up one stream larger once it is + // released. + if (_streams.empty() || _grant_ticket != _next_ticket) { rmm::cuda_set_device_raii set_device{_device_id}; return borrowed_stream(rmm::cuda_stream(_flags), std::bind_front(&exclusive_stream_pool::release_stream, this)); - } else { - _cv.wait(lock, [this]() { return !_streams.empty(); }); } + } else { + // FIFO ticket handoff: a released stream goes to the LONGEST-WAITING caller. The ticket is + // drawn under the lock, so arrival order is the service order; a caller that releases and + // immediately re-acquires draws a fresh ticket and queues behind every parked waiter + // instead of winning the wake-up race against them. When the pool has streams and no + // earlier waiter is unserved, the predicate is immediately true and nothing blocks. + const std::uint64_t ticket = _next_ticket++; + _cv.wait(lock, [&]() { return ticket == _grant_ticket && !_streams.empty(); }); + ++_grant_ticket; } // Acquire from the front; release_stream() returns to the back. This cycles through all // streams round-robin so every stream's prior async work has maximal time to drain before // it is handed out again, instead of hammering the most-recently-returned stream. auto stream = std::move(_streams.front()); _streams.pop_front(); + // A single release wakes every waiter (they must all re-check the head ticket); if streams + // remain for the next-in-line, pass the baton before leaving so it does not sleep until the + // next release. + if (!_streams.empty() && _grant_ticket != _next_ticket) { _cv.notify_all(); } return borrowed_stream(std::move(stream), std::bind_front(&exclusive_stream_pool::release_stream, this)); } @@ -104,7 +119,10 @@ void exclusive_stream_pool::release_stream(rmm::cuda_stream&& s) noexcept { std::lock_guard lock(_mutex); _streams.emplace_back(std::move(s)); - _cv.notify_one(); + // notify_all, not notify_one: only the head-ticket waiter may take the stream, and with one + // shared CV a notify_one could wake a non-head waiter that just goes back to sleep while the + // head never learns a stream arrived. + _cv.notify_all(); } } // namespace memory From a9b7aba8f1827598912a5f7fd8305f119ad747f8 Mon Sep 17 00:00:00 2001 From: Felipe Aramburu Date: Sat, 15 Aug 2026 04:10:32 -0500 Subject: [PATCH 4/5] fix(memory): per-space FIFO wait lists for blocking reservations memory_space::make_reservation parked on the space's notification channel with race-wins semantics: a release set one flag and notify_one'd an arbitrary waiter, and a fresh blocking caller's non-blocking fast path could claim the freed memory before any parked waiter retried. A heavy caller's release-and-re-request loop could therefore perpetually beat a light caller's single wait. notification_channel gains a FIFO wait list (scoped_waiter RAII ticket): NOTIFIED is only delivered to the head of the list, releases wake all waiters so the head definitely learns (non-heads re-sleep), a departing waiter passes the baton to the next in line, and shutdown() wakes every waiter instead of one. make_reservation keeps its fast path when nobody is parked, and otherwise joins the FIFO instead of barging. make_reservation_or_null / make_reservation_upto keep their try-semantics untouched. API shapes unchanged. Also removed: memory_reservation_manager's cross-space _wait_mutex / _wait_cv slow path. It was waited on but never notified anywhere, so any request with no candidate space (or with every candidate shutting down) hung forever; request_reservation returns nullptr for that case instead, mirroring make_reservation's shutdown contract. Co-Authored-By: Claude Fable 5 --- .../memory/memory_reservation_manager.hpp | 8 +- .../cucascade/memory/notification_channel.hpp | 50 +++++++++++- src/memory/memory_reservation_manager.cpp | 20 ++--- src/memory/memory_space.cpp | 28 +++++-- src/memory/notification_channel.cpp | 81 ++++++++++++++++--- 5 files changed, 150 insertions(+), 37 deletions(-) diff --git a/include/cucascade/memory/memory_reservation_manager.hpp b/include/cucascade/memory/memory_reservation_manager.hpp index 1beb3fc..3f4daaa 100644 --- a/include/cucascade/memory/memory_reservation_manager.hpp +++ b/include/cucascade/memory/memory_reservation_manager.hpp @@ -24,10 +24,8 @@ #include #include -#include #include #include -#include #include #include #include @@ -272,9 +270,9 @@ class memory_reservation_manager { void build_lookup_tables(); - // Synchronization for cross-space waiting when no memory_space can currently satisfy a request - mutable std::mutex _wait_mutex; - std::condition_variable _wait_cv; + // Blocking-until-memory-frees is per-space state: each memory_space's notification channel + // keeps a FIFO wait list and hands a released reservation to its longest-waiting caller. No + // cross-space wait state lives at the manager level. }; } // namespace memory diff --git a/include/cucascade/memory/notification_channel.hpp b/include/cucascade/memory/notification_channel.hpp index e2243b0..4601133 100644 --- a/include/cucascade/memory/notification_channel.hpp +++ b/include/cucascade/memory/notification_channel.hpp @@ -18,6 +18,8 @@ #pragma once #include +#include +#include #include #include #include @@ -38,9 +40,52 @@ struct notification_channel : std::enable_shared_from_this enum class wait_status { IDLE, NOTIFIED, SHUTDOWN }; + /** + * @brief RAII registration in the channel's FIFO wait list. + * + * A blocking caller registers ONCE for the whole wait (taking a ticket under the channel + * lock) and then calls wait() as many times as its retry loop needs. wait() only returns + * NOTIFIED to the waiter at the HEAD of the list, so a release notification is offered to + * the longest-waiting caller first — never to whichever waiter happens to win the wake-up + * race. A caller whose retry fails keeps its ticket (it stays head); a caller that is done + * destroys the waiter, which removes the ticket and passes the baton to the next in line. + * + * IDLE and SHUTDOWN are delivered to every waiter regardless of position: they mean no + * notification can be forthcoming, so queue order is moot. + */ + class scoped_waiter { + public: + explicit scoped_waiter(notification_channel& channel); + ~scoped_waiter(); + + scoped_waiter(const scoped_waiter&) = delete; + scoped_waiter& operator=(const scoped_waiter&) = delete; + scoped_waiter(scoped_waiter&&) = delete; + scoped_waiter& operator=(scoped_waiter&&) = delete; + + /// @brief True when no earlier-registered waiter is still unserved. + [[nodiscard]] bool is_head() const; + + /** + * @brief Block until this waiter is at the head of the list AND a notification arrived + * (NOTIFIED), or no notification can be forthcoming (IDLE / SHUTDOWN). + */ + wait_status wait(); + + private: + std::shared_ptr _channel; + std::uint64_t _ticket; + }; + ~notification_channel(); - wait_status wait(); + /** + * @brief True when at least one scoped_waiter is registered. + * + * Used by blocking callers to keep a fresh arrival from barging past parked waiters via a + * non-blocking fast path: when waiters exist, join the FIFO instead. + */ + [[nodiscard]] bool has_waiters() const; std::unique_ptr get_notifier(); @@ -58,6 +103,9 @@ struct notification_channel : std::enable_shared_from_this bool _has_been_notified{false}; std::size_t _n_active_notifiers{0}; bool _is_running{true}; + /// Live waiter tickets in arrival order; front() is the only waiter NOTIFIED may go to. + std::deque _wait_queue; + std::uint64_t _next_ticket{0}; }; using event_notifier = notification_channel::event_notifier; diff --git a/src/memory/memory_reservation_manager.cpp b/src/memory/memory_reservation_manager.cpp index db01035..854dea7 100644 --- a/src/memory/memory_reservation_manager.cpp +++ b/src/memory/memory_reservation_manager.cpp @@ -27,7 +27,6 @@ #include #include -#include #include #include @@ -146,22 +145,15 @@ memory_reservation_manager::~memory_reservation_manager() { shutdown(); } std::unique_ptr memory_reservation_manager::request_reservation( const reservation_request_strategy& request, size_t size) { - // Fast path: try to make a reservation immediately + // All blocking happens inside memory_space::make_reservation, which parks on the space's + // per-space FIFO wait list and which select_memory_space_and_make_reservation already calls + // for every candidate. Reaching this point without a reservation therefore means no + // candidate space exists for the request, or every candidate is shutting down — report that + // as nullptr, mirroring make_reservation's shutdown contract. if (auto res = select_memory_space_and_make_reservation(request, size); res.has_value()) { return std::move(res.value()); } - - // If none available, block until any memory_space can satisfy the request - std::unique_lock lock(_wait_mutex); - for (;;) { - if (auto res = select_memory_space_and_make_reservation(request, size); res.has_value()) { - // Release the wait lock before returning the reservation - lock.unlock(); - return std::move(res.value()); - } - // Wait until notified that memory may be available again - _wait_cv.wait(lock); - } + return nullptr; } const memory_space* memory_reservation_manager::get_memory_space(Tier tier, int32_t device_id) const diff --git a/src/memory/memory_space.cpp b/src/memory/memory_space.cpp index ca8d6b4..bebd99b 100644 --- a/src/memory/memory_space.cpp +++ b/src/memory/memory_space.cpp @@ -256,14 +256,32 @@ std::unique_ptr memory_space::make_reservation_upto(size_t size) std::unique_ptr memory_space::make_reservation(size_t size) { - std::unique_ptr res = make_reservation_or_null(size); - while (!res) { - auto status = _notification_channel->wait(); + // Fast path: with nobody parked on this space, an immediately-satisfiable reservation is + // granted without touching the wait list. + if (!_notification_channel->has_waiters()) { + if (auto res = make_reservation_or_null(size)) { return res; } + } + + // Blocking path: join the space's FIFO wait list. A released reservation is offered to the + // LONGEST-WAITING caller first (scoped_waiter only returns NOTIFIED at the head of the + // list), and the has_waiters() gate above keeps a fresh blocking caller from barging past + // parked waiters through the fast path, so a heavy caller's release-and-re-request loop + // cannot perpetually beat a light caller's single wait. Non-blocking callers + // (make_reservation_or_null / make_reservation_upto direct users) keep their try-semantics + // and may claim memory ahead of the queue; they never park, so nothing can starve behind + // them indefinitely. + notification_channel::scoped_waiter waiter(*_notification_channel); + for (;;) { + // Retry only at the head; an attempt from the middle of the queue would barge past the + // waiters registered earlier. + if (waiter.is_head()) { + if (auto res = make_reservation_or_null(size)) { return res; } + } + auto status = waiter.wait(); if (status == notification_channel::wait_status::SHUTDOWN) { return nullptr; } if (status == notification_channel::wait_status::IDLE) { return make_reservation_upto(size); } - res = make_reservation_or_null(size); + // NOTIFIED: we are the head and a release arrived — retry. } - return res; } rmm::cuda_stream_view memory_space::acquire_stream() const diff --git a/src/memory/notification_channel.cpp b/src/memory/notification_channel.cpp index 7fc8c5e..f4c018a 100644 --- a/src/memory/notification_channel.cpp +++ b/src/memory/notification_channel.cpp @@ -17,6 +17,8 @@ #include +#include + namespace cucascade { namespace memory { @@ -41,22 +43,72 @@ void notification_channel::acquire_notifier() } //===----------------------------------------------------------------------===// -// notification_channel +// notification_channel::scoped_waiter //===----------------------------------------------------------------------===// -notification_channel::~notification_channel() { shutdown(); } +notification_channel::scoped_waiter::scoped_waiter(notification_channel& channel) + : _channel(channel.shared_from_this()) +{ + std::lock_guard lock(_channel->_mutex); + _ticket = _channel->_next_ticket++; + _channel->_wait_queue.push_back(_ticket); +} -notification_channel::wait_status notification_channel::wait() +notification_channel::scoped_waiter::~scoped_waiter() { - std::unique_lock lock(_mutex); + auto& ch = *_channel; + { + std::lock_guard lock(ch._mutex); + auto it = std::find(ch._wait_queue.begin(), ch._wait_queue.end(), _ticket); + if (it != ch._wait_queue.end()) { ch._wait_queue.erase(it); } + if (ch._wait_queue.empty()) { return; } + // Pass the baton: this waiter usually leaves because its reservation was just granted, so + // the space's state changed and the new head must re-check it — a release that arrived + // during our retry was coalesced into the single notified flag we may have consumed. A + // fabricated notification costs the new head one cheap failed retry at worst; NOT waking it + // costs a stall until the next release. + ch._has_been_notified = true; + } + ch._cv.notify_all(); +} + +bool notification_channel::scoped_waiter::is_head() const +{ + std::lock_guard lock(_channel->_mutex); + return !_channel->_wait_queue.empty() && _channel->_wait_queue.front() == _ticket; +} + +notification_channel::wait_status notification_channel::scoped_waiter::wait() +{ + auto& ch = *_channel; + std::unique_lock lock(ch._mutex); bool notified = false; - _cv.wait(lock, [&, self = shared_from_this()] { - notified = std::exchange(_has_been_notified, false); - return notified || (_n_active_notifiers == 0) || not _is_running; + ch._cv.wait(lock, [&] { + // IDLE / SHUTDOWN break every waiter out regardless of position (nothing further can be + // posted). Only the HEAD may consume a notification: consuming it from the middle of the + // queue is precisely the wake-up race the FIFO exists to remove. + if (!ch._is_running || ch._n_active_notifiers == 0) { return true; } + if (!ch._wait_queue.empty() && ch._wait_queue.front() == _ticket) { + notified = std::exchange(ch._has_been_notified, false); + return notified; + } + return false; }); - return !_is_running ? wait_status::SHUTDOWN - : (notified) ? wait_status::NOTIFIED - : wait_status::IDLE; + return !ch._is_running ? wait_status::SHUTDOWN + : (notified) ? wait_status::NOTIFIED + : wait_status::IDLE; +} + +//===----------------------------------------------------------------------===// +// notification_channel +//===----------------------------------------------------------------------===// + +notification_channel::~notification_channel() { shutdown(); } + +bool notification_channel::has_waiters() const +{ + std::lock_guard lock(_mutex); + return !_wait_queue.empty(); } std::unique_ptr notification_channel::get_notifier() @@ -68,14 +120,19 @@ void notification_channel::shutdown() { std::lock_guard lock(_mutex); _is_running = false; - _cv.notify_one(); + // Every parked waiter must observe the shutdown, not just one. + _cv.notify_all(); } void notification_channel::notify() { std::lock_guard lock(_mutex); _has_been_notified = true; - _cv.notify_one(); + // notify_all so the HEAD waiter definitely wakes: with one shared CV a notify_one can land on + // a non-head waiter, which re-checks its position and goes back to sleep while the head never + // learns memory was released. Waiter counts per space are small, so the extra wake-ups are + // noise. + _cv.notify_all(); } void notification_channel::release_notifier() From ed8d31565068c0198d5c6d32bebeff4ffd319ae7 Mon Sep 17 00:00:00 2001 From: Felipe Aramburu Date: Sat, 15 Aug 2026 03:57:56 -0500 Subject: [PATCH 5/5] feat(data): shared-ownership repository snapshots + destructor-side leak accounting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Step 6 groundwork for concurrent query teardown in Sirius: - data_repository_manager::get_repositories() now returns std::vector> instead of raw pointers. A memory-pressure sweep holds the snapshot across blocking work (thread pool reserves, host/device copies); raw pointers dangled the moment a concurrent teardown cleared the map. Shared elements make the snapshot self-owning, so teardown no longer needs an external fence. - data_repository gains a destructor-side leak callback (set_leak_callback): under shared ownership a repository can outlive its manager — a borrower may hold it past the owning query's erase — so the destructor is the one place that reliably observes batches that died un-consumed. - data_repository_manager::set_leak_handler() attributes those reports to the {operator_id, port_id} the repository was registered under, applied to current and future repositories. clear_all_repositories() keeps its snapshot-based report for callers that still clear eagerly; shared-ownership teardown paths simply drop the manager instead. Co-Authored-By: Claude Fable 5 --- include/cucascade/data/data_repository.hpp | 44 ++++++++++- .../data/data_repository_manager.hpp | 73 +++++++++++++++++-- 2 files changed, 108 insertions(+), 9 deletions(-) diff --git a/include/cucascade/data/data_repository.hpp b/include/cucascade/data/data_repository.hpp index a1949d6..a053211 100644 --- a/include/cucascade/data/data_repository.hpp +++ b/include/cucascade/data/data_repository.hpp @@ -20,6 +20,7 @@ #include #include +#include #include #include #include @@ -58,8 +59,45 @@ class data_repository { /** * @brief Virtual destructor for proper cleanup of derived classes. + * + * If a leak callback is installed (see set_leak_callback) and the repository still + * holds data batches, the callback is invoked with the un-consumed count before the + * batches are released. Under shared-ownership teardown the repository dies when its + * LAST holder releases it — the owning manager may already be gone — so the + * destructor is the one place that reliably observes what died un-consumed. */ - virtual ~data_repository() = default; + virtual ~data_repository() + { + if (!_leak_callback) { return; } + // No lock: destruction implies exclusive access. + std::size_t remaining = 0; + for (const auto& partition : _data_batches) { + remaining += partition.size(); + } + if (remaining == 0) { return; } + try { + _leak_callback(remaining); + } catch (...) { // a reporting hook must never throw out of a destructor + } + } + + /** + * @brief Install a callback invoked by the destructor when batches die un-consumed. + * + * The callback receives the number of data batches still held at destruction time. + * It runs on whatever thread drops the last reference to the repository and must not + * throw (exceptions are swallowed). Typically installed by the owning manager so the + * report can be attributed to an {operator, port} — and by extension a query. + * + * @param callback The leak-report hook (empty disables reporting). + * + * @note Thread-safe operation protected by internal mutex + */ + void set_leak_callback(std::function callback) + { + std::lock_guard lock(_mutex); + _leak_callback = std::move(callback); + } /** * @brief Add a new data batch to this repository. @@ -314,6 +352,10 @@ class data_repository { mutable std::mutex _mutex; ///< Mutex for thread-safe access to repository operations std::vector>> _data_batches; ///< Container for data batch pointers (partitioned) + + private: + /// Invoked by the destructor with the count of batches that died un-consumed. + std::function _leak_callback; }; /// Compatibility alias, NOT a distinct type: the historical class that stored diff --git a/include/cucascade/data/data_repository_manager.hpp b/include/cucascade/data/data_repository_manager.hpp index 51447eb..cc40fdb 100644 --- a/include/cucascade/data/data_repository_manager.hpp +++ b/include/cucascade/data/data_repository_manager.hpp @@ -21,11 +21,13 @@ #include #include +#include #include #include #include #include #include +#include #include namespace cucascade { @@ -82,6 +84,12 @@ class data_repository_manager { public: using repository_type = data_repository; + /// Invoked (via each repository's destructor-side leak callback) when a repository + /// dies still holding un-consumed data batches. Receives the {operator_id, port_id} + /// the repository was registered under and the batch count. + using leak_handler_type = + std::function; + /** * @brief Default constructor - initializes empty repository manager. */ @@ -117,10 +125,38 @@ class data_repository_manager { std::lock_guard lock(_mutex); auto it = _repositories.find({operator_id, std::string(port_id)}); if (it != _repositories.end()) { throw std::runtime_error("Repository already exists"); } + if (_leak_handler && shared_repository) { + install_leak_callback(*shared_repository, operator_id, std::string(port_id)); + } _repositories[{operator_id, std::string(port_id)}] = std::move(shared_repository); } } + /** + * @brief Install the handler invoked when a repository dies still holding batches. + * + * Applied to every repository already registered and to every repository added later. + * Under shared ownership a repository can outlive its manager (a borrower — e.g. a + * memory-pressure sweep — may hold it past the manager's teardown), so leak accounting + * lives in the repository's own destructor; this handler is how the owner attributes + * that report to an {operator_id, port_id}. The handler runs on whatever thread drops + * the last repository reference and must not throw. + * + * @param handler The attribution hook (empty leaves repositories unhooked from now on; + * already-installed callbacks are not removed) + * + * @note Thread-safe operation + */ + void set_leak_handler(leak_handler_type handler) + { + std::lock_guard lock(_mutex); + _leak_handler = std::move(handler); + if (!_leak_handler) { return; } + for (auto& [key, repo] : _repositories) { + if (repo) { install_leak_callback(*repo, key.operator_id, key.port_id); } + } + } + /** * @brief Add a data_batch to specified operator repositories. * @@ -193,6 +229,11 @@ class data_repository_manager { * un-consumed data batches, this is a bug — it means some operator didn't fully * drain its input. * + * @note With a leak handler installed (set_leak_handler), a non-empty repository this + * destroys ALSO fires its destructor-side report — callers should rely on one + * mechanism or the other. Shared-ownership teardown paths simply drop the + * manager instead of calling this, leaving the accounting to the destructors. + * * @return Per-repository info for each repository that still had un-consumed batches. */ std::vector clear_all_repositories() @@ -210,27 +251,40 @@ class data_repository_manager { } /** - * @brief Get a snapshot of all current repository pointers. + * @brief Get a lifetime-safe snapshot of all current repositories. * - * Returns a vector of raw pointers to each non-null repository. The vector - * is built under the manager mutex, so callers can iterate it externally - * without holding the lock (the repositories themselves remain thread-safe). + * Returns shared ownership of each non-null repository. The vector is built under + * the manager mutex, so callers can iterate it externally without holding the lock — + * and because each element co-owns its repository, the snapshot stays valid across + * blocking work even if the manager is concurrently cleared or destroyed. (The old + * variant returned raw pointers, which dangled the moment a concurrent teardown + * destroyed the map's shared_ptrs — exactly what a long memory-pressure sweep racing + * a query's end would hit.) * - * @return std::vector Snapshot of non-null repository pointers + * @return std::vector> Snapshot of non-null repositories * @note Thread-safe — holds the manager mutex for the duration of collection. */ - std::vector get_repositories() + std::vector> get_repositories() { std::lock_guard lock(_mutex); - std::vector result; + std::vector> result; result.reserve(_repositories.size()); for (auto& [key, repo] : _repositories) { - if (repo) { result.push_back(repo.get()); } + if (repo) { result.push_back(repo); } } return result; } private: + /// Hook @p repo's destructor-side leak report up to _leak_handler with this key's + /// attribution. Caller holds _mutex (for _leak_handler); the callback captures a COPY + /// of the handler so it stays valid on whatever thread the repository finally dies. + void install_leak_callback(repository_type& repo, std::size_t operator_id, std::string port_id) + { + repo.set_leak_callback([handler = _leak_handler, operator_id, port = std::move(port_id)]( + std::size_t count) { handler(operator_id, port, count); }); + } + std::mutex _mutex; ///< Mutex for thread-safe access std::atomic _next_data_batch_id = 0; ///< Atomic counter for generating unique data batch identifiers @@ -239,6 +293,9 @@ class data_repository_manager { /// clear_all_repositories / add_new_repository (the manager remains the one /// logical owner; accessors only extend lifetime across their use). std::map> _repositories; + /// Attribution hook for repositories that die still holding batches; see + /// set_leak_handler(). + leak_handler_type _leak_handler; }; /// Compatibility alias, NOT a distinct type: kept so call sites written