mirror of
https://github.com/facebook/rocksdb.git
synced 2026-07-07 14:47:40 +08:00
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:
committed by
meta-codesync[bot]
parent
3695cb6767
commit
8f9cb1a708
+121
-45
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
};
|
||||
|
||||
|
||||
@@ -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
@@ -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);
|
||||
// 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<size_t> 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<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);
|
||||
// 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<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);
|
||||
}
|
||||
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<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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user