Introduce Memory restrictions for IO Dispatcher. (#14300)

Summary:
Introduction of memory limiter for IO Dispatch.

Currently, the user has no way of enacting policy with IO dispatcher. One important policy is the ability to restrict the amount of memory a multiscan or set of multiscans is allowed to pin. This PR introduces the max_prefetch_memory_bytes in the IODispatcherOptions, allowing for users to specify bounds on block cache memory usage.

There seems to be a minor performance increase however, I have found the scans to be a bit noisy. Each benchmark is run with a stride size of 30000 keys. This was done to ensure we maintain parity with trunk.
```
Configuration: 10 concurrent scans, 1024B values, 5242880 byte SST files
Scan sizes: 1024 keys = 1MiB, 2048 keys = 2MiB, 4096 keys = 4MiB per scan

| Keys/Scan | Mode  | Main (ops/sec)   | Main (us/op)     | limiter              (ops/sec) | limiter            (us/op) | Delta ops/sec |
|-----------|-------|------------------|------------------|----------------------|--------------------|---------------|
| 1024      | sync  |   151.6 +/- 8.0   | 6591.14 +/- 343.30 |    170.6 +/- 4.0      | 5855.32 +/- 136.19   | +12.00%       |
| 1024      | async |   156.4 +/- 24.7  | 6589.64 +/- 1345.73 |    173.8 +/- 2.7      | 5744.51 +/- 91.35   | +11.00%       |
| 2048      | sync  |    77.8 +/- 1.6   | 12785.64 +/- 286.49 |     87.6 +/- 3.4      | 11354.01 +/- 441.71   | +12.00%       |
| 2048      | async |    85.6 +/- 4.7   | 11658.11 +/- 618.49 |     91.4 +/- 1.2      | 10873.63 +/- 143.49   | +6.00%        |
| 4096      | sync  |    43.2 +/- 1.5   | 22932.27 +/- 730.66 |     43.8 +/- 0.7      | 22563.90 +/- 320.93   | +1.00%        |
| 4096      | async |    45.4 +/- 0.8   | 21875.64 +/- 357.04 |     46.2 +/- 0.7      | 21416.95 +/- 311.89   | +1.00%        |

```

Pull Request resolved: https://github.com/facebook/rocksdb/pull/14300

Reviewed By: anand1976

Differential Revision: D92316556

Pulled By: krhancoc

fbshipit-source-id: dc0b7958a33b8ef5fa5af82b1c6d960041837fc1
This commit is contained in:
Ryan Hancock
2026-02-06 11:29:30 -08:00
committed by meta-codesync[bot]
parent 3695cb6767
commit 8f9cb1a708
6 changed files with 1469 additions and 83 deletions
+121 -45
View File
@@ -6,8 +6,10 @@
#pragma once
#include <atomic>
#include <functional>
#include <memory>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#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<IOJob> job = std::make_shared<IOJob>();
// 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<IODispatcher> dispatcher(NewIODispatcher());
// std::shared_ptr<ReadSet> 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> 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<size_t>(job->block_handles.front().offset());
// CachableEntry<Block> 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<IOJob> job = std::make_shared<IOJob>();
* // 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<IODispatcher> dispatcher(NewIODispatcher());
* // std::shared_ptr<ReadSet> 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> 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<size_t>(job->block_handles.front().offset());
* // CachableEntry<Block> 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<IODispatcher> dispatcher(NewIODispatcher(opts));
* //
* // // Submit a job that needs more memory than available
* // // Partial prefetch will dispatch what fits immediately
* // std::shared_ptr<ReadSet> 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> 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<size_t, std::shared_ptr<AsyncIOState>> async_io_map_;
// For memory release notifications back to dispatcher (weak ref to avoid
// cycles)
std::weak_ptr<IODispatcherImplData> dispatcher_data_;
// Size of each block (parallel to pinned_blocks_) for memory accounting
std::vector<size_t> block_sizes_;
// Statistics counters
std::atomic<uint64_t> num_sync_reads_ = 0;
std::atomic<uint64_t> 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<std::atomic<bool>[]> pending_prefetch_flags_;
size_t pending_prefetch_flags_size_ = 0;
// Reference to pending request (for removal notification)
std::shared_ptr<PendingPrefetchRequest> pending_request_;
};
/*
@@ -218,8 +290,12 @@ class IODispatcher {
std::shared_ptr<ReadSet>* 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 {
+8
View File
@@ -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
};
+4
View File
@@ -292,6 +292,10 @@ const std::vector<std::pair<Tickers, std::string>> 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<std::pair<Histograms, std::string>> HistogramsNameMap = {
+437 -36
View File
@@ -14,12 +14,15 @@
#include "util/io_dispatcher_imp.h"
#include <deque>
#include <memory>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#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<Block>* 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<size_t> 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<ReadSet> read_set;
std::shared_ptr<IOJob> job;
// Pre-coalesced groups ready for dispatch (ordered by first block index)
std::deque<CoalescedPrefetchGroup> coalesced_groups;
// Individual block indices still pending (for RemoveFromPending lookup)
std::unordered_set<size_t> block_indices_to_prefetch;
std::atomic<size_t> 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<Impl> {
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<IOJob>& job,
std::shared_ptr<ReadSet>* 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<size_t> memory_used_{0}; // Atomic for lock-free accounting
std::atomic<bool> has_pending_requests_{false}; // Fast-path check
port::Mutex memory_mutex_; // Only for pending_prefetch_queue_ access
std::deque<std::shared_ptr<PendingPrefetchRequest>> pending_prefetch_queue_;
Statistics* statistics_ = nullptr;
private:
void PrepareIORequests(
const std::shared_ptr<IOJob>& job,
@@ -335,12 +433,214 @@ struct IODispatcherImpl::Impl {
const std::shared_ptr<ReadSet>& read_set,
std::vector<FSReadRequest>& read_reqs,
const std::vector<std::vector<size_t>>& 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<ReadSet>& read_set,
const std::shared_ptr<IOJob>& job,
const std::vector<size_t>& 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<CoalescedPrefetchGroup> PreCoalesceBlocks(
const std::shared_ptr<IOJob>& job, const std::shared_ptr<ReadSet>& rs,
const std::vector<size_t>& 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<PendingPrefetchRequest> 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<size_t> 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<size_t> 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<ReadSet>& read_set, const std::shared_ptr<IOJob>& job,
const std::vector<size_t>& block_indices) {
// Sync point for testing partial prefetch - passes number of blocks being
// dispatched
TEST_SYNC_POINT_CALLBACK("IODispatcherImpl::DispatchPrefetch:BlockCount",
const_cast<std::vector<size_t>*>(&block_indices));
// Prepare and execute IO for the given blocks
std::vector<FSReadRequest> read_reqs;
std::vector<std::vector<size_t>> 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<size_t> 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<FSReadRequest> sync_read_reqs;
std::vector<std::vector<size_t>> 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<IOJob>& job,
std::shared_ptr<ReadSet>* read_set) {
if (!read_set) {
@@ -353,6 +653,7 @@ Status IODispatcherImpl::Impl::SubmitJob(const std::shared_ptr<IOJob>& 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<IOJob>& job,
rs->num_cache_hits_ =
job->block_handles.size() - block_indices_to_read.size();
// Prepare read requests - coalesce adjacent blocks
std::vector<FSReadRequest> read_reqs;
std::vector<std::vector<size_t>> coalesced_block_indices;
PrepareIORequests(job, block_indices_to_read, job->block_handles, &read_reqs,
&coalesced_block_indices);
// 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<size_t> fallback_indices = ExecuteAsyncIO(
job, rs, read_reqs, coalesced_block_indices, &async_status);
if (!async_status.ok()) {
return async_status;
// 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;
}
// Fall back to sync IO for blocks where async is not supported
if (!fallback_indices.empty()) {
std::vector<FSReadRequest> sync_read_reqs;
std::vector<std::vector<size_t>> sync_coalesced_indices;
PrepareIORequests(job, fallback_indices, job->block_handles,
&sync_read_reqs, &sync_coalesced_indices);
// Store dispatcher reference for release callbacks
rs->dispatcher_data_ = shared_from_this();
Status s = ExecuteSyncIO(job, rs, sync_read_reqs, sync_coalesced_indices);
if (!s.ok()) {
return s;
}
rs->num_sync_reads_ += fallback_indices.size();
// 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_);
std::vector<size_t> blocks_to_dispatch;
std::deque<CoalescedPrefetchGroup> 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);
}
} else {
Status s = ExecuteSyncIO(job, rs, read_reqs, coalesced_block_indices);
if (!s.ok()) {
return s;
// Queue this group for later
groups_to_queue.push_back(std::move(group));
}
}
// 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<PendingPrefetchRequest>();
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<std::atomic<bool>[]>(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<CoalescedPrefetchGroup> IODispatcherImpl::Impl::PreCoalesceBlocks(
const std::shared_ptr<IOJob>& job, const std::shared_ptr<ReadSet>& rs,
const std::vector<size_t>& block_indices, size_t max_group_bytes) {
std::vector<CoalescedPrefetchGroup> 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<size_t> 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<size_t> IODispatcherImpl::Impl::ExecuteAsyncIO(
const std::shared_ptr<IOJob>& job, const std::shared_ptr<ReadSet>& read_set,
std::vector<FSReadRequest>& read_reqs,
@@ -648,7 +1041,11 @@ Status IODispatcherImpl::Impl::ExecuteSyncIO(
return Status::OK();
}
IODispatcherImpl::IODispatcherImpl() : impl_(new Impl()) {}
IODispatcherImpl::IODispatcherImpl()
: impl_(std::make_shared<Impl>(IODispatcherOptions())) {}
IODispatcherImpl::IODispatcherImpl(const IODispatcherOptions& options)
: impl_(std::make_shared<Impl>(options)) {}
IODispatcherImpl::~IODispatcherImpl() = default;
@@ -659,4 +1056,8 @@ Status IODispatcherImpl::SubmitJob(const std::shared_ptr<IOJob>& job,
IODispatcher* NewIODispatcher() { return new IODispatcherImpl(); }
IODispatcher* NewIODispatcher(const IODispatcherOptions& options) {
return new IODispatcherImpl(options);
}
} // namespace ROCKSDB_NAMESPACE
+3 -2
View File
@@ -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<IOJob>& job,
@@ -29,7 +30,7 @@ class IODispatcherImpl : public IODispatcher {
private:
struct Impl;
std::unique_ptr<Impl> impl_;
std::shared_ptr<Impl> impl_;
};
} // namespace ROCKSDB_NAMESPACE
+896
View File
@@ -12,6 +12,7 @@
#include <memory>
#include <mutex>
#include <thread>
#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<IODispatcher> dispatcher(NewIODispatcher(opts));
std::unique_ptr<BlockBasedTable> table;
std::vector<BlockHandle> 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<IOJob>();
job->block_handles = block_handles;
job->table = table.get();
job->job_options.read_options.async_io = false;
std::shared_ptr<ReadSet> 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> 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<IODispatcher> dispatcher(NewIODispatcher(opts));
std::unique_ptr<BlockBasedTable> table;
std::vector<BlockHandle> 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<IOJob>();
job1->block_handles = block_handles;
job1->table = table.get();
job1->job_options.read_options.async_io = false;
std::shared_ptr<ReadSet> 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<BlockBasedTable> table2;
std::vector<BlockHandle> block_handles2;
s = CreateAndOpenSST(30, &table2, &block_handles2);
ASSERT_OK(s);
auto job2 = std::make_shared<IOJob>();
job2->block_handles = block_handles2;
job2->table = table2.get();
job2->job_options.read_options.async_io = false;
std::shared_ptr<ReadSet> 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> 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<IODispatcher> dispatcher(NewIODispatcher(opts));
std::unique_ptr<BlockBasedTable> table;
std::vector<BlockHandle> 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<IOJob>();
job1->block_handles = block_handles;
job1->table = table.get();
job1->job_options.read_options.async_io = false;
std::shared_ptr<ReadSet> 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> 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<BlockBasedTable> table2;
std::vector<BlockHandle> block_handles2;
s = CreateAndOpenSST(20, &table2, &block_handles2);
ASSERT_OK(s);
auto job2 = std::make_shared<IOJob>();
job2->block_handles = block_handles2;
job2->table = table2.get();
job2->job_options.read_options.async_io = false;
std::shared_ptr<ReadSet> 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> 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<IODispatcher> dispatcher(NewIODispatcher(opts));
std::vector<std::shared_ptr<ReadSet>> read_sets;
std::vector<std::vector<BlockHandle>> all_block_handles;
// Create and submit multiple jobs
for (int i = 0; i < 3; i++) {
std::unique_ptr<BlockBasedTable> table;
std::vector<BlockHandle> block_handles;
Status s = CreateAndOpenSST(20 + i * 5, &table, &block_handles);
ASSERT_OK(s);
auto job = std::make_shared<IOJob>();
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<ReadSet> 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> 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<BlockBasedTable> table_new;
std::vector<BlockHandle> block_handles_new;
Status s = CreateAndOpenSST(25, &table_new, &block_handles_new);
ASSERT_OK(s);
auto job_new = std::make_shared<IOJob>();
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<ReadSet> 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> 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<IODispatcher> dispatcher(NewIODispatcher(opts));
std::unique_ptr<BlockBasedTable> table;
std::vector<BlockHandle> block_handles;
Status s = CreateAndOpenSST(50, &table, &block_handles);
ASSERT_OK(s);
auto job = std::make_shared<IOJob>();
job->block_handles = block_handles;
job->table = table.get();
job->job_options.read_options.async_io = false;
std::shared_ptr<ReadSet> 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> 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<IODispatcher> dispatcher(NewIODispatcher(opts));
// Create table outside the scope so it outlives the ReadSet
std::unique_ptr<BlockBasedTable> table;
std::vector<BlockHandle> 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<BlockBasedTable> table2;
std::vector<BlockHandle> block_handles2;
s = CreateAndOpenSST(30, &table2, &block_handles2);
ASSERT_OK(s);
std::shared_ptr<ReadSet> read_set2;
{
auto job = std::make_shared<IOJob>();
job->block_handles = block_handles;
job->table = table.get();
job->job_options.read_options.async_io = false;
std::shared_ptr<ReadSet> 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<IOJob>();
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> 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<IODispatcher> dispatcher(NewIODispatcher(opts));
std::unique_ptr<BlockBasedTable> table;
std::vector<BlockHandle> 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<std::vector<size_t>*>(arg);
blocks_dispatched_on_submit += indices->size();
});
SyncPoint::GetInstance()->EnableProcessing();
auto job = std::make_shared<IOJob>();
job->block_handles = block_handles;
job->table = table.get();
job->job_options.read_options.async_io = true; // Use async IO
std::shared_ptr<ReadSet> 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> 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<IODispatcher> dispatcher(NewIODispatcher(opts));
std::unique_ptr<BlockBasedTable> table;
std::vector<BlockHandle> 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<IOJob>();
job->block_handles = block_handles;
job->table = table.get();
job->job_options.read_options.async_io = true;
std::shared_ptr<ReadSet> 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> 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<BlockBasedTable> table;
std::vector<BlockHandle> 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<IODispatcher> 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<std::vector<size_t>*>(arg);
blocks_dispatched += indices->size();
});
SyncPoint::GetInstance()->EnableProcessing();
auto job = std::make_shared<IOJob>();
job->block_handles = block_handles;
job->table = table.get();
job->job_options.read_options.async_io = true;
std::shared_ptr<ReadSet> 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> 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<IODispatcher> dispatcher(NewIODispatcher(opts));
std::unique_ptr<BlockBasedTable> table;
std::vector<BlockHandle> block_handles;
Status s = CreateAndOpenSST(20, &table, &block_handles);
ASSERT_OK(s);
ASSERT_GE(block_handles.size(), 10);
auto job = std::make_shared<IOJob>();
job->block_handles = block_handles;
job->table = table.get();
job->job_options.read_options.async_io = true;
std::shared_ptr<ReadSet> 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> 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> 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<BlockBasedTable> table;
std::vector<BlockHandle> 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<IODispatcher> dispatcher(NewIODispatcher(opts));
auto job = std::make_shared<IOJob>();
job->block_handles = block_handles;
job->table = table.get();
job->job_options.read_options.async_io = true;
std::shared_ptr<ReadSet> 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> 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<IODispatcher> dispatcher(NewIODispatcher(opts));
std::unique_ptr<BlockBasedTable> table;
std::vector<BlockHandle> block_handles;
Status s = CreateAndOpenSST(30, &table, &block_handles);
ASSERT_OK(s);
ASSERT_GE(block_handles.size(), 10);
auto job = std::make_shared<IOJob>();
job->block_handles = block_handles;
job->table = table.get();
job->job_options.read_options.async_io = true;
std::shared_ptr<ReadSet> 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> 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<IODispatcher> dispatcher(NewIODispatcher(opts));
std::unique_ptr<BlockBasedTable> table;
std::vector<BlockHandle> block_handles;
Status s = CreateAndOpenSST(20, &table, &block_handles);
ASSERT_OK(s);
ASSERT_GE(block_handles.size(), 10);
auto job = std::make_shared<IOJob>();
job->block_handles = block_handles;
job->table = table.get();
job->job_options.read_options.async_io = false; // Sync IO
std::shared_ptr<ReadSet> 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> 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<BlockBasedTable> table;
std::vector<BlockHandle> 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<IODispatcher> dispatcher(NewIODispatcher(opts));
auto job = std::make_shared<IOJob>();
job->block_handles = block_handles;
job->table = table.get();
job->job_options.read_options.async_io = false; // Sync IO
std::shared_ptr<ReadSet> 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> 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<BlockBasedTable> table;
std::vector<BlockHandle> 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<IODispatcher> 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<IOJob>();
job->block_handles = block_handles;
job->table = table.get();
job->job_options.read_options.async_io = async;
std::shared_ptr<ReadSet> 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> 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<BlockBasedTable> table;
std::vector<BlockHandle> 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<IODispatcher> dispatcher(NewIODispatcher(opts));
// Track dispatch calls
std::vector<size_t> dispatch_counts;
SyncPoint::GetInstance()->SetCallBack(
"IODispatcherImpl::DispatchPrefetch:BlockCount", [&](void* arg) {
auto* indices = static_cast<std::vector<size_t>*>(arg);
dispatch_counts.push_back(indices->size());
});
SyncPoint::GetInstance()->EnableProcessing();
auto job = std::make_shared<IOJob>();
job->block_handles = block_handles;
job->table = table.get();
job->job_options.read_options.async_io = false; // Sync IO
std::shared_ptr<ReadSet> 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> 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> 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<BlockBasedTable> table;
std::vector<BlockHandle> 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<IODispatcher> dispatcher(NewIODispatcher(opts));
// Track how many blocks are in each dispatch call
std::vector<size_t> blocks_per_dispatch;
SyncPoint::GetInstance()->SetCallBack(
"IODispatcherImpl::DispatchPrefetch:BlockCount", [&](void* arg) {
auto* indices = static_cast<std::vector<size_t>*>(arg);
blocks_per_dispatch.push_back(indices->size());
});
SyncPoint::GetInstance()->EnableProcessing();
auto job = std::make_shared<IOJob>();
job->block_handles = block_handles;
job->table = table.get();
job->job_options.read_options.async_io = false;
std::shared_ptr<ReadSet> 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> 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) {