Skip to content

feat!: make read_only_data_batch truly const #192

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dhruv9vats wants to merge 3 commits into
base: main
Choose a base branch
from
+128 −222
Open

feat!: make read_only_data_batch truly const #192

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a8c9774
Make read_only_data_batch truly const
dhruv9vats Aug 19, 2026
4632ad5
keep mutable to read only via idle test
dhruv9vats Aug 19, 2026
08e1506
doc fix
dhruv9vats Aug 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
183 changes: 86 additions & 97 deletions include/cucascade/data/data_batch.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,6 @@
#include <mutex>
#include <optional>
#include <shared_mutex>
#include <stdexcept>
#include <unordered_map>
#include <utility>
#include <vector>
Expand Down Expand Up @@ -158,17 +157,19 @@ class data_batch : public std::enable_shared_from_this<data_batch> {
* @brief Transition from read-only back to idle (release shared lock).
*
* @param accessor Rvalue reference to the read-only accessor (consumed).
* @return The batch pointer, now in idle state.
* @return The batch pointer, now in idle state. Use `std::ignore =
* data_batch::to_idle(read_handle)` to simply let go of the read accessor.
*/
[[nodiscard]] static std::shared_ptr<data_batch> to_idle(read_only_data_batch&& accessor);
static std::shared_ptr<const data_batch> to_idle(read_only_data_batch&& accessor);

/**
* @brief Transition from mutable back to idle (release exclusive lock).
*
* @param accessor Rvalue reference to the mutable accessor (consumed).
* @return The batch pointer, now in idle state.
* @return The batch pointer, now in idle state. Use `std::ignore =
* data_batch::to_idle(write_handle)` to simply let go of the mutable accessor
*/
[[nodiscard]] static std::shared_ptr<data_batch> to_idle(mutable_data_batch&& accessor);
static std::shared_ptr<data_batch> to_idle(mutable_data_batch&& accessor);

// -- Non-static transitions (via shared_from_this) --
// The caller's shared_ptr is NOT consumed. These only work when the
Expand All @@ -182,7 +183,7 @@ class data_batch : public std::enable_shared_from_this<data_batch> {
*
* @return A read_only_data_batch holding the shared lock.
*/
[[nodiscard]] read_only_data_batch to_read_only();
[[nodiscard]] read_only_data_batch to_read_only() const;

/**
* @brief Transition from idle to mutable (exclusive lock) without consuming the caller's pointer.
Expand All @@ -204,7 +205,7 @@ class data_batch : public std::enable_shared_from_this<data_batch> {
* @return An optional containing the read-only accessor on success, or
* std::nullopt if the lock could not be acquired immediately.
*/
[[nodiscard]] std::optional<read_only_data_batch> try_to_read_only();
[[nodiscard]] std::optional<read_only_data_batch> try_to_read_only() const;

/**
* @brief Try to transition from idle to mutable (non-blocking).
Expand All @@ -216,36 +217,6 @@ class data_batch : public std::enable_shared_from_this<data_batch> {
*/
[[nodiscard]] std::optional<mutable_data_batch> try_to_mutable();

// -- Locked-to-locked static transitions --

/**
* @brief Transition from read-only to mutable (upgrade lock).
*
* Releases the shared lock, then acquires an exclusive lock (may block).
* Waits for all recorded asynchronous readers before acquiring the exclusive lock
* (re-checked under the lock) and exposing mutable access.
* The source accessor is consumed via move.
* NOTE: The transition is not atomic.
*
* @param accessor Rvalue reference to the read-only accessor (consumed).
* @return A mutable_data_batch holding the exclusive lock.
* @throws cucascade::cuda_error if a recorded reader event cannot be synchronized.
* @throws rmm::cuda_error if a reader event's CUDA device cannot be made current.
*/
[[nodiscard]] static mutable_data_batch readonly_to_mutable(read_only_data_batch&& accessor);

/**
* @brief Transition from mutable to read-only (downgrade lock).
*
* Releases the exclusive lock, then acquires a shared lock (may block).
* The source accessor is consumed via move.
* NOTE: The transition is not atomic.
*
* @param accessor Rvalue reference to the mutable accessor (consumed).
* @return A read_only_data_batch holding the shared lock.
*/
[[nodiscard]] static read_only_data_batch mutable_to_readonly(mutable_data_batch&& accessor);

private:
data_batch(uint64_t batch_id,
std::unique_ptr<idata_representation> data,
Expand Down Expand Up @@ -275,66 +246,74 @@ class data_batch : public std::enable_shared_from_this<data_batch> {
*/
void set_data(std::unique_ptr<idata_representation> data);

/**
* @brief Record completion of asynchronous work reading the current representation.
*
* GPU batches retain one event for every outstanding reader stream. Non-GPU tiers are a no-op.
* The caller must hold a shared lock while recording so the representation cannot change between
* issuing the read and recording its completion.
*
* @param reader_stream Stream on which work reading the batch was enqueued.
*/
void record_reader_event(rmm::cuda_stream_view reader_stream);

/**
* @brief Block until all recorded asynchronous readers have completed.
*
* Called while the batch's exclusive lock is held before mutable access is exposed.
*/
void synchronize_reader_events();

/**
* @brief Query and recycle completed reader events without blocking.
*
* @return true when no recorded asynchronous reader remains in flight.
* @throws cucascade::cuda_error if an event's query reports a CUDA failure.
*/
[[nodiscard]] bool reader_events_complete();

/**
* @brief Destructor-safe form of synchronize_reader_events().
*/
void synchronize_reader_events_no_throw() noexcept;

struct reader_event_pool {
// Events are stored as a pending prefix followed by reusable completed events. This bounds
// event creation by the peak number of overlapping reads registered on this CUDA device.
std::vector<cuda::cuda_event> events;
std::size_t pending_event_count{0};
};

/**
* @brief Move completed events from the pending prefix back into the reusable pool.
*
* Requires _reader_events_mutex to be held and the pool's CUDA device to be current.
* A query that reports failure is propagated rather than counted as still-pending: a failed
* event never completes, so absorbing it would stall every later mutable acquisition.
*
* @throws cucascade::cuda_error if an event's query reports a CUDA failure. @p pool is left
* unchanged from the failing event onward.
*/
void recycle_completed_reader_events(reader_event_pool& pool);
struct reader_event_pool_map {
// CUDA events are device-associated, so each device needs an independent reusable pool.
std::mutex reader_events_mutex;
std::unordered_map<int, reader_event_pool> reader_event_pools;

/**
* @brief Record completion of asynchronous work reading the current representation.
*
* GPU batches retain one event for every outstanding reader stream. Non-GPU tiers are a no-op.
* The caller must hold a shared lock while recording so the representation cannot change
* between issuing the read and recording its completion.
*
* @param reader_stream Stream on which work reading the batch was enqueued.
*/
void record_reader_event(rmm::cuda_stream_view reader_stream);

/**
* @brief Block until all recorded asynchronous readers have completed.
*
* Called while the batch's exclusive lock is held before mutable access is exposed.
*/
void synchronize_reader_events();

/**
* @brief Query and recycle completed reader events without blocking.
*
* @return true when no recorded asynchronous reader remains in flight.
* @throws cucascade::cuda_error if an event's query reports a CUDA failure.
*/
[[nodiscard]] bool reader_events_complete();

/**
* @brief Destructor-safe form of synchronize_reader_events().
*/
void synchronize_reader_events_no_throw() noexcept;

/**
* @brief Move completed events from the pending prefix back into the reusable pool.
*
* Requires _reader_events_mutex to be held and the pool's CUDA device to be current.
* A query that reports failure is propagated rather than counted as still-pending: a failed
* event never completes, so absorbing it would stall every later mutable acquisition.
*
* @throws cucascade::cuda_error if an event's query reports a CUDA failure. @p pool is left
* unchanged from the failing event onward.
*/
void recycle_completed_reader_events(reader_event_pool& pool) const;
};

const uint64_t _batch_id; ///< Immutable batch identifier
std::unique_ptr<idata_representation> _data; ///< Owned data representation
mutable std::shared_mutex _rw_mutex; ///< Reader-writer mutex
std::atomic<size_t> _subscriber_count{0}; ///< Atomic subscriber interest count
std::atomic<batch_state> _state{batch_state::idle}; ///< Observable lock state
std::atomic<size_t> _read_only_count{0}; ///< Count of active read_only_data_batch instances
const uint64_t _batch_id; ///< Immutable batch identifier
std::unique_ptr<idata_representation> _data; ///< Owned data representation
mutable std::shared_mutex _rw_mutex; ///< Reader-writer mutex
mutable std::atomic<size_t> _subscriber_count{0}; ///< Atomic subscriber interest count
mutable std::atomic<batch_state> _state{batch_state::idle}; ///< Observable lock state
mutable std::atomic<size_t> _read_only_count{
0}; ///< Count of active read_only_data_batch instances

// CUDA events are device-associated, so each device needs an independent reusable pool.
std::mutex _reader_events_mutex;
std::unordered_map<int, reader_event_pool> _reader_event_pools;
// mutable is for read_only case which holds `const data_batch`, but is OK because this
// internally uses a separate mutex.
mutable reader_event_pool_map _reader_event_pools;

std::unique_ptr<idata_batch_probe> _probe;
};
Expand All @@ -360,16 +339,19 @@ class read_only_data_batch {
// -- Named accessor methods --

/** @brief Get the batch identifier. */
uint64_t get_batch_id() const { return _batch->get_batch_id(); }
[[nodiscard]] uint64_t get_batch_id() const { return _batch->get_batch_id(); }

/** @brief Get the memory tier of the held data. */
memory::Tier get_current_tier() const { return _batch->get_current_tier(); }
[[nodiscard]] memory::Tier get_current_tier() const { return _batch->get_current_tier(); }

/** @brief Get a raw pointer to the data representation. */
[[nodiscard]] const idata_representation* get_data() const { return _batch->get_data(); }

/** @brief Get a raw pointer to the memory space. */
memory::memory_space* get_memory_space() const { return _batch->get_memory_space(); }
[[nodiscard]] memory::memory_space* get_memory_space() const
{
return _batch->get_memory_space();
}

/**
* @brief Get the writer event from the underlying representation, or nullptr.
Expand Down Expand Up @@ -418,7 +400,11 @@ class read_only_data_batch {
*/
void record_reader_event(rmm::cuda_stream_view reader_stream) const
{
_batch->record_reader_event(reader_stream);
// Host and disk representations do not expose stream-ordered device memory.
if (_batch->_data == nullptr || _batch->_data->get_current_tier() != memory::Tier::GPU) {
return;
}
_batch->_reader_event_pools.record_reader_event(reader_stream);
}

/**
Expand Down Expand Up @@ -499,13 +485,13 @@ class read_only_data_batch {
* @param parent Shared pointer to the parent data_batch (moved in).
* @param lock Shared lock already acquired on the parent's mutex.
*/
read_only_data_batch(std::shared_ptr<data_batch> parent,
read_only_data_batch(std::shared_ptr<const data_batch> parent,
std::shared_lock<std::shared_mutex> lock);

// INVARIANT: _batch must be declared before _lock -- destruction order is load-bearing.
// When destroyed, _lock releases the shared lock first, then _batch drops the parent
// reference. This prevents accessing a destroyed mutex.
std::shared_ptr<data_batch> _batch; ///< Parent lifetime (destroyed second)
std::shared_ptr<const data_batch> _batch; ///< Parent lifetime (destroyed second)
std::shared_lock<std::shared_mutex> _lock; ///< Shared lock (destroyed first)
};

Expand All @@ -523,16 +509,19 @@ class mutable_data_batch {
// -- Read methods (same as read_only) --

/** @brief Get the batch identifier. */
uint64_t get_batch_id() const { return _batch->get_batch_id(); }
[[nodiscard]] uint64_t get_batch_id() const { return _batch->get_batch_id(); }

/** @brief Get the memory tier of the held data. */
memory::Tier get_current_tier() const { return _batch->get_current_tier(); }
[[nodiscard]] memory::Tier get_current_tier() const { return _batch->get_current_tier(); }

/** @brief Get a raw pointer to the data representation. */
idata_representation* get_data() const { return _batch->get_data(); }
[[nodiscard]] idata_representation* get_data() const { return _batch->get_data(); }

/** @brief Get a raw pointer to the memory space. */
memory::memory_space* get_memory_space() const { return _batch->get_memory_space(); }
[[nodiscard]] memory::memory_space* get_memory_space() const
{
return _batch->get_memory_space();
}

// -- Write methods --

Expand Down
Loading
Loading