diff --git a/include/rocksdb/io_dispatcher.h b/include/rocksdb/io_dispatcher.h index 9c3fefd640..6354d72ad3 100644 --- a/include/rocksdb/io_dispatcher.h +++ b/include/rocksdb/io_dispatcher.h @@ -6,8 +6,10 @@ #pragma once #include +#include #include #include +#include #include #include "rocksdb/options.h" @@ -17,6 +19,22 @@ namespace ROCKSDB_NAMESPACE { class FileSystem; +class Statistics; + +// Forward declaration for internal implementation +struct IODispatcherImplData; +struct PendingPrefetchRequest; + +// Options for configuring IODispatcher behavior +struct IODispatcherOptions { + // Maximum memory (in bytes) for prefetching across all ReadSets. + // When this limit is reached, SubmitJob() blocks until memory is released. + // Set to 0 (default) for unlimited prefetch memory. + size_t max_prefetch_memory_bytes = 0; + + // Optional statistics for tracking memory limiter metrics + Statistics* statistics = nullptr; +}; /* * IODispatcher is a class that allows users to submit groups of IO jobs to be @@ -33,51 +51,88 @@ class FileSystem; * dispatcher, allowing for future ratelimiting and smarter dispatching policies * in the future. * -* Example: - // Submitting an IO job and reading blocks: - // - // std::shared_ptr job = std::make_shared(); - // job->table = table_reader; // Provided BlockBasedTable* - // job->job_options.io_coalesce_threshold = 32 * 1024; - // job->job_options.read_options = read_options; // Provided ReadOptions - // - // // Populate the job with block handles (e.g., from an index/iterator) - // job->block_handles.push_back(handle1); - // job->block_handles.push_back(handle2); - // job->block_handles.push_back(handle3); - // - // std::unique_ptr dispatcher(NewIODispatcher()); - // std::shared_ptr read_set; - // Status s = dispatcher->SubmitJob(job, &read_set); - // if (!s.ok()) { - // // Handle submit error - // } - // - // // Read by index - // for (size_t i = 1; i < job->block_handles.size(); ++i) { - // CachableEntry block_entry; - // Status rs = read_set->ReadIndex(i, &block_entry); - // if (!rs.ok()) { - // // Handle read error - // continue; - // } - // // Use block_entry (block contents are pinned here) - // } - // - // // Or read by byte offset - // { - // size_t offset = static_cast(job->block_handles.front().offset()); - // CachableEntry block_entry; - // Status rs = read_set->ReadOffset(offset, &block_entry); - // if (rs.ok()) { - // // Use block_entry - // } - // } - // - // // Stats - // uint64_t cache_hits = read_set->GetNumCacheHits(); - // uint64_t async_reads = read_set->GetNumAsyncReads(); - // uint64_t sync_reads = read_set->GetNumSyncReads(); + * Example 1: Basic Usage + * ---------------------- + * // Submitting an IO job and reading blocks: + * // + * // std::shared_ptr job = std::make_shared(); + * // job->table = table_reader; // Provided BlockBasedTable* + * // job->job_options.io_coalesce_threshold = 32 * 1024; + * // job->job_options.read_options = read_options; // Provided ReadOptions + * // + * // // Populate the job with block handles (e.g., from an index/iterator) + * // job->block_handles.push_back(handle1); + * // job->block_handles.push_back(handle2); + * // job->block_handles.push_back(handle3); + * // + * // std::unique_ptr dispatcher(NewIODispatcher()); + * // std::shared_ptr read_set; + * // Status s = dispatcher->SubmitJob(job, &read_set); + * // if (!s.ok()) { + * // // Handle submit error + * // } + * // + * // // Read by index + * // for (size_t i = 1; i < job->block_handles.size(); ++i) { + * // CachableEntry block_entry; + * // Status rs = read_set->ReadIndex(i, &block_entry); + * // if (!rs.ok()) { + * // // Handle read error + * // continue; + * // } + * // // Use block_entry (block contents are pinned here) + * // } + * // + * // // Or read by byte offset + * // { + * // size_t offset = + static_cast(job->block_handles.front().offset()); + * // CachableEntry block_entry; + * // Status rs = read_set->ReadOffset(offset, &block_entry); + * // if (rs.ok()) { + * // // Use block_entry + * // } + * // } + * // + * // // Stats + * // uint64_t cache_hits = read_set->GetNumCacheHits(); + * // uint64_t async_reads = read_set->GetNumAsyncReads(); + * // uint64_t sync_reads = read_set->GetNumSyncReads(); + * + * Example 2: Memory-Limited Prefetching + * ------------------------------------- + * // Configure a memory budget for prefetching to prevent unbounded memory use. + * // When the budget is exceeded, IODispatcher uses "partial prefetch": + * // - Dispatches as many blocks as fit in available memory (earlier first) + * // - Queues remaining blocks for later dispatch when memory is released + * // - Never blocks on SubmitJob - remaining blocks are read on-demand + * // + * // IODispatcherOptions opts; + * // opts.max_prefetch_memory_bytes = 64 * 1024 * 1024; // 64MB budget + * // opts.statistics = db_options.statistics.get(); // Optional metrics + * // + * // std::unique_ptr dispatcher(NewIODispatcher(opts)); + * // + * // // Submit a job that needs more memory than available + * // // Partial prefetch will dispatch what fits immediately + * // std::shared_ptr read_set; + * // Status s = dispatcher->SubmitJob(job, &read_set); // Never blocks + * // + * // // Read blocks in order - earlier blocks are more likely to be prefetched + * // for (size_t i = 0; i < job->block_handles.size(); ++i) { + * // CachableEntry block; + * // Status rs = read_set->ReadIndex(i, &block); + * // // Use block... + * // + * // // Release block when done to free memory for pending prefetches + * // read_set->ReleaseBlock(i); // Triggers dispatch of queued blocks + * // } + * // + * // Memory limiting statistics (when statistics is configured): + * // - PREFETCH_MEMORY_BYTES_GRANTED: Total bytes acquired for prefetching + * // - PREFETCH_MEMORY_BYTES_RELEASED: Total bytes released after use + * // - PREFETCH_MEMORY_REQUESTS_BLOCKED: Number of blocks that couldn't be + * // prefetched immediately due to memory pressure */ @@ -180,6 +235,13 @@ class ReadSet { // blocks are coalesced into a single IO request. std::unordered_map> async_io_map_; + // For memory release notifications back to dispatcher (weak ref to avoid + // cycles) + std::weak_ptr dispatcher_data_; + + // Size of each block (parallel to pinned_blocks_) for memory accounting + std::vector block_sizes_; + // Statistics counters std::atomic num_sync_reads_ = 0; std::atomic num_async_reads_ = 0; @@ -191,6 +253,16 @@ class ReadSet { // Perform synchronous read for a specific block Status SyncRead(size_t block_index); + + // Remove a block from pending prefetch (called by ReadIndex/ReleaseBlock) + void RemoveFromPending(size_t block_index); + + // Atomic flags indicating if block is pending prefetch (lock-free check) + std::unique_ptr[]> pending_prefetch_flags_; + size_t pending_prefetch_flags_size_ = 0; + + // Reference to pending request (for removal notification) + std::shared_ptr pending_request_; }; /* @@ -218,8 +290,12 @@ class IODispatcher { std::shared_ptr* read_set) = 0; }; +// Create IODispatcher with default options (no memory limit) IODispatcher* NewIODispatcher(); +// Create IODispatcher with custom options +IODispatcher* NewIODispatcher(const IODispatcherOptions& options); + // TrackingIODispatcher wraps another IODispatcher and tracks all ReadSets // created. This is useful for testing to verify IO statistics. class TrackingIODispatcher : public IODispatcher { diff --git a/include/rocksdb/statistics.h b/include/rocksdb/statistics.h index ae4ef57924..7cecac05f7 100644 --- a/include/rocksdb/statistics.h +++ b/include/rocksdb/statistics.h @@ -575,6 +575,14 @@ enum Tickers : uint32_t { // # of seeks that failed validation (out of order, etc.) MULTISCAN_SEEK_ERRORS, + // IODispatcher memory limiting statistics + // # of bytes granted to prefetch requests + PREFETCH_MEMORY_BYTES_GRANTED, + // # of bytes released from prefetch memory + PREFETCH_MEMORY_BYTES_RELEASED, + // # of prefetch requests that were blocked waiting for memory + PREFETCH_MEMORY_REQUESTS_BLOCKED, + TICKER_ENUM_MAX }; diff --git a/monitoring/statistics.cc b/monitoring/statistics.cc index ccc92bcb61..e6060cbeac 100644 --- a/monitoring/statistics.cc +++ b/monitoring/statistics.cc @@ -292,6 +292,10 @@ const std::vector> TickersNameMap = { {MULTISCAN_IO_COALESCED_NONADJACENT, "rocksdb.multiscan.io.coalesced.nonadjacent"}, {MULTISCAN_SEEK_ERRORS, "rocksdb.multiscan.seek.errors"}, + {PREFETCH_MEMORY_BYTES_GRANTED, "rocksdb.prefetch.memory.bytes.granted"}, + {PREFETCH_MEMORY_BYTES_RELEASED, "rocksdb.prefetch.memory.bytes.released"}, + {PREFETCH_MEMORY_REQUESTS_BLOCKED, + "rocksdb.prefetch.memory.requests.blocked"}, }; const std::vector> HistogramsNameMap = { diff --git a/util/io_dispatcher_imp.cc b/util/io_dispatcher_imp.cc index fd7d590185..2789414860 100644 --- a/util/io_dispatcher_imp.cc +++ b/util/io_dispatcher_imp.cc @@ -14,12 +14,15 @@ #include "util/io_dispatcher_imp.h" +#include #include #include #include #include #include "file/random_access_file_reader.h" +#include "monitoring/statistics_impl.h" +#include "port/port.h" #include "rocksdb/file_system.h" #include "rocksdb/io_dispatcher.h" #include "rocksdb/options.h" @@ -28,9 +31,19 @@ #include "table/block_based/cachable_entry.h" #include "table/block_based/reader_common.h" #include "table/format.h" +#include "test_util/sync_point.h" +#include "util/mutexlock.h" namespace ROCKSDB_NAMESPACE { +// IODispatcherImplData is the base that provides ReleaseMemory interface +// for ReadSets to call back when releasing blocks. Defined here so it's +// visible to ReadSet methods. +struct IODispatcherImplData { + virtual ~IODispatcherImplData() = default; + virtual void ReleaseMemory(size_t bytes) = 0; +}; + // Helper function to create and pin a block from a buffer // Used by both ReadSet::PollAndProcessAsyncIO and IODispatcherImpl::Impl static Status CreateAndPinBlockFromBuffer( @@ -98,6 +111,18 @@ struct AsyncIOState { // Must call AbortIO before deleting handles to avoid use-after-free when // io_uring completions arrive for deleted handles. ReadSet::~ReadSet() { + // Release memory for any blocks still pinned + // Note: block_sizes_[i] is only set for async IO reads where memory + // limiting applies. For sync reads, block_sizes_ remains 0, so this + // loop is effectively a no-op for sync reads. + if (auto dispatcher_data = dispatcher_data_.lock()) { + for (size_t i = 0; i < block_sizes_.size(); ++i) { + if (block_sizes_[i] > 0 && pinned_blocks_[i].GetValue()) { + dispatcher_data->ReleaseMemory(block_sizes_[i]); + } + } + } + if (async_io_map_.empty()) { return; } @@ -173,6 +198,9 @@ Status ReadSet::ReadIndex(size_t block_index, CachableEntry* out) { } // Case 3: Block needs synchronous read + // If this block was pending prefetch, remove it since we're reading it now + RemoveFromPending(block_index); + Status s = SyncRead(block_index); if (s.ok()) { *out = std::move(pinned_blocks_[block_index]); @@ -219,6 +247,22 @@ void ReadSet::ReleaseBlock(size_t block_index) { if (block_index >= pinned_blocks_.size()) { return; } + + // Remove from pending if applicable + RemoveFromPending(block_index); + + // Release memory BEFORE unpinning + // Note: block_sizes_[idx] is only set for async IO reads where memory + // limiting applies. For sync reads, block_sizes_ remains 0, so this + // check implicitly skips ReleaseMemory for sync reads. + if (pinned_blocks_[block_index].GetValue() && + block_index < block_sizes_.size() && block_sizes_[block_index] > 0) { + if (auto dispatcher_data = dispatcher_data_.lock()) { + dispatcher_data->ReleaseMemory(block_sizes_[block_index]); + } + block_sizes_[block_index] = 0; // Prevent double-release + } + // Unpin the block from cache pinned_blocks_[block_index].Reset(); // Clean up any pending async IO for this block @@ -300,9 +344,51 @@ Status ReadSet::SyncRead(size_t block_index) { /*async_read=*/false, /*use_block_cache_for_lookup=*/true); } -struct IODispatcherImpl::Impl { - Impl(); - ~Impl(); +// A pre-coalesced group of blocks for prefetching +struct CoalescedPrefetchGroup { + std::vector block_indices; // Blocks in this group (sorted by offset) + size_t total_bytes = 0; // Total bytes for this IO +}; + +// State for a pending memory request waiting to be granted +// Groups are pre-coalesced at queue time for efficient dispatch +struct PendingPrefetchRequest { + std::weak_ptr read_set; + std::shared_ptr job; + + // Pre-coalesced groups ready for dispatch (ordered by first block index) + std::deque coalesced_groups; + + // Individual block indices still pending (for RemoveFromPending lookup) + std::unordered_set block_indices_to_prefetch; + + std::atomic pending_bytes_{0}; // Track remaining bytes + mutable port::Mutex groups_mutex_; // Protects groups and set modifications +}; + +// Remove a block from pending prefetch (called when block is read or released) +void ReadSet::RemoveFromPending(size_t block_index) { + if (!pending_prefetch_flags_ || block_index >= pending_prefetch_flags_size_) { + return; + } + + // Atomic exchange - returns true only if it was previously true + if (!pending_prefetch_flags_[block_index].exchange(false)) { + return; // Already removed or never pending + } + + if (pending_request_) { + MutexLock lock(&pending_request_->groups_mutex_); + pending_request_->block_indices_to_prefetch.erase(block_index); + pending_request_->pending_bytes_ -= block_sizes_[block_index]; + } +} + +// IODispatcherImpl::Impl inherits from IODispatcherImplData +struct IODispatcherImpl::Impl : public IODispatcherImplData, + public std::enable_shared_from_this { + explicit Impl(const IODispatcherOptions& options); + ~Impl() override; // Non-copyable and non-movable Impl(const Impl&) = delete; @@ -313,6 +399,18 @@ struct IODispatcherImpl::Impl { Status SubmitJob(const std::shared_ptr& job, std::shared_ptr* read_set); + // Memory management methods - non-blocking + bool TryAcquireMemory(size_t bytes); + void ReleaseMemory(size_t bytes) override; + + // Memory limiting state + size_t max_prefetch_memory_bytes_ = 0; + std::atomic memory_used_{0}; // Atomic for lock-free accounting + std::atomic has_pending_requests_{false}; // Fast-path check + port::Mutex memory_mutex_; // Only for pending_prefetch_queue_ access + std::deque> pending_prefetch_queue_; + Statistics* statistics_ = nullptr; + private: void PrepareIORequests( const std::shared_ptr& job, @@ -335,12 +433,214 @@ struct IODispatcherImpl::Impl { const std::shared_ptr& read_set, std::vector& read_reqs, const std::vector>& coalesced_block_indices); + + // Try to dispatch pending prefetch requests when memory becomes available + void TryDispatchPendingPrefetches(); + + // Dispatch prefetch for a specific ReadSet (called when memory is available) + void DispatchPrefetch(const std::shared_ptr& read_set, + const std::shared_ptr& job, + const std::vector& block_indices); + + // Pre-coalesce blocks into groups, respecting max_group_bytes size limit. + // Returns groups ordered by first block index (earlier blocks first). + std::vector PreCoalesceBlocks( + const std::shared_ptr& job, const std::shared_ptr& rs, + const std::vector& block_indices, size_t max_group_bytes); }; -IODispatcherImpl::Impl::Impl() {} +IODispatcherImpl::Impl::Impl(const IODispatcherOptions& options) + : max_prefetch_memory_bytes_(options.max_prefetch_memory_bytes), + statistics_(options.statistics) {} IODispatcherImpl::Impl::~Impl() {} +bool IODispatcherImpl::Impl::TryAcquireMemory(size_t bytes) { + if (max_prefetch_memory_bytes_ == 0) { + return true; // No limit configured + } + + // Lock-free memory acquisition using compare-exchange + size_t current = memory_used_.load(std::memory_order_relaxed); + while (true) { + if (current + bytes > max_prefetch_memory_bytes_) { + // Not enough memory - caller should queue for later + RecordTick(statistics_, PREFETCH_MEMORY_REQUESTS_BLOCKED); + return false; + } + if (memory_used_.compare_exchange_weak(current, current + bytes, + std::memory_order_release, + std::memory_order_relaxed)) { + RecordTick(statistics_, PREFETCH_MEMORY_BYTES_GRANTED, bytes); + return true; + } + // current is updated by compare_exchange_weak on failure, retry + } +} + +void IODispatcherImpl::Impl::ReleaseMemory(size_t bytes) { + if (max_prefetch_memory_bytes_ == 0) { + return; // No limit configured + } + + // Lock-free memory release using atomic fetch_sub + size_t old_val = memory_used_.fetch_sub(bytes, std::memory_order_release); + assert(old_val >= bytes); + (void)old_val; // Suppress unused warning in release builds + RecordTick(statistics_, PREFETCH_MEMORY_BYTES_RELEASED, bytes); + + // Fast-path: skip dispatch attempt if no pending requests + // This avoids mutex contention in the common single-threaded iterator case + if (!has_pending_requests_.load(std::memory_order_acquire)) { + return; + } + + // Try to dispatch pending prefetches now that memory is available + TryDispatchPendingPrefetches(); +} + +void IODispatcherImpl::Impl::TryDispatchPendingPrefetches() { + // Process pending prefetch requests - dispatch entire coalesced groups + while (true) { + std::shared_ptr pending; + + { + MutexLock lock(&memory_mutex_); + if (pending_prefetch_queue_.empty()) { + has_pending_requests_.store(false, std::memory_order_release); + return; + } + + // Get the next pending request + pending = std::move(pending_prefetch_queue_.front()); + pending_prefetch_queue_.pop_front(); + } + + // Check if the ReadSet is still alive + auto read_set = pending->read_set.lock(); + if (!read_set) { + continue; // ReadSet was destroyed, skip this request + } + + // Try to acquire memory for coalesced groups (entire groups at a time) + std::vector blocks_to_dispatch; + bool has_remaining_groups = false; + + { + MutexLock lock(&pending->groups_mutex_); + + while (!pending->coalesced_groups.empty()) { + auto& group = pending->coalesced_groups.front(); + + // Filter out blocks that were already read (not in pending set anymore) + std::vector remaining_blocks; + size_t remaining_bytes = 0; + for (size_t idx : group.block_indices) { + if (pending->block_indices_to_prefetch.count(idx) > 0) { + remaining_blocks.push_back(idx); + remaining_bytes += read_set->block_sizes_[idx]; + } + } + + // Skip empty groups (all blocks were already read) + if (remaining_blocks.empty()) { + pending->coalesced_groups.pop_front(); + continue; + } + + // Try to acquire memory for remaining blocks only + if (TryAcquireMemory(remaining_bytes)) { + // Add all remaining blocks from this group to dispatch + for (size_t idx : remaining_blocks) { + blocks_to_dispatch.push_back(idx); + pending->block_indices_to_prefetch.erase(idx); + } + pending->pending_bytes_ -= remaining_bytes; + pending->coalesced_groups.pop_front(); + } else { + // Not enough memory for this group - update with remaining blocks + group.block_indices = std::move(remaining_blocks); + group.total_bytes = remaining_bytes; + has_remaining_groups = true; + break; + } + } + } + + // Save job before potential move of pending + auto job = pending->job; + + // Requeue if groups remain + if (has_remaining_groups) { + MutexLock lock(&memory_mutex_); + pending_prefetch_queue_.push_front(std::move(pending)); + } else { + // All groups dispatched, clear pending state + read_set->pending_request_.reset(); + } + + // Clear pending flags for dispatched blocks + if (read_set->pending_prefetch_flags_) { + for (size_t idx : blocks_to_dispatch) { + if (idx < read_set->pending_prefetch_flags_size_) { + read_set->pending_prefetch_flags_[idx].store(false); + } + } + } + + // Dispatch acquired blocks + if (!blocks_to_dispatch.empty()) { + DispatchPrefetch(read_set, job, blocks_to_dispatch); + } + + // If we dispatched nothing, stop (no memory available for any group) + if (blocks_to_dispatch.empty()) { + return; + } + } +} + +void IODispatcherImpl::Impl::DispatchPrefetch( + const std::shared_ptr& read_set, const std::shared_ptr& job, + const std::vector& block_indices) { + // Sync point for testing partial prefetch - passes number of blocks being + // dispatched + TEST_SYNC_POINT_CALLBACK("IODispatcherImpl::DispatchPrefetch:BlockCount", + const_cast*>(&block_indices)); + + // Prepare and execute IO for the given blocks + std::vector read_reqs; + std::vector> coalesced_block_indices; + PrepareIORequests(job, block_indices, job->block_handles, &read_reqs, + &coalesced_block_indices); + + if (job->job_options.read_options.async_io) { + Status async_status; + std::vector fallback_indices = ExecuteAsyncIO( + job, read_set, read_reqs, coalesced_block_indices, &async_status); + + // For blocks where async is not supported, do sync IO + if (!fallback_indices.empty()) { + std::vector sync_read_reqs; + std::vector> sync_coalesced_indices; + PrepareIORequests(job, fallback_indices, job->block_handles, + &sync_read_reqs, &sync_coalesced_indices); + // Prefetch errors are ignored - user will get the error when reading + Status s = + ExecuteSyncIO(job, read_set, sync_read_reqs, sync_coalesced_indices); + s.PermitUncheckedError(); + read_set->num_sync_reads_ += fallback_indices.size(); + } + // Async errors are also ignored - user will get the error when reading + async_status.PermitUncheckedError(); + } else { + // Prefetch errors are ignored - user will get the error when reading + Status s = ExecuteSyncIO(job, read_set, read_reqs, coalesced_block_indices); + s.PermitUncheckedError(); + read_set->num_sync_reads_ += block_indices.size(); + } +} + Status IODispatcherImpl::Impl::SubmitJob(const std::shared_ptr& job, std::shared_ptr* read_set) { if (!read_set) { @@ -353,6 +653,7 @@ Status IODispatcherImpl::Impl::SubmitJob(const std::shared_ptr& job, rs->job_ = job; rs->fs_ = job->table->get_rep()->ioptions.env->GetFileSystem(); rs->pinned_blocks_.resize(job->block_handles.size()); + rs->block_sizes_.resize(job->block_handles.size(), 0); // Build sorted index for O(log n) ReadOffset lookups via binary search. // sorted_block_indices_[i] = original index of i-th smallest block by offset. @@ -399,43 +700,74 @@ Status IODispatcherImpl::Impl::SubmitJob(const std::shared_ptr& job, rs->num_cache_hits_ = job->block_handles.size() - block_indices_to_read.size(); - // Prepare read requests - coalesce adjacent blocks - std::vector read_reqs; - std::vector> coalesced_block_indices; - PrepareIORequests(job, block_indices_to_read, job->block_handles, &read_reqs, - &coalesced_block_indices); + // Calculate block sizes for uncached blocks + for (const auto& idx : block_indices_to_read) { + size_t block_size = + BlockBasedTable::BlockSizeWithTrailer(job->block_handles[idx]); + rs->block_sizes_[idx] = block_size; + } - // Step 3: Execute IO requests based on JobOptions - if (job->job_options.read_options.async_io) { - // Try async IO - get back any blocks that need sync fallback (not - // supported) and surface any actual errors to caller - Status async_status; - std::vector fallback_indices = ExecuteAsyncIO( - job, rs, read_reqs, coalesced_block_indices, &async_status); - if (!async_status.ok()) { - return async_status; - } + // Store dispatcher reference for release callbacks + rs->dispatcher_data_ = shared_from_this(); - // Fall back to sync IO for blocks where async is not supported - if (!fallback_indices.empty()) { - std::vector sync_read_reqs; - std::vector> sync_coalesced_indices; - PrepareIORequests(job, fallback_indices, job->block_handles, - &sync_read_reqs, &sync_coalesced_indices); + // Pre-coalesce blocks into groups, respecting memory budget per group + // This ensures we dispatch meaningful IO sizes, not tiny single-block IOs + // Both memory-limited and non-memory-limited paths use the same coalescing + auto coalesced_groups = PreCoalesceBlocks(job, rs, block_indices_to_read, + max_prefetch_memory_bytes_); - Status s = ExecuteSyncIO(job, rs, sync_read_reqs, sync_coalesced_indices); - if (!s.ok()) { - return s; + std::vector blocks_to_dispatch; + std::deque groups_to_queue; + + // Try to acquire memory for entire coalesced groups + for (auto& group : coalesced_groups) { + if (TryAcquireMemory(group.total_bytes)) { + // Add all blocks from this group to dispatch + for (size_t idx : group.block_indices) { + blocks_to_dispatch.push_back(idx); } - rs->num_sync_reads_ += fallback_indices.size(); + } else { + // Queue this group for later + groups_to_queue.push_back(std::move(group)); } - } else { - Status s = ExecuteSyncIO(job, rs, read_reqs, coalesced_block_indices); - if (!s.ok()) { - return s; + } + + // Dispatch acquired blocks immediately + if (!blocks_to_dispatch.empty()) { + DispatchPrefetch(rs, job, blocks_to_dispatch); + } + + // Queue remaining groups for later (only applies when memory limiting) + if (!groups_to_queue.empty()) { + auto pending = std::make_shared(); + pending->read_set = rs; + pending->job = job; + + size_t pending_bytes = 0; + for (const auto& group : groups_to_queue) { + for (size_t idx : group.block_indices) { + pending->block_indices_to_prefetch.insert(idx); + } + pending_bytes += group.total_bytes; + } + pending->coalesced_groups = std::move(groups_to_queue); + pending->pending_bytes_ = pending_bytes; + + // Set up pending flags for queued blocks only + size_t num_blocks = job->block_handles.size(); + rs->pending_prefetch_flags_ = + std::make_unique[]>(num_blocks); + rs->pending_prefetch_flags_size_ = num_blocks; + for (size_t idx : pending->block_indices_to_prefetch) { + rs->pending_prefetch_flags_[idx].store(true); + } + rs->pending_request_ = pending; + + { + MutexLock lock(&memory_mutex_); + pending_prefetch_queue_.push_back(std::move(pending)); + has_pending_requests_.store(true, std::memory_order_release); } - // We bump this for sync reads - rs->num_sync_reads_ += block_indices_to_read.size(); } *read_set = std::move(rs); @@ -503,6 +835,67 @@ void IODispatcherImpl::Impl::PrepareIORequests( } } +std::vector IODispatcherImpl::Impl::PreCoalesceBlocks( + const std::shared_ptr& job, const std::shared_ptr& rs, + const std::vector& block_indices, size_t max_group_bytes) { + std::vector groups; + + if (block_indices.empty()) { + return groups; + } + + const auto& block_handles = job->block_handles; + const uint64_t coalesce_threshold = job->job_options.io_coalesce_threshold; + + // Sort block indices by offset for coalescing + std::vector sorted_indices = block_indices; + std::sort(sorted_indices.begin(), sorted_indices.end(), + [&block_handles](size_t a, size_t b) { + return block_handles[a].offset() < block_handles[b].offset(); + }); + + // Build coalesced groups respecting max_group_bytes + groups.emplace_back(); + + for (size_t idx : sorted_indices) { + size_t block_size = rs->block_sizes_[idx]; + + // Skip blocks that are individually larger than the memory budget + // These will be read synchronously when needed (via ReadIndex fallback) + if (max_group_bytes > 0 && block_size > max_group_bytes) { + continue; + } + + // Check if we need to start a new group + bool start_new_group = false; + + if (!groups.back().block_indices.empty()) { + // Check gap with previous block + size_t last_idx = groups.back().block_indices.back(); + const auto& last_handle = block_handles[last_idx]; + uint64_t last_end = last_handle.offset() + + BlockBasedTable::BlockSizeWithTrailer(last_handle); + uint64_t current_start = block_handles[idx].offset(); + + if (current_start > last_end + coalesce_threshold) { + start_new_group = true; // Gap too large + } else if (max_group_bytes > 0 && + groups.back().total_bytes + block_size > max_group_bytes) { + start_new_group = true; // Would exceed size limit + } + } + + if (start_new_group) { + groups.emplace_back(); + } + + groups.back().block_indices.push_back(idx); + groups.back().total_bytes += block_size; + } + + return groups; +} + std::vector IODispatcherImpl::Impl::ExecuteAsyncIO( const std::shared_ptr& job, const std::shared_ptr& read_set, std::vector& read_reqs, @@ -648,7 +1041,11 @@ Status IODispatcherImpl::Impl::ExecuteSyncIO( return Status::OK(); } -IODispatcherImpl::IODispatcherImpl() : impl_(new Impl()) {} +IODispatcherImpl::IODispatcherImpl() + : impl_(std::make_shared(IODispatcherOptions())) {} + +IODispatcherImpl::IODispatcherImpl(const IODispatcherOptions& options) + : impl_(std::make_shared(options)) {} IODispatcherImpl::~IODispatcherImpl() = default; @@ -659,4 +1056,8 @@ Status IODispatcherImpl::SubmitJob(const std::shared_ptr& job, IODispatcher* NewIODispatcher() { return new IODispatcherImpl(); } +IODispatcher* NewIODispatcher(const IODispatcherOptions& options) { + return new IODispatcherImpl(options); +} + } // namespace ROCKSDB_NAMESPACE diff --git a/util/io_dispatcher_imp.h b/util/io_dispatcher_imp.h index 3324705ada..c4e52b86d5 100644 --- a/util/io_dispatcher_imp.h +++ b/util/io_dispatcher_imp.h @@ -21,7 +21,8 @@ namespace ROCKSDB_NAMESPACE { class IODispatcherImpl : public IODispatcher { public: - explicit IODispatcherImpl(); + IODispatcherImpl(); + explicit IODispatcherImpl(const IODispatcherOptions& options); ~IODispatcherImpl() override; Status SubmitJob(const std::shared_ptr& job, @@ -29,7 +30,7 @@ class IODispatcherImpl : public IODispatcher { private: struct Impl; - std::unique_ptr impl_; + std::shared_ptr impl_; }; } // namespace ROCKSDB_NAMESPACE diff --git a/util/io_dispatcher_test.cc b/util/io_dispatcher_test.cc index 7f8e0a9311..d5ea665e5c 100644 --- a/util/io_dispatcher_test.cc +++ b/util/io_dispatcher_test.cc @@ -12,6 +12,7 @@ #include #include +#include #include "db/db_test_util.h" #include "db/dbformat.h" @@ -23,6 +24,7 @@ #include "table/block_based/block_based_table_builder.h" #include "table/block_based/block_based_table_factory.h" #include "table/block_based/block_based_table_reader.h" +#include "test_util/sync_point.h" // Enable io_uring support for this test extern "C" bool RocksDbIOUringEnable() { return true; } @@ -896,6 +898,900 @@ TEST_F(IODispatcherTest, VerifyReadRequestDetails) { } } +// Test that memory limiting blocks when the limit is exceeded +TEST_F(IODispatcherTest, MemoryLimitBlocksWhenExceeded) { + // Create dispatcher with a small memory limit (1MB) + IODispatcherOptions opts; + opts.max_prefetch_memory_bytes = 1 * 1024 * 1024; // 1MB + std::unique_ptr dispatcher(NewIODispatcher(opts)); + + std::unique_ptr table; + std::vector block_handles; + Status s = CreateAndOpenSST(50, &table, &block_handles); + ASSERT_OK(s); + ASSERT_GT(block_handles.size(), 0); + + // Submit a job - should succeed immediately (non-blocking) + auto job = std::make_shared(); + job->block_handles = block_handles; + job->table = table.get(); + job->job_options.read_options.async_io = false; + + std::shared_ptr read_set; + s = dispatcher->SubmitJob(job, &read_set); + ASSERT_OK(s); + ASSERT_NE(read_set, nullptr); + + // Read all blocks - they may be read synchronously if prefetch was deferred + for (size_t i = 0; i < block_handles.size(); ++i) { + CachableEntry block; + Status read_status = read_set->ReadIndex(i, &block); + ASSERT_OK(read_status); + ASSERT_NE(block.GetValue(), nullptr); + } +} + +// Test that SubmitJob never blocks even when memory is exhausted +TEST_F(IODispatcherTest, SubmitJobNeverBlocks) { + // Create dispatcher with a tiny memory limit + IODispatcherOptions opts; + opts.max_prefetch_memory_bytes = 1024; // 1KB - very small + std::unique_ptr dispatcher(NewIODispatcher(opts)); + + std::unique_ptr table; + std::vector block_handles; + Status s = CreateAndOpenSST(50, &table, &block_handles); + ASSERT_OK(s); + ASSERT_GT(block_handles.size(), 0); + + // Submit first job - uses up all memory + auto job1 = std::make_shared(); + job1->block_handles = block_handles; + job1->table = table.get(); + job1->job_options.read_options.async_io = false; + + std::shared_ptr read_set1; + s = dispatcher->SubmitJob(job1, &read_set1); + ASSERT_OK(s); // Should succeed immediately + + // Submit second job - should also succeed immediately (not block) + std::unique_ptr table2; + std::vector block_handles2; + s = CreateAndOpenSST(30, &table2, &block_handles2); + ASSERT_OK(s); + + auto job2 = std::make_shared(); + job2->block_handles = block_handles2; + job2->table = table2.get(); + job2->job_options.read_options.async_io = false; + + std::shared_ptr read_set2; + s = dispatcher->SubmitJob(job2, &read_set2); + ASSERT_OK(s); // Should succeed immediately - prefetch is just deferred + + // Reads work - blocks are fetched synchronously on demand + for (size_t i = 0; i < block_handles2.size(); ++i) { + CachableEntry block; + Status read_status = read_set2->ReadIndex(i, &block); + ASSERT_OK(read_status); + ASSERT_NE(block.GetValue(), nullptr); + } +} + +// Test that releasing blocks triggers pending prefetches +TEST_F(IODispatcherTest, BlockReleaseTriggersWaitingJob) { + // Create dispatcher with a small memory limit + IODispatcherOptions opts; + opts.max_prefetch_memory_bytes = 100 * 1024; // 100KB + std::unique_ptr dispatcher(NewIODispatcher(opts)); + + std::unique_ptr table; + std::vector block_handles; + Status s = CreateAndOpenSST(30, &table, &block_handles); + ASSERT_OK(s); + ASSERT_GT(block_handles.size(), 0); + + // Submit first job + auto job1 = std::make_shared(); + job1->block_handles = block_handles; + job1->table = table.get(); + job1->job_options.read_options.async_io = false; + + std::shared_ptr read_set1; + s = dispatcher->SubmitJob(job1, &read_set1); + ASSERT_OK(s); + ASSERT_NE(read_set1, nullptr); + + // Read all blocks from first job + for (size_t i = 0; i < block_handles.size(); ++i) { + CachableEntry block; + Status read_status = read_set1->ReadIndex(i, &block); + ASSERT_OK(read_status); + } + + // Submit second job - prefetch will be deferred due to memory limit + std::unique_ptr table2; + std::vector block_handles2; + s = CreateAndOpenSST(20, &table2, &block_handles2); + ASSERT_OK(s); + + auto job2 = std::make_shared(); + job2->block_handles = block_handles2; + job2->table = table2.get(); + job2->job_options.read_options.async_io = false; + + std::shared_ptr read_set2; + s = dispatcher->SubmitJob(job2, &read_set2); + ASSERT_OK(s); // Should succeed immediately + ASSERT_NE(read_set2, nullptr); + + // Release blocks from first job - this should trigger pending prefetches + for (size_t i = 0; i < block_handles.size(); ++i) { + read_set1->ReleaseBlock(i); + } + + // Read all blocks from second job - should work + for (size_t i = 0; i < block_handles2.size(); ++i) { + CachableEntry block; + Status read_status = read_set2->ReadIndex(i, &block); + ASSERT_OK(read_status); + ASSERT_NE(block.GetValue(), nullptr); + } +} + +// Test that multiple ReadSets share the memory budget +TEST_F(IODispatcherTest, MultipleReadSetsShareMemoryBudget) { + IODispatcherOptions opts; + opts.max_prefetch_memory_bytes = 10 * 1024 * 1024; // 10MB + std::unique_ptr dispatcher(NewIODispatcher(opts)); + + std::vector> read_sets; + std::vector> all_block_handles; + + // Create and submit multiple jobs + for (int i = 0; i < 3; i++) { + std::unique_ptr table; + std::vector block_handles; + + Status s = CreateAndOpenSST(20 + i * 5, &table, &block_handles); + ASSERT_OK(s); + + auto job = std::make_shared(); + job->block_handles = block_handles; + job->table = table.get(); + job->job_options.read_options.async_io = false; + tables_.push_back(std::move(table)); + + all_block_handles.push_back(block_handles); + std::shared_ptr read_set; + s = dispatcher->SubmitJob(job, &read_set); + ASSERT_OK(s); + read_sets.push_back(read_set); + } + + // Verify all ReadSets can read their blocks + for (size_t i = 0; i < read_sets.size(); ++i) { + for (size_t j = 0; j < all_block_handles[i].size(); ++j) { + CachableEntry block; + Status read_status = read_sets[i]->ReadIndex(j, &block); + ASSERT_OK(read_status); + ASSERT_NE(block.GetValue(), nullptr); + } + } + + // Release all blocks from first ReadSet + for (size_t i = 0; i < all_block_handles[0].size(); ++i) { + read_sets[0]->ReleaseBlock(i); + } + + // Create another job - should work because first ReadSet released memory + std::unique_ptr table_new; + std::vector block_handles_new; + Status s = CreateAndOpenSST(25, &table_new, &block_handles_new); + ASSERT_OK(s); + + auto job_new = std::make_shared(); + job_new->block_handles = block_handles_new; + job_new->table = table_new.get(); + job_new->job_options.read_options.async_io = false; + + std::shared_ptr read_set_new; + s = dispatcher->SubmitJob(job_new, &read_set_new); + ASSERT_OK(s); + ASSERT_NE(read_set_new, nullptr); + + for (size_t i = 0; i < block_handles_new.size(); ++i) { + CachableEntry block; + Status read_status = read_set_new->ReadIndex(i, &block); + ASSERT_OK(read_status); + ASSERT_NE(block.GetValue(), nullptr); + } +} + +// Test that no memory limiting is applied when max_prefetch_memory_bytes is 0 +TEST_F(IODispatcherTest, NoMemoryLimitWhenZero) { + IODispatcherOptions opts; + opts.max_prefetch_memory_bytes = 0; // No limit + std::unique_ptr dispatcher(NewIODispatcher(opts)); + + std::unique_ptr table; + std::vector block_handles; + Status s = CreateAndOpenSST(50, &table, &block_handles); + ASSERT_OK(s); + + auto job = std::make_shared(); + job->block_handles = block_handles; + job->table = table.get(); + job->job_options.read_options.async_io = false; + + std::shared_ptr read_set; + s = dispatcher->SubmitJob(job, &read_set); + ASSERT_OK(s); + ASSERT_NE(read_set, nullptr); + + for (size_t i = 0; i < block_handles.size(); ++i) { + CachableEntry block; + Status read_status = read_set->ReadIndex(i, &block); + ASSERT_OK(read_status); + ASSERT_NE(block.GetValue(), nullptr); + } +} + +// Test memory release on ReadSet destruction triggers pending prefetches +TEST_F(IODispatcherTest, MemoryReleasedOnReadSetDestruction) { + IODispatcherOptions opts; + opts.max_prefetch_memory_bytes = 100 * 1024; // 100KB + std::unique_ptr dispatcher(NewIODispatcher(opts)); + + // Create table outside the scope so it outlives the ReadSet + std::unique_ptr table; + std::vector block_handles; + Status s = CreateAndOpenSST(30, &table, &block_handles); + ASSERT_OK(s); + + // Second table - created now so it's available after first ReadSet is + // destroyed + std::unique_ptr table2; + std::vector block_handles2; + s = CreateAndOpenSST(30, &table2, &block_handles2); + ASSERT_OK(s); + + std::shared_ptr read_set2; + + { + auto job = std::make_shared(); + job->block_handles = block_handles; + job->table = table.get(); + job->job_options.read_options.async_io = false; + + std::shared_ptr read_set; + s = dispatcher->SubmitJob(job, &read_set); + ASSERT_OK(s); + ASSERT_NE(read_set, nullptr); + + // Submit second job while first is still alive - prefetch will be deferred + auto job2 = std::make_shared(); + job2->block_handles = block_handles2; + job2->table = table2.get(); + job2->job_options.read_options.async_io = false; + + s = dispatcher->SubmitJob(job2, &read_set2); + ASSERT_OK(s); // Should succeed immediately + ASSERT_NE(read_set2, nullptr); + + // First ReadSet goes out of scope here and should release all memory, + // which triggers pending prefetches for second ReadSet + } + + // Read all blocks from second job - should work because first ReadSet + // released its memory on destruction + for (size_t i = 0; i < block_handles2.size(); ++i) { + CachableEntry block; + Status read_status = read_set2->ReadIndex(i, &block); + ASSERT_OK(read_status); + ASSERT_NE(block.GetValue(), nullptr); + } +} + +// Test that partial prefetch dispatches as many blocks as memory allows +// and queues the rest for later dispatch +TEST_F(IODispatcherTest, PartialPrefetchDispatchesWhatFits) { + // Skip this test if io_uring is not available since partial prefetch + // only applies to async IO + if (!kIOUringPresent) { + return; // io_uring not available, skip async IO test + } + + // Create dispatcher with memory limit that allows only some blocks + // Each block is ~16KB, so 50KB allows roughly 3 blocks + IODispatcherOptions opts; + opts.max_prefetch_memory_bytes = 50 * 1024; // 50KB + std::unique_ptr dispatcher(NewIODispatcher(opts)); + + std::unique_ptr table; + std::vector block_handles; + // Create 10 blocks - only ~3 should fit in memory + Status s = CreateAndOpenSST(10, &table, &block_handles); + ASSERT_OK(s); + ASSERT_GE(block_handles.size(), 5); + + // Use sync point to count blocks dispatched during SubmitJob + size_t blocks_dispatched_on_submit = 0; + SyncPoint::GetInstance()->SetCallBack( + "IODispatcherImpl::DispatchPrefetch:BlockCount", [&](void* arg) { + auto* indices = static_cast*>(arg); + blocks_dispatched_on_submit += indices->size(); + }); + SyncPoint::GetInstance()->EnableProcessing(); + + auto job = std::make_shared(); + job->block_handles = block_handles; + job->table = table.get(); + job->job_options.read_options.async_io = true; // Use async IO + + std::shared_ptr read_set; + s = dispatcher->SubmitJob(job, &read_set); + ASSERT_OK(s); + ASSERT_NE(read_set, nullptr); + + // With partial prefetch, we expect SOME blocks to have been dispatched + // (the ones that fit in memory), but not ALL blocks + // This is the key assertion: partial prefetch means > 0 blocks dispatched + // even though total memory needed exceeds the limit + EXPECT_GT(blocks_dispatched_on_submit, 0) + << "Expected some blocks to be dispatched with partial prefetch"; + EXPECT_LT(blocks_dispatched_on_submit, block_handles.size()) + << "Expected not all blocks to be dispatched (memory limit should apply)"; + + SyncPoint::GetInstance()->DisableProcessing(); + SyncPoint::GetInstance()->ClearAllCallBacks(); + + // Now read all blocks - remaining blocks will be fetched on demand + for (size_t i = 0; i < block_handles.size(); ++i) { + CachableEntry block; + Status read_status = read_set->ReadIndex(i, &block); + ASSERT_OK(read_status); + ASSERT_NE(block.GetValue(), nullptr); + } + + // Verify all blocks were ultimately read + uint64_t total_reads = read_set->GetNumSyncReads() + + read_set->GetNumAsyncReads() + + read_set->GetNumCacheHits(); + EXPECT_EQ(total_reads, block_handles.size()); +} + +// Test that earlier block indices are prioritized in partial prefetch +TEST_F(IODispatcherTest, PartialPrefetchPrioritizesEarlierIndices) { + // Skip this test if io_uring is not available + if (!kIOUringPresent) { + return; // io_uring not available, skip async IO test + } + + // Create dispatcher with memory limit that allows only 1-2 blocks + IODispatcherOptions opts; + opts.max_prefetch_memory_bytes = 20 * 1024; // 20KB - room for ~1 block + std::unique_ptr dispatcher(NewIODispatcher(opts)); + + std::unique_ptr table; + std::vector block_handles; + Status s = CreateAndOpenSST(10, &table, &block_handles); + ASSERT_OK(s); + ASSERT_GE(block_handles.size(), 5); + + tracking_fs_->ClearReadOps(); + + auto job = std::make_shared(); + job->block_handles = block_handles; + job->table = table.get(); + job->job_options.read_options.async_io = true; + + std::shared_ptr read_set; + s = dispatcher->SubmitJob(job, &read_set); + ASSERT_OK(s); + + // Get the async reads that were dispatched + auto read_ops = tracking_fs_->GetReadOps(); + + // Find the offset of the first async read + uint64_t first_async_offset = UINT64_MAX; + for (const auto& op : read_ops) { + if (op.type == ReadOp::kReadAsync && !op.requests.empty()) { + first_async_offset = std::min(first_async_offset, op.requests[0].first); + } + } + + // The first async read should be for the first block (lowest offset) + // This verifies that earlier indices are prioritized + if (first_async_offset != UINT64_MAX) { + EXPECT_EQ(first_async_offset, block_handles[0].offset()) + << "Expected first async read to be for the first block (earliest " + "index)"; + } + + // Read all blocks to complete the test + for (size_t i = 0; i < block_handles.size(); ++i) { + CachableEntry block; + Status read_status = read_set->ReadIndex(i, &block); + ASSERT_OK(read_status); + ASSERT_NE(block.GetValue(), nullptr); + } +} + +// Test that blocks larger than the memory budget are excluded from prefetch +// and fall back to synchronous read +TEST_F(IODispatcherTest, OversizedBlocksFallbackToSyncRead) { + // Skip this test if io_uring is not available since we need async IO + if (!kIOUringPresent) { + return; + } + + std::unique_ptr table; + std::vector block_handles; + Status s = CreateAndOpenSST(10, &table, &block_handles); + ASSERT_OK(s); + ASSERT_GE(block_handles.size(), 3); + + // Calculate the size of a single block + size_t single_block_size = + BlockBasedTable::BlockSizeWithTrailer(block_handles[0]); + + // Create dispatcher with memory limit smaller than a single block + // This means ALL blocks are "oversized" and should fall back to sync read + IODispatcherOptions opts; + opts.max_prefetch_memory_bytes = single_block_size / 2; // Half a block + std::unique_ptr dispatcher(NewIODispatcher(opts)); + + // Track dispatches - with oversized blocks, nothing should be dispatched + size_t blocks_dispatched = 0; + SyncPoint::GetInstance()->SetCallBack( + "IODispatcherImpl::DispatchPrefetch:BlockCount", [&](void* arg) { + auto* indices = static_cast*>(arg); + blocks_dispatched += indices->size(); + }); + SyncPoint::GetInstance()->EnableProcessing(); + + auto job = std::make_shared(); + job->block_handles = block_handles; + job->table = table.get(); + job->job_options.read_options.async_io = true; + + std::shared_ptr read_set; + s = dispatcher->SubmitJob(job, &read_set); + ASSERT_OK(s); + ASSERT_NE(read_set, nullptr); + + // No blocks should have been dispatched since they're all oversized + EXPECT_EQ(blocks_dispatched, 0) + << "Expected no blocks to be dispatched when all blocks are oversized"; + + SyncPoint::GetInstance()->DisableProcessing(); + SyncPoint::GetInstance()->ClearAllCallBacks(); + + // All blocks should still be readable via sync fallback + for (size_t i = 0; i < block_handles.size(); ++i) { + CachableEntry block; + Status read_status = read_set->ReadIndex(i, &block); + ASSERT_OK(read_status); + ASSERT_NE(block.GetValue(), nullptr); + } + + // All reads should be sync since blocks couldn't be prefetched + EXPECT_GT(read_set->GetNumSyncReads(), 0) + << "Expected sync reads for oversized blocks"; +} + +// Test that reading blocks before prefetch dispatch correctly updates +// memory accounting for coalesced groups +TEST_F(IODispatcherTest, PartialReadsUpdateCoalescedGroups) { + // Skip this test if io_uring is not available + if (!kIOUringPresent) { + return; + } + + // Create dispatcher with memory limit that allows only some blocks + IODispatcherOptions opts; + opts.max_prefetch_memory_bytes = 50 * 1024; // 50KB + std::unique_ptr dispatcher(NewIODispatcher(opts)); + + std::unique_ptr table; + std::vector block_handles; + Status s = CreateAndOpenSST(20, &table, &block_handles); + ASSERT_OK(s); + ASSERT_GE(block_handles.size(), 10); + + auto job = std::make_shared(); + job->block_handles = block_handles; + job->table = table.get(); + job->job_options.read_options.async_io = true; + + std::shared_ptr read_set; + s = dispatcher->SubmitJob(job, &read_set); + ASSERT_OK(s); + ASSERT_NE(read_set, nullptr); + + // Read some blocks directly (simulating on-demand access before prefetch) + // This removes them from pending and should update coalesced group accounting + for (size_t i = 0; i < 5 && i < block_handles.size(); ++i) { + CachableEntry block; + Status read_status = read_set->ReadIndex(i, &block); + ASSERT_OK(read_status); + ASSERT_NE(block.GetValue(), nullptr); + } + + // Release the blocks we read - this frees memory + for (size_t i = 0; i < 5 && i < block_handles.size(); ++i) { + read_set->ReleaseBlock(i); + } + + // Now read the remaining blocks - these should work correctly + // The key test: memory accounting should be correct even though some blocks + // were removed from pending groups before dispatch + for (size_t i = 5; i < block_handles.size(); ++i) { + CachableEntry block; + Status read_status = read_set->ReadIndex(i, &block); + ASSERT_OK(read_status) << "Failed to read block " << i; + ASSERT_NE(block.GetValue(), nullptr) << "Block " << i << " is null"; + } + + // Verify all remaining blocks were read successfully + uint64_t total_reads = read_set->GetNumSyncReads() + + read_set->GetNumAsyncReads() + + read_set->GetNumCacheHits(); + // We read 5 blocks initially, then the remaining blocks + EXPECT_GE(total_reads, block_handles.size() - 5) + << "Expected at least the remaining blocks to be counted"; +} + +// Test that a mix of oversized and normal blocks works correctly +TEST_F(IODispatcherTest, MixedOversizedAndNormalBlocks) { + // Skip this test if io_uring is not available + if (!kIOUringPresent) { + return; + } + + std::unique_ptr table; + std::vector block_handles; + Status s = CreateAndOpenSST(10, &table, &block_handles); + ASSERT_OK(s); + ASSERT_GE(block_handles.size(), 5); + + // Calculate the size of a typical block + size_t typical_block_size = + BlockBasedTable::BlockSizeWithTrailer(block_handles[0]); + + // Create dispatcher with memory limit that allows exactly 2 typical blocks + // This means groups of 3+ blocks become "oversized" as a group + IODispatcherOptions opts; + opts.max_prefetch_memory_bytes = typical_block_size * 2; + std::unique_ptr dispatcher(NewIODispatcher(opts)); + + auto job = std::make_shared(); + job->block_handles = block_handles; + job->table = table.get(); + job->job_options.read_options.async_io = true; + + std::shared_ptr read_set; + s = dispatcher->SubmitJob(job, &read_set); + ASSERT_OK(s); + ASSERT_NE(read_set, nullptr); + + // All blocks should be readable regardless of prefetch status + for (size_t i = 0; i < block_handles.size(); ++i) { + CachableEntry block; + Status read_status = read_set->ReadIndex(i, &block); + ASSERT_OK(read_status) << "Failed to read block " << i; + ASSERT_NE(block.GetValue(), nullptr) << "Block " << i << " is null"; + } + + // Verify total reads match + uint64_t total_reads = read_set->GetNumSyncReads() + + read_set->GetNumAsyncReads() + + read_set->GetNumCacheHits(); + EXPECT_EQ(total_reads, block_handles.size()); +} + +// Test that memory is properly accounted when groups are partially consumed +TEST_F(IODispatcherTest, MemoryAccountingWithPartialGroupConsumption) { + // Skip this test if io_uring is not available + if (!kIOUringPresent) { + return; + } + + // Create dispatcher with a specific memory limit + IODispatcherOptions opts; + opts.max_prefetch_memory_bytes = 100 * 1024; // 100KB + std::unique_ptr dispatcher(NewIODispatcher(opts)); + + std::unique_ptr table; + std::vector block_handles; + Status s = CreateAndOpenSST(30, &table, &block_handles); + ASSERT_OK(s); + ASSERT_GE(block_handles.size(), 10); + + auto job = std::make_shared(); + job->block_handles = block_handles; + job->table = table.get(); + job->job_options.read_options.async_io = true; + + std::shared_ptr read_set; + s = dispatcher->SubmitJob(job, &read_set); + ASSERT_OK(s); + ASSERT_NE(read_set, nullptr); + + // Read blocks one at a time and release them + // This tests that RemoveFromPending correctly updates pending state + // and that TryDispatchPendingPrefetches filters correctly + for (size_t i = 0; i < block_handles.size(); ++i) { + CachableEntry block; + Status read_status = read_set->ReadIndex(i, &block); + ASSERT_OK(read_status) << "Failed to read block " << i; + ASSERT_NE(block.GetValue(), nullptr) << "Block " << i << " is null"; + + // Release the block immediately after reading + read_set->ReleaseBlock(i); + } + + // Verify total reads match + uint64_t total_reads = read_set->GetNumSyncReads() + + read_set->GetNumAsyncReads() + + read_set->GetNumCacheHits(); + EXPECT_EQ(total_reads, block_handles.size()); +} + +// Test that sync prefetching respects memory limits +TEST_F(IODispatcherTest, SyncPrefetchWithMemoryLimit) { + // Create dispatcher with a small memory limit + IODispatcherOptions opts; + opts.max_prefetch_memory_bytes = 50 * 1024; // 50KB + std::unique_ptr dispatcher(NewIODispatcher(opts)); + + std::unique_ptr table; + std::vector block_handles; + Status s = CreateAndOpenSST(20, &table, &block_handles); + ASSERT_OK(s); + ASSERT_GE(block_handles.size(), 10); + + auto job = std::make_shared(); + job->block_handles = block_handles; + job->table = table.get(); + job->job_options.read_options.async_io = false; // Sync IO + + std::shared_ptr read_set; + s = dispatcher->SubmitJob(job, &read_set); + ASSERT_OK(s); + ASSERT_NE(read_set, nullptr); + + // All blocks should be readable even with memory limits + for (size_t i = 0; i < block_handles.size(); ++i) { + CachableEntry block; + Status read_status = read_set->ReadIndex(i, &block); + ASSERT_OK(read_status) << "Failed to read block " << i; + ASSERT_NE(block.GetValue(), nullptr) << "Block " << i << " is null"; + } + + // Verify all were sync reads + EXPECT_GT(read_set->GetNumSyncReads(), 0) + << "Expected sync reads with async_io=false"; + EXPECT_EQ(read_set->GetNumAsyncReads(), 0) + << "Expected no async reads with async_io=false"; +} + +// Test that oversized blocks work correctly with sync IO +TEST_F(IODispatcherTest, OversizedBlocksWithSyncIO) { + std::unique_ptr table; + std::vector block_handles; + Status s = CreateAndOpenSST(10, &table, &block_handles); + ASSERT_OK(s); + ASSERT_GE(block_handles.size(), 3); + + // Calculate the size of a single block + size_t single_block_size = + BlockBasedTable::BlockSizeWithTrailer(block_handles[0]); + + // Create dispatcher with memory limit smaller than a single block + // This means ALL blocks are "oversized" + IODispatcherOptions opts; + opts.max_prefetch_memory_bytes = single_block_size / 2; // Half a block + std::unique_ptr dispatcher(NewIODispatcher(opts)); + + auto job = std::make_shared(); + job->block_handles = block_handles; + job->table = table.get(); + job->job_options.read_options.async_io = false; // Sync IO + + std::shared_ptr read_set; + s = dispatcher->SubmitJob(job, &read_set); + ASSERT_OK(s); + ASSERT_NE(read_set, nullptr); + + // All blocks should still be readable via sync fallback + for (size_t i = 0; i < block_handles.size(); ++i) { + CachableEntry block; + Status read_status = read_set->ReadIndex(i, &block); + ASSERT_OK(read_status) << "Failed to read block " << i; + ASSERT_NE(block.GetValue(), nullptr) << "Block " << i << " is null"; + } + + // All reads should be sync + EXPECT_GT(read_set->GetNumSyncReads(), 0) + << "Expected sync reads for oversized blocks"; +} + +// Test that a single block larger than total memory budget still works +TEST_F(IODispatcherTest, SingleBlockLargerThanTotalMemory) { + std::unique_ptr table; + std::vector block_handles; + Status s = CreateAndOpenSST(5, &table, &block_handles); + ASSERT_OK(s); + ASSERT_GE(block_handles.size(), 1); + + // Set memory limit to 1 byte - smaller than any block + IODispatcherOptions opts; + opts.max_prefetch_memory_bytes = 1; + std::unique_ptr dispatcher(NewIODispatcher(opts)); + + // Test with both sync and async modes + for (bool async : {false, true}) { + // Skip async if io_uring not available + if (async && !kIOUringPresent) { + continue; + } + + auto job = std::make_shared(); + job->block_handles = block_handles; + job->table = table.get(); + job->job_options.read_options.async_io = async; + + std::shared_ptr read_set; + s = dispatcher->SubmitJob(job, &read_set); + ASSERT_OK(s) << "SubmitJob failed with async=" << async; + ASSERT_NE(read_set, nullptr); + + // All blocks should be readable + for (size_t i = 0; i < block_handles.size(); ++i) { + CachableEntry block; + Status read_status = read_set->ReadIndex(i, &block); + ASSERT_OK(read_status) + << "Failed to read block " << i << " with async=" << async; + ASSERT_NE(block.GetValue(), nullptr) + << "Block " << i << " is null with async=" << async; + } + } +} + +// Test that sync prefetching defers later groups and dispatches them +// when memory is released +TEST_F(IODispatcherTest, SyncPrefetchDefersAndDispatchesLaterGroups) { + std::unique_ptr table; + std::vector block_handles; + // Create 10+ blocks so we have enough to test deferred dispatch + Status s = CreateAndOpenSST(20, &table, &block_handles); + ASSERT_OK(s); + ASSERT_GE(block_handles.size(), 10); + + // Calculate typical block size + size_t typical_block_size = + BlockBasedTable::BlockSizeWithTrailer(block_handles[0]); + + // Set memory limit to fit approximately 3 blocks + // This should cause groups to be split and some deferred + IODispatcherOptions opts; + opts.max_prefetch_memory_bytes = typical_block_size * 3; + std::unique_ptr dispatcher(NewIODispatcher(opts)); + + // Track dispatch calls + std::vector dispatch_counts; + SyncPoint::GetInstance()->SetCallBack( + "IODispatcherImpl::DispatchPrefetch:BlockCount", [&](void* arg) { + auto* indices = static_cast*>(arg); + dispatch_counts.push_back(indices->size()); + }); + SyncPoint::GetInstance()->EnableProcessing(); + + auto job = std::make_shared(); + job->block_handles = block_handles; + job->table = table.get(); + job->job_options.read_options.async_io = false; // Sync IO + + std::shared_ptr read_set; + s = dispatcher->SubmitJob(job, &read_set); + ASSERT_OK(s); + ASSERT_NE(read_set, nullptr); + + // After SubmitJob, some blocks should have been dispatched (first group) + // and remaining groups should be queued + size_t initial_dispatch_count = dispatch_counts.size(); + EXPECT_GT(initial_dispatch_count, 0) + << "Expected at least one dispatch during SubmitJob"; + + // Read and release first few blocks - this should trigger deferred dispatch + for (size_t i = 0; i < 3 && i < block_handles.size(); ++i) { + CachableEntry block; + Status read_status = read_set->ReadIndex(i, &block); + ASSERT_OK(read_status); + ASSERT_NE(block.GetValue(), nullptr); + // Release to free memory + read_set->ReleaseBlock(i); + } + + // After releasing blocks, more dispatches should have occurred + // as the pending queue gets processed + size_t dispatch_count_after_release = dispatch_counts.size(); + EXPECT_GE(dispatch_count_after_release, initial_dispatch_count) + << "Expected more dispatches after releasing blocks"; + + SyncPoint::GetInstance()->DisableProcessing(); + SyncPoint::GetInstance()->ClearAllCallBacks(); + + // All remaining blocks should still be readable + for (size_t i = 3; i < block_handles.size(); ++i) { + CachableEntry block; + Status read_status = read_set->ReadIndex(i, &block); + ASSERT_OK(read_status) << "Failed to read block " << i; + ASSERT_NE(block.GetValue(), nullptr) << "Block " << i << " is null"; + } +} + +// Test that coalesced groups are properly split based on memory budget +TEST_F(IODispatcherTest, CoalescedGroupsSplitByMemoryBudget) { + std::unique_ptr table; + std::vector block_handles; + Status s = CreateAndOpenSST(15, &table, &block_handles); + ASSERT_OK(s); + ASSERT_GE(block_handles.size(), 10); + + // Calculate typical block size + size_t typical_block_size = + BlockBasedTable::BlockSizeWithTrailer(block_handles[0]); + + // Set memory limit to fit exactly 5 blocks + // With 10+ blocks, we should get at least 2 groups + IODispatcherOptions opts; + opts.max_prefetch_memory_bytes = typical_block_size * 5; + std::unique_ptr dispatcher(NewIODispatcher(opts)); + + // Track how many blocks are in each dispatch call + std::vector blocks_per_dispatch; + SyncPoint::GetInstance()->SetCallBack( + "IODispatcherImpl::DispatchPrefetch:BlockCount", [&](void* arg) { + auto* indices = static_cast*>(arg); + blocks_per_dispatch.push_back(indices->size()); + }); + SyncPoint::GetInstance()->EnableProcessing(); + + auto job = std::make_shared(); + job->block_handles = block_handles; + job->table = table.get(); + job->job_options.read_options.async_io = false; + + std::shared_ptr read_set; + s = dispatcher->SubmitJob(job, &read_set); + ASSERT_OK(s); + + // First dispatch should have at most 5 blocks (memory limit) + ASSERT_GT(blocks_per_dispatch.size(), 0); + EXPECT_LE(blocks_per_dispatch[0], 5) + << "First dispatch should be limited by memory budget"; + + // Read and release all blocks to trigger remaining dispatches + for (size_t i = 0; i < block_handles.size(); ++i) { + CachableEntry block; + Status read_status = read_set->ReadIndex(i, &block); + ASSERT_OK(read_status); + read_set->ReleaseBlock(i); + } + + SyncPoint::GetInstance()->DisableProcessing(); + SyncPoint::GetInstance()->ClearAllCallBacks(); + + // Verify each dispatch was limited by memory budget + for (size_t i = 0; i < blocks_per_dispatch.size(); ++i) { + EXPECT_LE(blocks_per_dispatch[i], 5) + << "Dispatch " << i << " exceeded memory budget"; + } +} + } // namespace ROCKSDB_NAMESPACE int main(int argc, char** argv) {