database/rocksdb
1
0
Fork 0
mirror of https://github.com/facebook/rocksdb.git synced 2026-10-23 08:09:02 +03:00
Code Activity

Add opt-in stats for async RocksDB reads (#15017)

Browse Source 
Summary:
Pull Request resolved: https://github.com/facebook/rocksdb/pull/15017

## Summary

RocksDB async `GetAsync` and `MultiGetAsync` callbacks now default to a completion-only path. Callbacks can override `EnableStats()` when they need request-specific `PerfContext` or `IOStatsContext`; otherwise RocksDB avoids creating a coroutine stats context and returns null stats pointers. Callbacks must not initiate another async read from `OnComplete()`.

RocksDB stats normally use thread-local storage, which does not compose with interleaved coroutine reads. A stats-enabled request captures only the caller's configuration: perf level, per-level perf enablement, and whether io stats are disabled. It then starts with empty request contexts. `onSet()` loads that request's contexts into TLS, while `onUnset()` saves the accumulated contexts back into request data. On resume, the same contexts are loaded again. This isolates request counters without preserving or restoring pre-existing executor TLS state. The synchronous fallback similarly resets TLS before executing a stats-enabled read.

`PerfContext` and `IOStatsContext` gain `Merge()` support for caller-side aggregation. `PerfContext` move construction and assignment now transfer ownership of per-level state rather than deep-copying it, avoiding allocations during coroutine context handoffs. A fbcode-wide audit found no caller that reads a moved-from `PerfContext`.

`db_bench` opts in only for coroutine jobs that print per-job perf context and merges every completed request's callback stats into the job aggregate.

## Example: Two Async Gets on One Thread

Both callbacks enable stats. Each request starts with empty counters and carries its contexts across suspensions.

```text
Shared executor thread
|
|  onSet(A): load A's empty contexts into TLS
|  A records 3 reads
|  onUnset(A): save TLS contexts back into A
|
|  onSet(B): load B's empty contexts into TLS
|  B records 4 reads
|  onUnset(B): save TLS contexts back into B
|
|  onSet(A): load A's contexts containing 3 reads
|  A records 2 more reads
|  callback A receives block_read_count = 5
|
|  onSet(B): load B's contexts containing 4 reads
|  callback B receives block_read_count = 4
```

Isolation comes from replacing TLS with the active request's saved contexts on every `onSet()`.

### Timers

Existing wall-clock timers remain in coroutine frames, so timers spanning `co_await` include suspended time and measure end-to-end latency.

Request-level `get_cpu_nanos` is suspension-aware: it starts in `onSet()` and stops in `onUnset()` or `Finalize()`, charging only time while the request is active.

Nested CPU timers remain disabled because normal RAII CPU timers would also charge executor work performed while their coroutine is suspended. `MultiGetAsync` fanout additionally requires per-child timer state to avoid double counting.

## Benchmarks

Benchmarked to confirm no performance regressions when perf is disabled, and also to see the overhead of enabling perf.

```bash
db_bench --benchmarks=<readrandomcoroutine|multireadrandomcoroutine> --db=<db_path> --coro_jobs_per_thread=10 --use_existing_db=true --num=2000000 --value_size=1024 --compression_type=none --threads=8 --batch_size=16 --multiread_batched=true --use_direct_reads=true --cache_size=1428076626 --reads=320000 --perf_level=<1|3> --seed=<seed>
```

| Workload | Before change | Perf level 1 | Before -> 1 | Perf level 3 | 1 -> 3 |
| -------- | ------------- | ------------ | ----------- | ------------ | ------ |
| Async Get | 475.6K ops/s | 476.7K ops/s | +0.2% | 470.5K ops/s | -1.3% |
| Async MultiGet | 567.3K ops/s | 567.4K ops/s | +0.0% | 552.2K ops/s | -2.7% |

Reviewed By: xingbowang

Differential Revision: D112634874

fbshipit-source-id: 94e240aeb1496fa12dac995b5aa5a9c3543cc5fc
This commit is contained in:
Josh Kang
2026-07-28 09:32:26 -07:00
committed by meta-codesync[bot]
parent f5c0c3e274
commit 7247edbf0a
22 changed files with 960 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
BUCK
View File

@@ -268,6 +268,7 @@ cpp_library_wrapper(name="rocksdb_lib", srcs=[
"util/compression.cc",
"util/compression_context_cache.cc",
"util/concurrent_task_limiter_impl.cc",
"util/coro_stats_util.cc",
"util/crc32c.cc",
"util/crc32c_arm64.cc",
"util/data_structure.cc",
@@ -392,6 +393,7 @@ cpp_library_wrapper(name="rocksdb_lib", srcs=[
"//folly/io/async:async_base",
"//folly/io/async:event_base_manager",
"//folly/io/async:io_uring_backend",
"//folly/io/async:request_context",
"//folly/synchronization:distributed_mutex",
"//folly:executor",
], headers=glob(["**/*.h"]), link_whole=False, extra_test_libs=False)

1
CMakeLists.txt
View File

@@ -1002,6 +1002,7 @@ set(SOURCES
util/simple_mixed_compressor.cc
util/compression_context_cache.cc
util/concurrent_task_limiter_impl.cc
util/coro_stats_util.cc
util/crc32c.cc
util/data_structure.cc
util/dynamic_bloom.cc

1
buckifier/buckify_rocksdb.py
View File

@@ -159,6 +159,7 @@ def generate_buck(repo_path, deps_map):
"//folly/io/async:async_base",
"//folly/io/async:event_base_manager",
"//folly/io/async:io_uring_backend",
"//folly/io/async:request_context",
"//folly/synchronization:distributed_mutex",
],
headers=LiteralValue("glob([\"**/*.h\"])")

178
db/db_impl/db_impl.cc
View File

@@ -17,6 +17,7 @@
#if USE_COROUTINES
#include "folly/Executor.h"
#include "folly/coro/Invoke.h"
#include "folly/executors/IOExecutor.h"
#include "folly/io/async/EventBase.h"
#endif // USE_COROUTINES
@@ -123,6 +124,7 @@
#include "util/cast_util.h"
#include "util/coding.h"
#include "util/compression.h"
#include "util/coro_stats_util.h"
#include "util/crc32c.h"
#include "util/defer.h"
#include "util/distributed_mutex.h"
@@ -7513,40 +7515,58 @@ void DBImpl::TriggerPeriodicCompaction() {
}
}
namespace {
void ResetThreadLocalStatsForAsyncCallback() {
#ifndef NPERF_CONTEXT
get_perf_context()->Reset();
#endif
#ifndef NIOSTATS_CONTEXT
get_iostats_context()->Reset();
#endif
}
const PerfContext* CurrentPerfContextForAsyncCallback(bool stats_enabled) {
if (!stats_enabled) {
return nullptr;
}
#ifdef NPERF_CONTEXT
return nullptr;
#else
return get_perf_context();
#endif
}
const IOStatsContext* CurrentIOStatsContextForAsyncCallback(
bool stats_enabled) {
if (!stats_enabled) {
return nullptr;
}
#ifdef NIOSTATS_CONTEXT
return nullptr;
#else
return get_iostats_context();
#endif
}
#if USE_COROUTINES
std::optional<CoroutineStatsConfig> CaptureCoroutineStatsConfigForCallback(
bool stats_enabled) {
if (!stats_enabled) {
return std::nullopt;
}
return CaptureCoroutineStatsConfig();
}
#endif // USE_COROUTINES
} // namespace
void DBImpl::GetAsync(const ReadOptions& options,
ColumnFamilyHandle* column_family, const Slice& key,
PinnableSlice* value, std::string* timestamp,
Status& status, AsyncCallback& callback) {
#if USE_COROUTINES
auto* read_executor =
immutable_db_options_.env->GetFileSystem()->GetReadExecutor();
if (read_executor != nullptr) {
auto* read_event_base = read_executor->getEventBase();
assert(read_event_base != nullptr);
auto task = [](DBImpl* db, ReadOptions task_options,
ColumnFamilyHandle* task_column_family, Slice task_key,
PinnableSlice* task_value, std::string* task_timestamp,
Status& task_status,
AsyncCallback& task_callback) -> folly::coro::Task<void> {
task_status = co_await folly::coro::co_nothrow(
db->GetCoroutine(task_options, task_column_family, task_key,
task_value, task_timestamp));
task_callback.OnComplete();
}(this, options, column_family, key, value, timestamp, status, callback);
folly::coro::co_withExecutor(
folly::Executor::getKeepAliveToken(read_event_base), std::move(task))
.start();
return;
}
#endif // USE_COROUTINES
DB::GetAsync(options, column_family, key, value, timestamp, status, callback);
}
void DBImpl::MultiGetAsync(const ReadOptions& options, const size_t num_keys,
ColumnFamilyHandle** column_families,
const Slice* keys, PinnableSlice* values,
std::string* timestamps, Status* statuses,
const bool sorted_input, AsyncCallback& callback) {
const bool stats_enabled = callback.EnableStats();
#if USE_COROUTINES
auto* read_executor =
immutable_db_options_.env->GetFileSystem()->GetReadExecutor();
@@ -7554,25 +7574,97 @@ void DBImpl::MultiGetAsync(const ReadOptions& options, const size_t num_keys,
auto* read_event_base = read_executor->getEventBase();
assert(read_event_base != nullptr);
auto task =
[](DBImpl* db, ReadOptions task_options, size_t task_num_keys,
ColumnFamilyHandle** task_column_families, const Slice* task_keys,
PinnableSlice* task_values, std::string* task_timestamps,
Status* task_statuses, bool task_sorted_input,
AsyncCallback& task_callback) -> folly::coro::Task<void> {
co_await folly::coro::co_nothrow(db->MultiGetCoroutine(
task_options, task_num_keys, task_column_families, task_keys,
task_values, task_timestamps, task_statuses, task_sorted_input));
task_callback.OnComplete();
}(this, options, num_keys, column_families, keys, values, timestamps,
statuses, sorted_input, callback);
[](std::optional<CoroutineStatsConfig> task_stats_config, DBImpl* db,
ReadOptions task_options, ColumnFamilyHandle* task_column_family,
Slice task_key, PinnableSlice* task_value,
std::string* task_timestamp, Status& task_status,
AsyncCallback& task_callback) mutable -> folly::coro::Task<void> {
std::optional<CoroutineStatsContextScope> coroutine_stats_scope;
if (task_stats_config.has_value()) {
coroutine_stats_scope.emplace(std::move(*task_stats_config),
db->immutable_db_options_.env);
}
task_status = co_await folly::coro::co_nothrow(
db->GetCoroutine(task_options, task_column_family, task_key,
task_value, task_timestamp));
if (coroutine_stats_scope.has_value()) {
coroutine_stats_scope->Finalize();
}
task_callback.OnComplete(
CurrentPerfContextForAsyncCallback(coroutine_stats_scope.has_value()),
CurrentIOStatsContextForAsyncCallback(
coroutine_stats_scope.has_value()));
}(CaptureCoroutineStatsConfigForCallback(stats_enabled), this, options,
column_family, key, value,
timestamp, status, callback);
folly::coro::co_withExecutor(
folly::Executor::getKeepAliveToken(read_event_base), std::move(task))
.start();
return;
}
#endif // USE_COROUTINES
DB::MultiGetAsync(options, num_keys, column_families, keys, values,
timestamps, statuses, sorted_input, callback);
if (stats_enabled) {
// Mirror behavior of coroutine version where each coroutine gets a clean
// context
ResetThreadLocalStatsForAsyncCallback();
}
status = Get(options, column_family, key, value, timestamp);
callback.OnComplete(CurrentPerfContextForAsyncCallback(stats_enabled),
CurrentIOStatsContextForAsyncCallback(stats_enabled));
}
void DBImpl::MultiGetAsync(const ReadOptions& options, const size_t num_keys,
ColumnFamilyHandle** column_families,
const Slice* keys, PinnableSlice* values,
std::string* timestamps, Status* statuses,
const bool sorted_input, AsyncCallback& callback) {
const bool stats_enabled = callback.EnableStats();
#if USE_COROUTINES
auto* read_executor =
immutable_db_options_.env->GetFileSystem()->GetReadExecutor();
if (read_executor != nullptr) {
auto* read_event_base = read_executor->getEventBase();
assert(read_event_base != nullptr);
auto task = [](std::optional<CoroutineStatsConfig> task_stats_config,
DBImpl* db, ReadOptions task_options, size_t task_num_keys,
ColumnFamilyHandle** task_column_families,
const Slice* task_keys, PinnableSlice* task_values,
std::string* task_timestamps, Status* task_statuses,
bool task_sorted_input, AsyncCallback& task_callback) mutable
-> folly::coro::Task<void> {
std::optional<CoroutineStatsContextScope> coroutine_stats_scope;
if (task_stats_config.has_value()) {
coroutine_stats_scope.emplace(std::move(*task_stats_config),
db->immutable_db_options_.env);
}
co_await folly::coro::co_nothrow(db->MultiGetCoroutine(
task_options, task_num_keys, task_column_families, task_keys,
task_values, task_timestamps, task_statuses, task_sorted_input));
if (coroutine_stats_scope.has_value()) {
coroutine_stats_scope->Finalize();
}
task_callback.OnComplete(
CurrentPerfContextForAsyncCallback(coroutine_stats_scope.has_value()),
CurrentIOStatsContextForAsyncCallback(
coroutine_stats_scope.has_value()));
}(CaptureCoroutineStatsConfigForCallback(stats_enabled), this, options,
num_keys, column_families, keys, values, timestamps, statuses,
sorted_input, callback);
folly::coro::co_withExecutor(
folly::Executor::getKeepAliveToken(read_event_base), std::move(task))
.start();
return;
}
#endif // USE_COROUTINES
if (stats_enabled) {
// Mirror behavior of coroutine version where each coroutine gets a clean
// context
ResetThreadLocalStatsForAsyncCallback();
}
MultiGet(options, num_keys, column_families, keys, values, timestamps,
statuses, sorted_input);
callback.OnComplete(CurrentPerfContextForAsyncCallback(stats_enabled),
CurrentIOStatsContextForAsyncCallback(stats_enabled));
}
void DBImpl::TrackOrUntrackFiles(

8
db/db_impl/db_impl_readonly_sync_and_async.h
View File

@@ -43,9 +43,9 @@ DEFINE_SYNC_AND_ASYNC(Status, DBImplReadOnly::GetImpl)
#if defined(WITHOUT_COROUTINES)
PERF_CPU_TIMER_GUARD(get_cpu_nanos, immutable_db_options_.clock);
PERF_TIMER_GUARD(get_snapshot_time);
#endif // defined(WITHOUT_COROUTINES)
#endif
StopWatch sw(immutable_db_options_.clock, stats_, DB_GET);
PERF_TIMER_GUARD(get_snapshot_time);
const Comparator* ucmp = get_impl_options.column_family->GetComparator();
assert(ucmp);
@@ -92,9 +92,7 @@ DEFINE_SYNC_AND_ASYNC(Status, DBImplReadOnly::GetImpl)
SequenceNumber max_covering_tombstone_seq = 0;
LookupKey lkey(key, snapshot, read_options.timestamp);
#if defined(WITHOUT_COROUTINES)
PERF_TIMER_STOP(get_snapshot_time);
#endif // defined(WITHOUT_COROUTINES)
std::optional<VersionBlobFetcher> memtable_blob_fetcher;
if (cfd->ioptions().enable_blob_direct_write ||
cfd->GetLatestMutableCFOptions().enable_blob_files) {
@@ -122,9 +120,7 @@ DEFINE_SYNC_AND_ASYNC(Status, DBImplReadOnly::GetImpl)
&is_blob_index, get_impl_options.value_found);
RecordTick(stats_, MEMTABLE_HIT);
} else {
#if defined(WITHOUT_COROUTINES)
PERF_TIMER_GUARD(get_from_output_files_time);
#endif // defined(WITHOUT_COROUTINES)
PinnedIteratorsManager pinned_iters_mgr;
CO_AWAIT(super_version->current->Get, read_options, lkey,
get_impl_options.value, get_impl_options.columns, ts, &s,

10
db/db_impl/db_impl_secondary_sync_and_async.h
View File

@@ -42,9 +42,9 @@ DEFINE_SYNC_AND_ASYNC(Status, DBImplSecondary::GetImpl)
#if defined(WITHOUT_COROUTINES)
PERF_CPU_TIMER_GUARD(get_cpu_nanos, immutable_db_options_.clock);
PERF_TIMER_GUARD(get_snapshot_time);
#endif // defined(WITHOUT_COROUTINES)
#endif
StopWatch sw(immutable_db_options_.clock, stats_, DB_GET);
PERF_TIMER_GUARD(get_snapshot_time);
const Comparator* ucmp = get_impl_options.column_family->GetComparator();
assert(ucmp);
@@ -98,9 +98,7 @@ DEFINE_SYNC_AND_ASYNC(Status, DBImplSecondary::GetImpl)
SequenceNumber max_covering_tombstone_seq = 0;
LookupKey lkey(key, snapshot, read_options.timestamp);
#if defined(WITHOUT_COROUTINES)
PERF_TIMER_STOP(get_snapshot_time);
#endif // defined(WITHOUT_COROUTINES)
bool done = false;
std::optional<VersionBlobFetcher> memtable_blob_fetcher;
if (cfd->ioptions().enable_blob_direct_write ||
@@ -179,9 +177,7 @@ DEFINE_SYNC_AND_ASYNC(Status, DBImplSecondary::GetImpl)
CO_RETURN s;
}
if (!done) {
#if defined(WITHOUT_COROUTINES)
PERF_TIMER_GUARD(get_from_output_files_time);
#endif // defined(WITHOUT_COROUTINES)
PinnedIteratorsManager pinned_iters_mgr;
CO_AWAIT(super_version->current->Get, read_options, lkey,
get_impl_options.value, get_impl_options.columns, ts, &s,
@@ -193,9 +189,7 @@ DEFINE_SYNC_AND_ASYNC(Status, DBImplSecondary::GetImpl)
RecordTick(stats_, MEMTABLE_MISS);
}
{
#if defined(WITHOUT_COROUTINES)
PERF_TIMER_GUARD(get_post_process_time);
#endif // defined(WITHOUT_COROUTINES)
#if defined(WITH_COROUTINES)
CleanupSuperVersion(super_version);
#else

19
db/db_impl/db_impl_sync_and_async.h
View File

@@ -47,9 +47,10 @@ DEFINE_SYNC_AND_ASYNC(Status, DBImpl::Get)
if (read_executor != nullptr) {
auto* read_event_base = read_executor->getEventBase();
assert(read_event_base != nullptr);
co_return co_await folly::coro::co_nothrow(folly::coro::co_withExecutor(
Status s = co_await folly::coro::co_nothrow(folly::coro::co_withExecutor(
folly::Executor::getKeepAliveToken(read_event_base),
GetImplCoroutine(read_options, column_family, key, value, timestamp)));
co_return s;
}
#endif
CO_RETURN GetImpl(read_options, column_family, key, value, timestamp);
@@ -99,11 +100,9 @@ DEFINE_SYNC_AND_ASYNC(Status, DBImpl::GetImpl)
#if defined(WITHOUT_COROUTINES)
PERF_CPU_TIMER_GUARD(get_cpu_nanos, immutable_db_options_.clock);
#endif // defined(WITHOUT_COROUTINES)
#endif
StopWatch sw(immutable_db_options_.clock, stats_, DB_GET);
#if defined(WITHOUT_COROUTINES)
PERF_TIMER_GUARD(get_snapshot_time);
#endif // defined(WITHOUT_COROUTINES)
auto cfh = static_cast_with_check<ColumnFamilyHandleImpl>(
get_impl_options.column_family);
@@ -208,9 +207,7 @@ DEFINE_SYNC_AND_ASYNC(Status, DBImpl::GetImpl)
// s is both in/out. When in, s could either be OK or MergeInProgress.
// merge_operands will contain the sequence of merges in the latter case.
LookupKey lkey(key, snapshot, read_options.timestamp);
#if defined(WITHOUT_COROUTINES)
PERF_TIMER_STOP(get_snapshot_time);
#endif // defined(WITHOUT_COROUTINES)
bool skip_memtable = (read_options.read_tier == kPersistedTier &&
has_unpersisted_data_.load(std::memory_order_relaxed));
@@ -301,9 +298,7 @@ DEFINE_SYNC_AND_ASYNC(Status, DBImpl::GetImpl)
TEST_SYNC_POINT("DBImpl::GetImpl:PostMemTableGet:1");
PinnedIteratorsManager pinned_iters_mgr;
if (!done) {
#if defined(WITHOUT_COROUTINES)
PERF_TIMER_GUARD(get_from_output_files_time);
#endif // defined(WITHOUT_COROUTINES)
CO_AWAIT(
sv->current->Get, read_options, lkey, get_impl_options.value,
get_impl_options.columns, timestamp, &s, &merge_context,
@@ -326,9 +321,7 @@ DEFINE_SYNC_AND_ASYNC(Status, DBImpl::GetImpl)
}
{
#if defined(WITHOUT_COROUTINES)
PERF_TIMER_GUARD(get_post_process_time);
#endif // defined(WITHOUT_COROUTINES)
RecordTick(stats_, NUMBER_KEYS_READ);
size_t size = 0;
@@ -521,9 +514,7 @@ DEFINE_SYNC_AND_ASYNC(Status, DBImpl::MultiGetImpl)
}
}
if (lookup_current) {
#if defined(WITHOUT_COROUTINES)
PERF_TIMER_GUARD(get_from_output_files_time);
#endif // defined(WITHOUT_COROUTINES)
CO_AWAIT(super_version->current->MultiGet, read_options, &range,
callback);
}
@@ -542,9 +533,7 @@ DEFINE_SYNC_AND_ASYNC(Status, DBImpl::MultiGetImpl)
}
// Post processing (decrement reference counts and record statistics)
#if defined(WITHOUT_COROUTINES)
PERF_TIMER_GUARD(get_post_process_time);
#endif // defined(WITHOUT_COROUTINES)
size_t num_found = 0;
uint64_t bytes_read = 0;
// value_size_soft_limit was enforced above (and inside per-file MultiGet) on
@@ -627,9 +616,7 @@ DEFINE_SYNC_AND_ASYNC(Status, DBImpl::MultiGetImpl)
RecordTick(stats_, NUMBER_MULTIGET_BYTES_READ, bytes_read);
RecordInHistogram(stats_, BYTES_PER_MULTIGET, bytes_read);
PERF_COUNTER_ADD(multiget_read_bytes, bytes_read);
#if defined(WITHOUT_COROUTINES)
PERF_TIMER_STOP(get_post_process_time);
#endif // defined(WITHOUT_COROUTINES)
CO_RETURN s;
}

280
db/perf_context_test.cc
View File

@@ -7,23 +7,41 @@
#include <algorithm>
#include <iostream>
#include <memory>
#include <thread>
#include <vector>
#include "env/composite_env_wrapper.h"
#include "monitoring/histogram.h"
#include "monitoring/instrumented_mutex.h"
#include "monitoring/perf_context_imp.h"
#include "monitoring/thread_status_util.h"
#include "port/port.h"
#include "rocksdb/db.h"
#include "rocksdb/env.h"
#include "rocksdb/file_system.h"
#include "rocksdb/iostats_context.h"
#include "rocksdb/memtablerep.h"
#include "rocksdb/slice_transform.h"
#include "rocksdb/system_clock.h"
#include "test_util/mock_time_env.h"
#include "test_util/testharness.h"
#include "util/stop_watch.h"
#include "util/string_util.h"
#include "utilities/merge_operators.h"
#if defined(USE_COROUTINES)
#include "folly/Executor.h"
#include "folly/coro/BlockingWait.h"
#include "folly/coro/Collect.h"
#include "folly/coro/CurrentExecutor.h"
#include "folly/coro/Task.h"
#include "folly/executors/IOThreadPoolExecutor.h"
#include "folly/io/async/EventBase.h"
#include "folly/io/async/Request.h"
#include "util/coro_stats_util.h"
#endif
bool FLAGS_random_key = false;
bool FLAGS_use_set_based_memetable = false;
int FLAGS_total_keys = 100;
@@ -64,7 +82,267 @@ std::unique_ptr<DB> OpenDb(bool read_only = false) {
return db;
}
class PerfContextTest : public testing::Test {};
class PerfContextTest : public testing::Test {
protected:
void SetUp() override {
SetPerfLevel(PerfLevel::kEnableCount);
get_perf_context()->ClearPerLevelPerfContext();
get_perf_context()->Reset();
get_iostats_context()->Reset();
get_iostats_context()->disable_iostats = false;
}
};
#if defined(USE_COROUTINES)
TEST_F(PerfContextTest, CoroutineStatsContextScopeCollectsStats) {
SetPerfLevel(PerfLevel::kEnableTimeAndCPUTimeExceptForMutex);
get_perf_context()->EnablePerLevelPerfContext();
folly::IOThreadPoolExecutor executor(1);
folly::EventBase* event_base = executor.getEventBase();
ASSERT_NE(nullptr, event_base);
auto clock = std::make_shared<MockSystemClock>(SystemClock::Default());
clock->SetCurrentTime(1);
CoroutineStatsConfig stats_config = CaptureCoroutineStatsConfig();
folly::coro::blockingWait(folly::coro::co_withExecutor(
folly::Executor::getKeepAliveToken(event_base),
[event_base, clock = clock.get(),
stats_config =
std::move(stats_config)]() mutable -> folly::coro::Task<void> {
EXPECT_TRUE(event_base->isInEventBaseThread());
CoroutineStatsContextScope scope(std::move(stats_config),
Env::Default());
EXPECT_EQ(PerfLevel::kEnableTimeAndCPUTimeExceptForMutex,
GetPerfLevel());
EXPECT_EQ(0, get_perf_context()->block_read_count);
EXPECT_TRUE(get_perf_context()->per_level_perf_context_enabled);
get_perf_context()->block_read_count += 3;
(*get_perf_context()->level_to_perf_context)[2].block_cache_hit_count +=
5;
EXPECT_EQ(0, get_iostats_context()->bytes_read);
get_iostats_context()->bytes_read += 7;
get_iostats_context()->cpu_read_nanos += 11;
get_iostats_context()
->file_io_stats_by_temperature.hot_file_read_count += 1;
{
PERF_TIMER_GUARD_WITH_CLOCK(block_read_time, clock);
clock->SleepForMicroseconds(7);
}
co_await folly::coro::co_reschedule_on_current_executor;
EXPECT_TRUE(event_base->isInEventBaseThread());
EXPECT_EQ(3, get_perf_context()->block_read_count);
EXPECT_TRUE(get_perf_context()->per_level_perf_context_enabled);
get_perf_context()->block_read_count += 5;
(*get_perf_context()->level_to_perf_context)[2].block_cache_hit_count +=
2;
EXPECT_EQ(7, get_iostats_context()->bytes_read);
get_iostats_context()->bytes_read += 11;
get_iostats_context()->cpu_read_nanos += 13;
get_iostats_context()
->file_io_stats_by_temperature.hot_file_read_count += 2;
scope.Finalize();
auto verify_stats = [] {
const PerfContext* captured_perf_context = get_perf_context();
ASSERT_NE(nullptr, captured_perf_context);
EXPECT_EQ(8, captured_perf_context->block_read_count);
EXPECT_EQ(7000, captured_perf_context->block_read_time);
ASSERT_NE(nullptr, captured_perf_context->level_to_perf_context);
EXPECT_EQ(7, (*captured_perf_context->level_to_perf_context)[2]
.block_cache_hit_count);
const IOStatsContext* iostats_context = get_iostats_context();
ASSERT_NE(nullptr, iostats_context);
EXPECT_EQ(18, iostats_context->bytes_read);
EXPECT_EQ(24, iostats_context->cpu_read_nanos);
EXPECT_EQ(3, iostats_context->file_io_stats_by_temperature
.hot_file_read_count);
};
verify_stats();
}()));
}
TEST_F(PerfContextTest, CoroutineStatsContextsRemainIsolatedWhenInterleaved) {
CoroutineStatsConfig first_stats_config = CaptureCoroutineStatsConfig();
SetPerfLevel(PerfLevel::kEnableTimeExceptForMutex);
get_perf_context()->EnablePerLevelPerfContext();
get_iostats_context()->disable_iostats = true;
CoroutineStatsConfig second_stats_config = CaptureCoroutineStatsConfig();
folly::IOThreadPoolExecutor executor(1);
folly::EventBase* event_base = executor.getEventBase();
ASSERT_NE(nullptr, event_base);
folly::coro::blockingWait(folly::coro::co_withExecutor(
folly::Executor::getKeepAliveToken(event_base),
[first_stats_config = std::move(first_stats_config),
second_stats_config = std::move(
second_stats_config)]() mutable -> folly::coro::Task<void> {
get_perf_context()->Reset();
get_perf_context()->ClearPerLevelPerfContext();
get_perf_context()->block_read_count = 100;
get_iostats_context()->Reset();
get_iostats_context()->bytes_read = 200;
auto make_task = [](CoroutineStatsConfig stats_config, uint64_t first,
uint64_t second) -> folly::coro::Task<uint64_t> {
const PerfLevel expected_perf_level = stats_config.perf_level;
const bool expected_per_level_perf_context_enabled =
stats_config.per_level_perf_context_enabled;
const bool expected_iostats_disabled = stats_config.iostats_disabled;
CoroutineStatsContextScope scope(std::move(stats_config),
Env::Default());
EXPECT_EQ(expected_perf_level, GetPerfLevel());
EXPECT_EQ(expected_per_level_perf_context_enabled,
get_perf_context()->per_level_perf_context_enabled);
EXPECT_EQ(expected_iostats_disabled,
get_iostats_context()->disable_iostats);
EXPECT_EQ(0, get_perf_context()->block_read_count);
EXPECT_EQ(0, get_iostats_context()->bytes_read);
get_perf_context()->block_read_count += first;
get_iostats_context()->bytes_read += first;
co_await folly::coro::co_reschedule_on_current_executor;
EXPECT_EQ(expected_perf_level, GetPerfLevel());
EXPECT_EQ(expected_per_level_perf_context_enabled,
get_perf_context()->per_level_perf_context_enabled);
EXPECT_EQ(expected_iostats_disabled,
get_iostats_context()->disable_iostats);
EXPECT_EQ(first, get_perf_context()->block_read_count);
EXPECT_EQ(first, get_iostats_context()->bytes_read);
get_perf_context()->block_read_count += second;
get_iostats_context()->bytes_read += second;
scope.Finalize();
EXPECT_EQ(first + second, get_iostats_context()->bytes_read);
co_return get_perf_context()->block_read_count;
};
std::vector<folly::coro::Task<uint64_t>> tasks;
tasks.push_back(make_task(std::move(first_stats_config), 1, 10));
tasks.push_back(make_task(std::move(second_stats_config), 2, 20));
std::vector<uint64_t> results =
co_await folly::coro::collectAllRange(std::move(tasks));
EXPECT_EQ((std::vector<uint64_t>{11, 22}), results);
get_perf_context()->Reset();
get_iostats_context()->Reset();
co_return;
}()));
}
TEST_F(PerfContextTest, CoroutineStatsContextScopeCollectsGetCpuNanos) {
SetPerfLevel(PerfLevel::kEnableTimeAndCPUTimeExceptForMutex);
auto clock = std::make_shared<MockSystemClock>(SystemClock::Default());
CompositeEnvWrapper env(Env::Default(), clock);
clock->SetCurrentCPUTimeNanos(0);
{
CoroutineStatsContextScope scope(CaptureCoroutineStatsConfig(), &env);
clock->MockSleepForCPUNanos(17);
clock->MockSleepForCPUNanos(5);
// Time while the request context is unset is not charged to the request.
{
folly::RequestContextScopeGuard context_scope;
clock->MockSleepForCPUNanos(13);
}
clock->MockSleepForCPUNanos(7);
scope.Finalize();
EXPECT_EQ(29, get_perf_context()->get_cpu_nanos);
}
}
#endif
TEST_F(PerfContextTest, AsyncCallbackStatsStartEmptyOnSyncFallback) {
class NoReadExecutorFileSystem final : public FileSystemWrapper {
public:
explicit NoReadExecutorFileSystem(const std::shared_ptr<FileSystem>& target)
: FileSystemWrapper(target) {}
const char* Name() const override { return "NoReadExecutorFileSystem"; }
folly::IOExecutor* GetReadExecutor() override { return nullptr; }
};
class StatsCallback final : public DB::AsyncCallback {
public:
bool EnableStats() const override { return true; }
void OnComplete(const PerfContext* callback_perf_context,
const IOStatsContext* callback_iostats_context) override {
ASSERT_NE(nullptr, callback_perf_context);
ASSERT_NE(nullptr, callback_iostats_context);
called = true;
get_read_bytes = callback_perf_context->get_read_bytes;
multiget_read_bytes = callback_perf_context->multiget_read_bytes;
}
bool called = false;
uint64_t get_read_bytes = 0;
uint64_t multiget_read_bytes = 0;
};
auto file_system =
std::make_shared<NoReadExecutorFileSystem>(FileSystem::Default());
std::unique_ptr<Env> env = NewCompositeEnv(file_system);
Options options;
options.create_if_missing = true;
options.env = env.get();
std::unique_ptr<DB> db;
ASSERT_OK(DB::Open(options, kDbName, &db));
ASSERT_OK(db->Put(WriteOptions(), "key", "value"));
ASSERT_OK(db->Put(WriteOptions(), "key2", "value2"));
get_perf_context()->get_read_bytes = 13;
StatsCallback callback;
Status status;
std::string value;
db->GetAsync(ReadOptions(), "key", &value, status, callback);
EXPECT_TRUE(callback.called);
ASSERT_OK(status);
EXPECT_EQ("value", value);
EXPECT_EQ(value.size(), callback.get_read_bytes);
get_perf_context()->Reset();
get_perf_context()->multiget_read_bytes = 29;
StatsCallback multiget_callback;
ColumnFamilyHandle* column_families[] = {db->DefaultColumnFamily(),
db->DefaultColumnFamily()};
const std::string key_strings[] = {"key", "key2"};
const Slice keys[] = {key_strings[0], key_strings[1]};
PinnableSlice values[2];
Status statuses[2];
db->MultiGetAsync(ReadOptions(), 2, column_families, keys, values, statuses,
/*sorted_input=*/false, multiget_callback);
EXPECT_TRUE(multiget_callback.called);
EXPECT_OK(statuses[0]);
EXPECT_OK(statuses[1]);
EXPECT_EQ("value", values[0].ToString());
EXPECT_EQ("value2", values[1].ToString());
EXPECT_EQ(values[0].size() + values[1].size(),
multiget_callback.multiget_read_bytes);
db.reset();
ASSERT_OK(DestroyDB(kDbName, options));
}
TEST_F(PerfContextTest, SeekIntoDeletion) {
ASSERT_OK(DestroyDB(kDbName, Options()));

3
db_stress_tool/db_stress_common.cc
View File

@@ -42,7 +42,8 @@ namespace {
class BlockingAsyncCallback : public DB::AsyncCallback {
public:
void OnComplete() override {
void OnComplete(const PerfContext* /*perf_context*/,
const IOStatsContext* /*iostats_context*/) override {
std::lock_guard<std::mutex> lock(mu_);
done_ = true;
cv_.notify_one();

12
file/random_access_file_reader_sync_and_async.h
View File

@@ -94,10 +94,7 @@ DEFINE_SYNC_AND_ASYNC(IOStatus, RandomAccessFileReader::Read)
GetFileReadHistograms(stats_, opts.io_activity),
(stats_ != nullptr) ? &elapsed : nullptr, true /*overwrite*/,
true /*delay_enabled*/);
#if defined(WITHOUT_COROUTINES)
auto prev_perf_level = GetPerfLevel();
IOSTATS_TIMER_GUARD(read_nanos);
#endif
if (use_direct_io() && is_aligned == false) {
size_t aligned_offset =
TruncateToPageBoundary(alignment, static_cast<size_t>(offset));
@@ -267,9 +264,6 @@ DEFINE_SYNC_AND_ASYNC(IOStatus, RandomAccessFileReader::Read)
*result = Slice(res_scratch, io_s.ok() ? pos : 0);
}
RecordIOStats(stats_, file_temperature_, is_last_level_, result->size());
#if defined(WITHOUT_COROUTINES)
SetPerfLevel(prev_perf_level);
#endif
}
if (stats_ != nullptr && file_read_hist_ != nullptr) {
file_read_hist_->Add(elapsed);
@@ -324,10 +318,7 @@ DEFINE_SYNC_AND_ASYNC(IOStatus, RandomAccessFileReader::MultiRead)
GetFileReadHistograms(stats_, opts.io_activity),
(stats_ != nullptr) ? &elapsed : nullptr, true /*overwrite*/,
true /*delay_enabled*/);
#if defined(WITHOUT_COROUTINES)
auto prev_perf_level = GetPerfLevel();
IOSTATS_TIMER_GUARD(read_nanos);
#endif
FSReadRequest* fs_reqs = read_reqs;
size_t num_fs_reqs = num_reqs;
@@ -465,9 +456,6 @@ DEFINE_SYNC_AND_ASYNC(IOStatus, RandomAccessFileReader::MultiRead)
}
RecordIOStats(stats_, file_temperature_, is_last_level_, result_size);
}
#if defined(WITHOUT_COROUTINES)
SetPerfLevel(prev_perf_level);
#endif
}
if (stats_ != nullptr && file_read_hist_ != nullptr) {
file_read_hist_->Add(elapsed);

36
include/rocksdb/db.h
View File

@@ -44,7 +44,9 @@ struct DBOptions;
struct ExternalSstFileInfo;
struct FlushOptions;
struct FlushWALOptions;
struct IOStatsContext;
struct Options;
struct PerfContext;
struct ReadOptions;
struct TableProperties;
struct WriteOptions;
@@ -1115,11 +1117,28 @@ class DB {
//
// Callers must keep the DB, callback, inputs, and output buffers alive until
// the callback returns. The callback may run inline before the async method
// returns, or later from the implementation's completion path.
// returns, or later from the implementation's completion path. Callbacks must
// not invoke another async read.
//
// STATS:
// Callbacks can opt into perf and io metrics by overriding `EnableStats`.
// When provided, the returned contexts contain metrics for this request
// only. Most metrics should generally be available. Some scoped CPU metrics
// may be missing (e.g. `block_read_cpu_time`). Each returned context only has
// stats for a single operation, unlike the sync version which can re-use the
// same context for multiple operations. DO NOT use get_perf_context() or
// get_iostats_context() for statistics as you would for sync versions.
//
// Also enabling perf for async reads is generally more expensive because each
// request needs a separate stats context, rather than relying on the
// traditional TLS stats. Stats configuration (e.g. perf level) can be set
// before calling async read, and the config is saved by the coroutine.
class AsyncCallback {
public:
virtual ~AsyncCallback() = default;
virtual void OnComplete() = 0;
virtual bool EnableStats() const { return false; }
virtual void OnComplete(const PerfContext* callback_perf_context,
const IOStatsContext* callback_iostats_context) = 0;
};
virtual void GetAsync(const ReadOptions& options,
@@ -1127,7 +1146,8 @@ class DB {
PinnableSlice* value, std::string* timestamp,
Status& status, AsyncCallback& callback) {
status = Get(options, column_family, key, value, timestamp);
callback.OnComplete();
callback.OnComplete(/*callback_perf_context=*/nullptr,
/*callback_iostats_context=*/nullptr);
}
virtual void GetAsync(const ReadOptions& options,
@@ -1145,14 +1165,17 @@ class DB {
PinnableSlice* value() { return &pinnable_value_; }
void OnComplete() override {
bool EnableStats() const override { return callback_.EnableStats(); }
void OnComplete(const PerfContext* callback_perf_context,
const IOStatsContext* callback_iostats_context) override {
std::unique_ptr<CallbackWrapper> self(this);
if (status_.ok() && pinnable_value_.IsPinned()) {
value_->assign(pinnable_value_.data(), pinnable_value_.size());
}
AsyncCallback& callback = callback_;
self.reset();
callback.OnComplete();
callback.OnComplete(callback_perf_context, callback_iostats_context);
}
private:
@@ -1202,7 +1225,8 @@ class DB {
const bool sorted_input, AsyncCallback& callback) {
MultiGet(options, num_keys, column_families, keys, values, timestamps,
statuses, sorted_input);
callback.OnComplete();
callback.OnComplete(/*callback_perf_context=*/nullptr,
/*callback_iostats_context=*/nullptr);
}
virtual void MultiGetAsync(const ReadOptions& options, const size_t num_keys,

6
include/rocksdb/iostats_context.h
View File

@@ -74,11 +74,14 @@ struct FileIOByTemperature {
unknown_non_last_level_read_count = 0;
unknown_last_level_read_count = 0;
}
void Merge(const FileIOByTemperature& other);
};
struct IOStatsContext {
// reset all io-stats counter to zero
void Reset();
void Merge(const IOStatsContext& other);
std::string ToString(bool exclude_zero_counters = false) const;
@@ -108,7 +111,8 @@ struct IOStatsContext {
uint64_t logger_nanos;
// CPU time spent in write() and pwrite()
uint64_t cpu_write_nanos;
// CPU time spent in read() and pread()
// CPU time spent in read() and pread(). Not supported for async read
// requests.
uint64_t cpu_read_nanos;
FileIOByTemperature file_io_stats_by_temperature;

10
include/rocksdb/perf_context.h
View File

@@ -60,6 +60,7 @@ struct PerfContextByLevelBase {
// PerfContextByLevel
struct PerfContextByLevel : public PerfContextByLevelBase {
void Reset(); // reset all performance counters to zero
void Merge(const PerfContextByLevel& other);
};
/*
@@ -76,7 +77,8 @@ struct PerfContextBase {
uint64_t block_read_count; // total number of block reads (with IO)
uint64_t block_read_byte; // total number of bytes from block reads
uint64_t block_read_time; // total nanos spent on block reads
// total cpu time in nanos spent on block reads
// total cpu time in nanos spent on block reads. Not supported for async read
// requests.
uint64_t block_read_cpu_time;
uint64_t block_cache_index_hit_count; // total number of index block hits
// total number of standalone handles lookup from secondary cache
@@ -105,8 +107,8 @@ struct PerfContextBase {
// bytes for vals after compression in secondary cache
uint64_t compressed_sec_cache_compressed_bytes;
uint64_t block_checksum_time; // total nanos spent on block checksum
uint64_t block_decompress_time; // total nanos spent on block decompression
uint64_t block_checksum_time; // total nanos spent on block checksum
uint64_t block_decompress_time; // total nanos spent on block decompression
uint64_t block_decompress_count; // total number of block decompressions
uint64_t get_read_bytes; // bytes for vals returned by Get
@@ -328,8 +330,10 @@ struct PerfContext : public PerfContextBase {
PerfContext(const PerfContext&);
PerfContext& operator=(const PerfContext&);
PerfContext(PerfContext&&) noexcept;
PerfContext& operator=(PerfContext&&) noexcept;
void Reset(); // reset all performance counters to zero
void Merge(const PerfContext& other);
std::string ToString(bool exclude_zero_counters = false) const;

41
monitoring/iostats_context.cc
View File

@@ -20,6 +20,27 @@ thread_local IOStatsContext iostats_context;
IOStatsContext* get_iostats_context() { return &iostats_context; }
void FileIOByTemperature::Merge(const FileIOByTemperature& other) {
#ifndef NIOSTATS_CONTEXT
hot_file_bytes_read += other.hot_file_bytes_read;
warm_file_bytes_read += other.warm_file_bytes_read;
cool_file_bytes_read += other.cool_file_bytes_read;
cold_file_bytes_read += other.cold_file_bytes_read;
ice_file_bytes_read += other.ice_file_bytes_read;
unknown_non_last_level_bytes_read += other.unknown_non_last_level_bytes_read;
unknown_last_level_bytes_read += other.unknown_last_level_bytes_read;
hot_file_read_count += other.hot_file_read_count;
warm_file_read_count += other.warm_file_read_count;
cool_file_read_count += other.cool_file_read_count;
cold_file_read_count += other.cold_file_read_count;
ice_file_read_count += other.ice_file_read_count;
unknown_non_last_level_read_count += other.unknown_non_last_level_read_count;
unknown_last_level_read_count += other.unknown_last_level_read_count;
#else
(void)other;
#endif
}
void IOStatsContext::Reset() {
#ifndef NIOSTATS_CONTEXT
thread_pool_id = Env::Priority::TOTAL;
@@ -39,6 +60,26 @@ void IOStatsContext::Reset() {
#endif //! NIOSTATS_CONTEXT
}
void IOStatsContext::Merge(const IOStatsContext& other) {
#ifndef NIOSTATS_CONTEXT
bytes_read += other.bytes_read;
bytes_written += other.bytes_written;
open_nanos += other.open_nanos;
allocate_nanos += other.allocate_nanos;
write_nanos += other.write_nanos;
read_nanos += other.read_nanos;
range_sync_nanos += other.range_sync_nanos;
prepare_write_nanos += other.prepare_write_nanos;
fsync_nanos += other.fsync_nanos;
logger_nanos += other.logger_nanos;
cpu_write_nanos += other.cpu_write_nanos;
cpu_read_nanos += other.cpu_read_nanos;
file_io_stats_by_temperature.Merge(other.file_io_stats_by_temperature);
#else
(void)other;
#endif
}
#define IOSTATS_CONTEXT_OUTPUT(counter) \
if (!exclude_zero_counters || counter > 0) { \
ss << #counter << " = " << counter << ", "; \

57
monitoring/perf_context.cc
View File

@@ -5,6 +5,7 @@
//
#include <sstream>
#include <utility>
#include "monitoring/perf_context_imp.h"
@@ -224,12 +225,8 @@ PerfContext::PerfContext(const PerfContext& other) {
#endif
}
PerfContext::PerfContext(PerfContext&& other) noexcept {
#ifdef NPERF_CONTEXT
(void)other;
#else
copyMetrics(&other);
#endif
PerfContext::PerfContext(PerfContext&& other) noexcept : PerfContext() {
*this = std::move(other);
}
PerfContext& PerfContext::operator=(const PerfContext& other) {
@@ -241,6 +238,24 @@ PerfContext& PerfContext::operator=(const PerfContext& other) {
return *this;
}
PerfContext& PerfContext::operator=(PerfContext&& other) noexcept {
#ifdef NPERF_CONTEXT
(void)other;
#else
if (this != &other) {
ClearPerLevelPerfContext();
#define EMIT_MOVE_FIELDS(x) x = other.x;
DEF_PERF_CONTEXT_METRICS(EMIT_MOVE_FIELDS)
#undef EMIT_MOVE_FIELDS
level_to_perf_context = other.level_to_perf_context;
per_level_perf_context_enabled = other.per_level_perf_context_enabled;
other.level_to_perf_context = nullptr;
other.per_level_perf_context_enabled = false;
}
#endif
return *this;
}
void PerfContext::copyMetrics(const PerfContext* other) noexcept {
#ifdef NPERF_CONTEXT
(void)other;
@@ -272,6 +287,26 @@ void PerfContext::Reset() {
#endif
}
void PerfContext::Merge(const PerfContext& other) {
#ifndef NPERF_CONTEXT
#define EMIT_FIELDS(x) x += other.x;
DEF_PERF_CONTEXT_METRICS(EMIT_FIELDS)
#undef EMIT_FIELDS
if (other.level_to_perf_context != nullptr) {
if (level_to_perf_context == nullptr) {
level_to_perf_context = new std::map<uint32_t, PerfContextByLevel>();
}
for (const auto& kv : *other.level_to_perf_context) {
(*level_to_perf_context)[kv.first].Merge(kv.second);
}
}
per_level_perf_context_enabled =
per_level_perf_context_enabled || other.per_level_perf_context_enabled;
#else
(void)other;
#endif
}
void PerfContextByLevel::Reset() {
#ifndef NPERF_CONTEXT
#define EMIT_FIELDS(x) x = 0;
@@ -280,6 +315,16 @@ void PerfContextByLevel::Reset() {
#endif
}
void PerfContextByLevel::Merge(const PerfContextByLevel& other) {
#ifndef NPERF_CONTEXT
#define EMIT_FIELDS(x) x += other.x;
DEF_PERF_CONTEXT_LEVEL_METRICS(EMIT_FIELDS)
#undef EMIT_FIELDS
#else
(void)other;
#endif
}
std::string PerfContext::ToString(bool exclude_zero_counters) const {
#ifdef NPERF_CONTEXT
(void)exclude_zero_counters;

1
src.mk
View File

@@ -253,6 +253,7 @@ LIB_SOURCES = \
util/compression.cc \
util/compression_context_cache.cc \
util/concurrent_task_limiter_impl.cc \
util/coro_stats_util.cc \
util/crc32c.cc \
util/crc32c_arm64.cc \
util/data_structure.cc \

2
table/block_based/block_based_table_reader_sync_and_async.h
View File

@@ -331,9 +331,7 @@ DEFINE_SYNC_AND_ASYNC(TBlockIter*, BlockBasedTable::NewDataBlockIterator)(
using IterBlocklike =
std::conditional_t<std::is_same_v<TBlockIter, DataBlockIter>, Block_kData,
Block_kIndex>;
#if defined(WITHOUT_COROUTINES)
PERF_TIMER_GUARD(new_table_block_iter_nanos);
#endif // defined(WITHOUT_COROUTINES)
TBlockIter* iter = input_iter != nullptr ? input_iter : new TBlockIter;
if (!s.ok()) {

6
table/block_fetcher_sync_and_async.h
View File

@@ -35,8 +35,8 @@ DEFINE_SYNC_AND_ASYNC(void, BlockFetcher::ReadBlock)(bool retry) {
AlignedBufferAllocationContext direct_io_context{
&direct_io_buffer_,
direct_io_allocator ? &direct_io_allocator : nullptr};
#if defined(WITHOUT_COROUTINES)
PERF_TIMER_GUARD(block_read_time);
#if defined(WITHOUT_COROUTINES)
PERF_CPU_TIMER_GUARD(
block_read_cpu_time,
ioptions_.env ? ioptions_.env->GetSystemClock().get() : nullptr);
@@ -47,8 +47,8 @@ DEFINE_SYNC_AND_ASYNC(void, BlockFetcher::ReadBlock)(bool retry) {
PERF_COUNTER_ADD(block_read_count, 1);
used_buf_ = const_cast<char*>(slice_.data());
} else if (use_fs_scratch_) {
#if defined(WITHOUT_COROUTINES)
PERF_TIMER_GUARD(block_read_time);
#if defined(WITHOUT_COROUTINES)
PERF_CPU_TIMER_GUARD(
block_read_cpu_time,
ioptions_.env ? ioptions_.env->GetSystemClock().get() : nullptr);
@@ -71,8 +71,8 @@ DEFINE_SYNC_AND_ASYNC(void, BlockFetcher::ReadBlock)(bool retry) {
CO_RETURN;
}
#if defined(WITHOUT_COROUTINES)
PERF_TIMER_GUARD(block_read_time);
#if defined(WITHOUT_COROUTINES)
PERF_CPU_TIMER_GUARD(
block_read_cpu_time,
ioptions_.env ? ioptions_.env->GetSystemClock().get() : nullptr);

28
test_util/mock_time_env.h
View File

@@ -41,6 +41,20 @@ class MockSystemClock : public SystemClockWrapper {
return current_time_us_ * 1000;
}
uint64_t CPUMicros() override {
if (!mock_cpu_time_.load(std::memory_order_relaxed)) {
return SystemClockWrapper::CPUMicros();
}
return current_cpu_time_ns_.load(std::memory_order_relaxed) / 1000;
}
uint64_t CPUNanos() override {
if (!mock_cpu_time_.load(std::memory_order_relaxed)) {
return SystemClockWrapper::CPUNanos();
}
return current_cpu_time_ns_.load(std::memory_order_relaxed);
}
uint64_t RealNowMicros() { return target_->NowMicros(); }
void SetCurrentTime(uint64_t time_sec) {
@@ -69,6 +83,18 @@ class MockSystemClock : public SystemClockWrapper {
current_time_us_.fetch_add(micros);
}
void SetCurrentCPUTimeNanos(uint64_t nanos) {
current_cpu_time_ns_.store(nanos, std::memory_order_relaxed);
mock_cpu_time_.store(true, std::memory_order_relaxed);
}
void MockSleepForCPUNanos(uint64_t nanos) {
assert(mock_cpu_time_.load(std::memory_order_relaxed));
assert(current_cpu_time_ns_.load(std::memory_order_relaxed) <=
std::numeric_limits<uint64_t>::max() - nanos);
current_cpu_time_ns_.fetch_add(nanos, std::memory_order_relaxed);
}
bool TimedWait(port::CondVar* cv,
std::chrono::microseconds deadline) override {
uint64_t now_micros = NowMicros();
@@ -103,6 +129,8 @@ class MockSystemClock : public SystemClockWrapper {
private:
std::atomic<uint64_t> current_time_us_{0};
std::atomic<uint64_t> current_cpu_time_ns_{0};
std::atomic<bool> mock_cpu_time_{false};
static constexpr uint64_t kMicrosInSecond = 1000U * 1000U;
};

109
tools/db_bench_tool.cc
View File

@@ -112,10 +112,12 @@
#if USE_COROUTINES
#include "folly/coro/Baton.h"
#include "folly/coro/CurrentExecutor.h"
#include "folly/coro/Task.h"
#include "folly/executors/IOExecutor.h"
#include "folly/futures/Future.h"
#include "folly/io/async/EventBase.h"
#include "folly/io/async/Request.h"
#endif // USE_COROUTINES
using GFLAGS_NAMESPACE::ParseCommandLineFlags;
@@ -7609,37 +7611,83 @@ class Benchmark {
}
struct AsyncAwaitCallback : public DB::AsyncCallback {
void OnComplete() override { baton.post(); }
AsyncAwaitCallback(
PerfContext* job_perf_context,
std::shared_ptr<folly::RequestContext> caller_request_context)
: job_perf_context_(job_perf_context),
caller_request_context_(std::move(caller_request_context)) {}
bool EnableStats() const override { return job_perf_context_ != nullptr; }
void OnComplete(const PerfContext* callback_perf_context,
const IOStatsContext* /*iostats_context*/) override {
if (job_perf_context_ != nullptr && callback_perf_context != nullptr) {
job_perf_context_->Merge(*callback_perf_context);
}
folly::RequestContextScopeGuard context_scope(caller_request_context_);
baton.post();
}
PerfContext* job_perf_context_;
std::shared_ptr<folly::RequestContext> caller_request_context_;
folly::coro::Baton baton;
};
void PrepareCoroutineJobPerfContext(PerfContext* job_perf_context) {
SetPerfLevel(static_cast<PerfLevel>(FLAGS_perf_level));
if (job_perf_context == nullptr) {
return;
}
#ifndef NPERF_CONTEXT
get_perf_context()->EnablePerLevelPerfContext();
#endif
}
folly::coro::Task<Status> AwaitGetAsync(DB* db, const ReadOptions& options,
ColumnFamilyHandle* cfh,
const Slice& key,
PinnableSlice* value,
std::string* timestamp) {
AsyncAwaitCallback callback;
std::string* timestamp,
PerfContext* job_perf_context) {
PrepareCoroutineJobPerfContext(job_perf_context);
AsyncAwaitCallback callback(job_perf_context,
folly::RequestContext::saveContext());
Status status;
db->GetAsync(options, cfh, key, value, timestamp, status, callback);
co_await callback.baton;
co_await folly::coro::co_reschedule_on_current_executor;
co_return status;
}
folly::coro::Task<void> AwaitMultiGetAsync(
DB* db, const ReadOptions& options, size_t num_keys,
ColumnFamilyHandle** column_families, const Slice* keys,
PinnableSlice* values, Status* statuses) {
AsyncAwaitCallback callback;
PinnableSlice* values, Status* statuses, PerfContext* job_perf_context) {
PrepareCoroutineJobPerfContext(job_perf_context);
AsyncAwaitCallback callback(job_perf_context,
folly::RequestContext::saveContext());
db->MultiGetAsync(options, num_keys, column_families, keys, values,
statuses, callback);
co_await callback.baton;
co_await folly::coro::co_reschedule_on_current_executor;
co_return;
}
void SetCoroutineJobPerfContextMessage(
std::vector<std::string>* job_perf_contexts, int job_id,
const char* benchmark_name, const PerfContext* job_perf_context) {
if (job_perf_contexts == nullptr || job_perf_context == nullptr) {
return;
}
(*job_perf_contexts)[job_id] = std::string(benchmark_name) + " job " +
std::to_string(job_id) + " PERF_CONTEXT:\n" +
job_perf_context->ToString();
}
folly::coro::Task<void> ReadRandomCoroutineJob(
uint64_t seed, int64_t ops, std::atomic<int64_t>* total_ops,
std::atomic<int64_t>* total_found, std::atomic<int64_t>* total_bytes) {
std::atomic<int64_t>* total_found, std::atomic<int64_t>* total_bytes,
int job_id, std::vector<std::string>* job_perf_contexts) {
Random64 rng(seed);
ReadOptions options = read_options_;
std::unique_ptr<const char[]> key_guard;
@@ -7654,6 +7702,7 @@ class Benchmark {
int64_t job_ops = 0;
int64_t job_found = 0;
int64_t job_bytes = 0;
PerfContext job_perf_context;
for (int64_t done = 0; done < ops; ++done) {
DBWithColumnFamilies* db_with_cfh = SelectDBWithCfh(rng.Next());
const int64_t key_rand = GetRandomKey(&rng);
@@ -7675,8 +7724,9 @@ class Benchmark {
cfh = db_with_cfh->db->DefaultColumnFamily();
}
Status s = co_await AwaitGetAsync(db_with_cfh->db, options, cfh, key,
&pinnable_val, ts_ptr);
Status s = co_await AwaitGetAsync(
db_with_cfh->db, options, cfh, key, &pinnable_val, ts_ptr,
job_perf_contexts != nullptr ? &job_perf_context : nullptr);
++job_ops;
if (s.ok()) {
++job_found;
@@ -7691,6 +7741,8 @@ class Benchmark {
total_ops->fetch_add(job_ops, std::memory_order_relaxed);
total_found->fetch_add(job_found, std::memory_order_relaxed);
total_bytes->fetch_add(job_bytes, std::memory_order_relaxed);
SetCoroutineJobPerfContextMessage(job_perf_contexts, job_id,
"readrandomcoroutine", &job_perf_context);
co_return;
}
@@ -7701,7 +7753,8 @@ class Benchmark {
folly::coro::Task<void> MultiGetAsyncCoroutineJob(
DB* db, ColumnFamilyHandle* cfh, uint64_t seed, int64_t ops,
std::atomic<int64_t>* total_ops, std::atomic<int64_t>* total_found,
std::atomic<int64_t>* total_bytes) {
std::atomic<int64_t>* total_bytes, int job_id,
std::vector<std::string>* job_perf_contexts) {
Random64 rng(seed);
ReadOptions options = read_options_;
std::vector<std::unique_ptr<const char[]>> key_guards(entries_per_batch_);
@@ -7717,14 +7770,17 @@ class Benchmark {
int64_t job_ops = 0;
int64_t job_found = 0;
int64_t job_bytes = 0;
PerfContext job_perf_context;
for (int64_t done = 0; done < ops; done += entries_per_batch_) {
for (int64_t i = 0; i < entries_per_batch_; ++i) {
GenerateKeyFromInt(GetRandomKey(&rng), FLAGS_num, &keys[i]);
statuses[i] = Status::OK();
values[i].Reset();
}
co_await AwaitMultiGetAsync(db, options, entries_per_batch_, cfs.data(),
keys.data(), values.get(), statuses.data());
co_await AwaitMultiGetAsync(
db, options, entries_per_batch_, cfs.data(), keys.data(),
values.get(), statuses.data(),
job_perf_contexts != nullptr ? &job_perf_context : nullptr);
job_ops += entries_per_batch_;
for (int64_t i = 0; i < entries_per_batch_; ++i) {
if (statuses[i].ok()) {
@@ -7740,6 +7796,9 @@ class Benchmark {
total_ops->fetch_add(job_ops, std::memory_order_relaxed);
total_found->fetch_add(job_found, std::memory_order_relaxed);
total_bytes->fetch_add(job_bytes, std::memory_order_relaxed);
SetCoroutineJobPerfContextMessage(job_perf_contexts, job_id,
"multireadrandomcoroutine",
&job_perf_context);
co_return;
}
#endif // USE_COROUTINES
@@ -7760,6 +7819,14 @@ class Benchmark {
std::atomic<int64_t> total_ops{0};
std::atomic<int64_t> total_found{0};
std::atomic<int64_t> total_bytes{0};
std::vector<std::string> job_perf_contexts;
std::vector<std::string>* job_perf_contexts_ptr = nullptr;
#ifndef NPERF_CONTEXT
if (FLAGS_perf_level > ROCKSDB_NAMESPACE::PerfLevel::kDisable) {
job_perf_contexts.resize(num_jobs);
job_perf_contexts_ptr = &job_perf_contexts;
}
#endif
// Spread reads_ * pool_threads ops across the jobs
const int64_t total_reads = reads_ * pool_threads;
@@ -7778,7 +7845,8 @@ class Benchmark {
folly::coro::co_withExecutor(
folly::Executor::getKeepAliveToken(read_executor->getEventBase()),
ReadRandomCoroutineJob(thread->rand.Next(), ops, &total_ops,
&total_found, &total_bytes))
&total_found, &total_bytes, j,
job_perf_contexts_ptr))
.start());
}
for (auto& f : futures) {
@@ -7791,6 +7859,9 @@ class Benchmark {
total_found.load(), total_ops.load());
thread->stats.AddBytes(total_bytes.load());
thread->stats.AddMessage(msg);
if (!job_perf_contexts.empty() && !job_perf_contexts.front().empty()) {
thread->stats.AddMessage(job_perf_contexts.front());
}
#else // USE_COROUTINES
(void)thread;
fprintf(stderr,
@@ -7819,6 +7890,14 @@ class Benchmark {
std::atomic<int64_t> total_ops{0};
std::atomic<int64_t> total_found{0};
std::atomic<int64_t> total_bytes{0};
std::vector<std::string> job_perf_contexts;
std::vector<std::string>* job_perf_contexts_ptr = nullptr;
#ifndef NPERF_CONTEXT
if (FLAGS_perf_level > ROCKSDB_NAMESPACE::PerfLevel::kDisable) {
job_perf_contexts.resize(num_jobs);
job_perf_contexts_ptr = &job_perf_contexts;
}
#endif
const int64_t total_reads = reads_ * pool_threads;
if (total_reads % num_jobs != 0) {
@@ -7841,7 +7920,8 @@ class Benchmark {
folly::coro::co_withExecutor(
folly::Executor::getKeepAliveToken(read_executor->getEventBase()),
MultiGetAsyncCoroutineJob(db, cfh, thread->rand.Next(), ops,
&total_ops, &total_found, &total_bytes))
&total_ops, &total_found, &total_bytes,
j, job_perf_contexts_ptr))
.start());
}
for (auto& f : futures) {
@@ -7854,6 +7934,9 @@ class Benchmark {
total_found.load(), total_ops.load());
thread->stats.AddBytes(total_bytes.load());
thread->stats.AddMessage(msg);
if (!job_perf_contexts.empty() && !job_perf_contexts.front().empty()) {
thread->stats.AddMessage(job_perf_contexts.front());
}
#else // USE_COROUTINES
(void)thread;
fprintf(stderr,

213
util/coro_stats_util.cc Normal file
View File

@@ -0,0 +1,213 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under both the GPLv2 (found in the
// COPYING file in the root directory) and Apache 2.0 License
// (found in the LICENSE.Apache file in the root directory).
#include "util/coro_stats_util.h"
#if defined(USE_COROUTINES)
#include <cassert>
#include <cstdint>
#include <memory>
#include <utility>
#include "folly/io/async/Request.h"
#include "monitoring/iostats_context_imp.h"
#include "monitoring/perf_context_imp.h"
#include "rocksdb/env.h"
namespace ROCKSDB_NAMESPACE {
namespace {
struct CoroutineStatsRequestDataTraits {
static const folly::RequestToken kToken;
};
const folly::RequestToken CoroutineStatsRequestDataTraits::kToken(
"rocksdb_coroutine_stats_context");
class CoroutineStatsRequestData;
#ifndef NDEBUG
CoroutineStatsRequestData* GetCoroutineStatsData();
#endif
#ifndef NPERF_CONTEXT
struct CoroutineStatsThreadLocalState {
CoroutineStatsRequestData* data = nullptr;
uint64_t get_cpu_nanos_start = 0;
};
thread_local CoroutineStatsThreadLocalState coroutine_stats_thread_local_state;
#endif
class CoroutineStatsRequestData : public folly::RequestData {
public:
explicit CoroutineStatsRequestData(CoroutineStatsConfig stats_config,
Env* env)
: env_(env), perf_level_(stats_config.perf_level) {
(void)env_;
assert(env_ != nullptr);
#ifndef NDEBUG
owner_thread_id_ = env_->GetThreadID();
#endif
assert(CapturedPerfLevel() > PerfLevel::kUninitialized);
assert(CapturedPerfLevel() < PerfLevel::kOutOfBounds);
#ifndef NPERF_CONTEXT
if (stats_config.per_level_perf_context_enabled) {
perf_context_.EnablePerLevelPerfContext();
}
#endif
#ifndef NIOSTATS_CONTEXT
iostats_context_.Reset();
iostats_context_.disable_iostats = stats_config.iostats_disabled;
#endif
}
bool hasCallback() override { return true; }
void onSet() override {
AssertSameThread();
SetPerfLevel(CapturedPerfLevel());
LoadThreadLocalStats();
StartGetCpuTimer();
}
void onUnset() override {
StopGetCpuTimer();
SaveThreadLocalStats();
}
bool PerfLevelAtLeast(PerfLevel perf_level) const {
return CapturedPerfLevel() >= perf_level;
}
PerfLevel CapturedPerfLevel() const { return perf_level_; }
void Finalize() {
assert(GetCoroutineStatsData() == this);
assert(!finalized_);
finalized_ = true;
StopGetCpuTimer();
}
private:
void LoadThreadLocalStats() {
#ifndef NPERF_CONTEXT
*get_perf_context() = std::move(perf_context_);
#endif
#ifndef NIOSTATS_CONTEXT
*get_iostats_context() = std::move(iostats_context_);
#endif
}
void SaveThreadLocalStats() {
#ifndef NPERF_CONTEXT
perf_context_ = std::move(*get_perf_context());
#endif
#ifndef NIOSTATS_CONTEXT
iostats_context_ = std::move(*get_iostats_context());
#endif
}
// Folly event bases resume requests on the thread where they suspended.
void AssertSameThread() {
#ifndef NDEBUG
const uint64_t current_thread_id = env_->GetThreadID();
assert(owner_thread_id_ == current_thread_id);
#endif
}
void StartGetCpuTimer() {
#ifndef NPERF_CONTEXT
if (finalized_) {
return;
}
if (!PerfLevelAtLeast(PerfLevel::kEnableTimeAndCPUTimeExceptForMutex)) {
return;
}
coroutine_stats_thread_local_state.data = this;
coroutine_stats_thread_local_state.get_cpu_nanos_start =
env_->GetSystemClock()->CPUNanos();
#endif
}
void StopGetCpuTimer() {
#ifndef NPERF_CONTEXT
if (coroutine_stats_thread_local_state.data != this) {
return;
}
const uint64_t start =
coroutine_stats_thread_local_state.get_cpu_nanos_start;
coroutine_stats_thread_local_state = {};
const uint64_t now = env_->GetSystemClock()->CPUNanos();
if (now >= start) {
get_perf_context()->get_cpu_nanos += now - start;
}
#endif
}
Env* const env_;
PerfLevel perf_level_;
#ifndef NPERF_CONTEXT
PerfContext perf_context_;
#endif
#ifndef NIOSTATS_CONTEXT
IOStatsContext iostats_context_;
#endif
bool finalized_ = false;
#ifndef NDEBUG
uint64_t owner_thread_id_ = 0;
#endif
};
#ifndef NDEBUG
CoroutineStatsRequestData* GetCoroutineStatsData() {
auto* context = folly::RequestContext::try_get();
if (context == nullptr) {
return nullptr;
}
return static_cast<CoroutineStatsRequestData*>(
context->getThreadCachedContextData<CoroutineStatsRequestDataTraits>());
}
#endif
} // namespace
CoroutineStatsConfig CaptureCoroutineStatsConfig() {
CoroutineStatsConfig stats_config;
stats_config.perf_level = GetPerfLevel();
#ifndef NPERF_CONTEXT
stats_config.per_level_perf_context_enabled =
get_perf_context()->per_level_perf_context_enabled;
#endif
#ifndef NIOSTATS_CONTEXT
stats_config.iostats_disabled = get_iostats_context()->disable_iostats;
#endif
return stats_config;
}
CoroutineStatsContextScope::CoroutineStatsContextScope(
CoroutineStatsConfig stats_config, Env* env)
: request_data_(nullptr) {
assert(GetCoroutineStatsData() == nullptr); // NO nesting allowed
auto data =
std::make_unique<CoroutineStatsRequestData>(std::move(stats_config), env);
request_data_ = data.get();
guard_ = std::make_unique<folly::ShallowCopyRequestContextScopeGuard>(
CoroutineStatsRequestDataTraits::kToken, std::move(data));
}
CoroutineStatsContextScope::~CoroutineStatsContextScope() = default;
void CoroutineStatsContextScope::Finalize() const {
auto* data = static_cast<CoroutineStatsRequestData*>(request_data_);
assert(data != nullptr);
data->Finalize();
}
} // namespace ROCKSDB_NAMESPACE
#endif // defined(USE_COROUTINES)

58
util/coro_stats_util.h Normal file
View File

@@ -0,0 +1,58 @@
// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is licensed under both the GPLv2 (found in the
// COPYING file in the root directory) and Apache 2.0 License
// (found in the LICENSE.Apache file in the root directory).
#pragma once
#if defined(USE_COROUTINES)
#include <memory>
#include "rocksdb/perf_level.h"
#include "rocksdb/rocksdb_namespace.h"
namespace folly {
class RequestData;
struct ShallowCopyRequestContextScopeGuard;
} // namespace folly
namespace ROCKSDB_NAMESPACE {
class Env;
struct CoroutineStatsConfig {
PerfLevel perf_level = PerfLevel::kEnableCount;
bool per_level_perf_context_enabled = false;
bool iostats_disabled = false;
};
CoroutineStatsConfig CaptureCoroutineStatsConfig();
// Installs request-scoped perf and IO stats while a coroutine is active. Folly
// request context is used to save the request stats on suspension and reload
// them with the captured configuration on resumption, so multiple coroutines
// can share a single executor thread, but each own separate stats contexts.
// Call Finalize() once operation has completed, to ensure stats are flushed.
class CoroutineStatsContextScope {
public:
explicit CoroutineStatsContextScope(CoroutineStatsConfig stats_config,
Env* env);
~CoroutineStatsContextScope();
CoroutineStatsContextScope(const CoroutineStatsContextScope&) = delete;
CoroutineStatsContextScope& operator=(const CoroutineStatsContextScope&) =
delete;
CoroutineStatsContextScope(CoroutineStatsContextScope&&) = delete;
CoroutineStatsContextScope& operator=(CoroutineStatsContextScope&&) = delete;
void Finalize() const;
private:
folly::RequestData* request_data_;
std::unique_ptr<folly::ShallowCopyRequestContextScopeGuard> guard_;
};
} // namespace ROCKSDB_NAMESPACE
#endif // USE_COROUTINES