Concurrency-safe repository ownership and fair blocking primitives

docs/data-management.md

@@ -448,8 +448,9 @@ manager.add_new_repository(0, "output", std::make_unique<data_repository>());
 manager.add_new_repository(1, "input", std::make_unique<data_repository>());
 manager.add_new_repository(1, "output", std::make_unique<data_repository>());
 
-// 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

include/cucascade/data/data_repository.hpp

@@ -20,6 +20,7 @@
 #include <cucascade/data/data_batch.hpp>
 
 #include <condition_variable>
+#include <functional>
 #include <memory>
 #include <mutex>
 #include <optional>
@@ -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<void(std::size_t)> callback)
+  {
+    std::lock_guard<std::mutex> lock(_mutex);
+    _leak_callback = std::move(callback);
+  }
 
   /**
    * @brief Add a new data batch to this repository.
@@ -314,8 +352,16 @@ class data_repository {
   mutable std::mutex _mutex;  ///< Mutex for thread-safe access to repository operations
   std::vector<std::vector<std::shared_ptr<data_batch>>>
     _data_batches;  ///< Container for data batch pointers (partitioned)
+
+ private:
+  /// Invoked by the destructor with the count of batches that died un-consumed.
+  std::function<void(std::size_t)> _leak_callback;
 };
 
+/// Compatibility alias, NOT a distinct type: the historical class that stored
+/// `shared_ptr<data_batch>` 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

include/cucascade/data/data_repository_manager.hpp

@@ -27,6 +27,7 @@
 #include <mutex>
 #include <string>
 #include <unordered_map>
+#include <utility>
 #include <vector>
 
 namespace cucascade {
@@ -83,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<void(std::size_t operator_id, const std::string& port_id, std::size_t count)>;
+
   /**
    * @brief Default constructor - initializes empty repository manager.
    */
@@ -110,12 +117,43 @@ class data_repository_manager {
                           std::string_view port_id,
                           std::unique_ptr<repository_type> repository)
   {
-    std::unique_ptr<repository_type> 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<repository_type> shared_repository{std::move(repository)};
     {
       std::lock_guard<std::mutex> 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);
+      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<std::mutex> 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); }
     }
   }
 
@@ -139,21 +177,26 @@ 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<repository_type>& Reference to the repository
+   * @return std::shared_ptr<repository_type> 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<repository_type>& get_repository(size_t operator_id, std::string_view port_id)
+  std::shared_ptr<repository_type> get_repository_shared(size_t operator_id,
+                                                         std::string_view port_id)
   {
+    std::lock_guard<std::mutex> lock(_mutex);
     return _repositories.at({operator_id, std::string(port_id)});
   }
 
@@ -186,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<leaked_repository_info> clear_all_repositories()
@@ -203,51 +251,58 @@ class data_repository_manager {
   }
 
   /**
-   * @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<void(repository_type*)> visitor)
-  {
-    std::lock_guard<std::mutex> lock(_mutex);
-    for (auto& [key, repo] : _repositories) {
-      if (repo) { visitor(repo.get()); }
-    }
-  }
-
-  /**
-   * @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<repository_type*> Snapshot of non-null repository pointers
+   * @return std::vector<std::shared_ptr<repository_type>> Snapshot of non-null repositories
    * @note Thread-safe — holds the manager mutex for the duration of collection.
    */
-  std::vector<repository_type*> get_repositories()
+  std::vector<std::shared_ptr<repository_type>> get_repositories()
   {
     std::lock_guard<std::mutex> lock(_mutex);
-    std::vector<repository_type*> result;
+    std::vector<std::shared_ptr<repository_type>> 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<uint64_t> _next_data_batch_id =
     0;  ///< Atomic counter for generating unique data batch identifiers
-  std::map<operator_port_key, std::unique_ptr<repository_type>>
-    _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<operator_port_key, std::shared_ptr<repository_type>> _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
+/// 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

include/cucascade/memory/memory_reservation_manager.hpp

@@ -24,10 +24,8 @@
 #include <rmm/cuda_device.hpp>
 #include <rmm/cuda_stream_view.hpp>
 
-#include <condition_variable>
 #include <filesystem>
 #include <memory>
-#include <mutex>
 #include <optional>
 #include <span>
 #include <string>
@@ -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

include/cucascade/memory/notification_channel.hpp

@@ -18,6 +18,8 @@
 #pragma once
 
 #include <condition_variable>
+#include <cstdint>
+#include <deque>
 #include <memory>
 #include <mutex>
 #include <utility>
@@ -38,9 +40,52 @@ struct notification_channel : std::enable_shared_from_this<notification_channel>
 
   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<notification_channel> _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<event_notifier> get_notifier();
 
@@ -58,6 +103,9 @@ struct notification_channel : std::enable_shared_from_this<notification_channel>
   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<std::uint64_t> _wait_queue;
+  std::uint64_t _next_ticket{0};
 };
 
 using event_notifier = notification_channel::event_notifier;