diff --git a/BUCK b/BUCK index 6a17780bcc..afd3eee076 100644 --- a/BUCK +++ b/BUCK @@ -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) diff --git a/CMakeLists.txt b/CMakeLists.txt index 9d2febd9a9..7b684311de 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -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 diff --git a/buckifier/buckify_rocksdb.py b/buckifier/buckify_rocksdb.py index 0fb48f2e30..0c9983e387 100755 --- a/buckifier/buckify_rocksdb.py +++ b/buckifier/buckify_rocksdb.py @@ -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\"])") diff --git a/db/db_impl/db_impl.cc b/db/db_impl/db_impl.cc index 12928088a1..d9fc97b53c 100644 --- a/db/db_impl/db_impl.cc +++ b/db/db_impl/db_impl.cc @@ -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 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 { - 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 { - 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 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 { + std::optional 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 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 { + std::optional 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( diff --git a/db/db_impl/db_impl_readonly_sync_and_async.h b/db/db_impl/db_impl_readonly_sync_and_async.h index 088449a90d..80a12fa4b9 100644 --- a/db/db_impl/db_impl_readonly_sync_and_async.h +++ b/db/db_impl/db_impl_readonly_sync_and_async.h @@ -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 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, diff --git a/db/db_impl/db_impl_secondary_sync_and_async.h b/db/db_impl/db_impl_secondary_sync_and_async.h index bb6984fe3f..3af9156565 100644 --- a/db/db_impl/db_impl_secondary_sync_and_async.h +++ b/db/db_impl/db_impl_secondary_sync_and_async.h @@ -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 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 diff --git a/db/db_impl/db_impl_sync_and_async.h b/db/db_impl/db_impl_sync_and_async.h index 984827a3dd..65bbf80bcd 100644 --- a/db/db_impl/db_impl_sync_and_async.h +++ b/db/db_impl/db_impl_sync_and_async.h @@ -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( 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; } diff --git a/db/perf_context_test.cc b/db/perf_context_test.cc index fbd38c6c26..7e6844aecc 100644 --- a/db/perf_context_test.cc +++ b/db/perf_context_test.cc @@ -7,23 +7,41 @@ #include #include +#include #include #include +#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 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(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 { + 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 { + 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 { + 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> 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 results = + co_await folly::coro::collectAllRange(std::move(tasks)); + + EXPECT_EQ((std::vector{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(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& 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(FileSystem::Default()); + std::unique_ptr env = NewCompositeEnv(file_system); + Options options; + options.create_if_missing = true; + options.env = env.get(); + + std::unique_ptr 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())); diff --git a/db_stress_tool/db_stress_common.cc b/db_stress_tool/db_stress_common.cc index 53d7a545dd..0bac46b095 100644 --- a/db_stress_tool/db_stress_common.cc +++ b/db_stress_tool/db_stress_common.cc @@ -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 lock(mu_); done_ = true; cv_.notify_one(); diff --git a/file/random_access_file_reader_sync_and_async.h b/file/random_access_file_reader_sync_and_async.h index bd8fc781d4..0db0f2da73 100644 --- a/file/random_access_file_reader_sync_and_async.h +++ b/file/random_access_file_reader_sync_and_async.h @@ -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(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); diff --git a/include/rocksdb/db.h b/include/rocksdb/db.h index cecd406765..982e6b2f10 100644 --- a/include/rocksdb/db.h +++ b/include/rocksdb/db.h @@ -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 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, diff --git a/include/rocksdb/iostats_context.h b/include/rocksdb/iostats_context.h index 8fce6181c0..550fd16bb8 100644 --- a/include/rocksdb/iostats_context.h +++ b/include/rocksdb/iostats_context.h @@ -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; diff --git a/include/rocksdb/perf_context.h b/include/rocksdb/perf_context.h index e446bbc798..91bb1e8f56 100644 --- a/include/rocksdb/perf_context.h +++ b/include/rocksdb/perf_context.h @@ -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; diff --git a/monitoring/iostats_context.cc b/monitoring/iostats_context.cc index 9f96655a6b..f618fe1b2f 100644 --- a/monitoring/iostats_context.cc +++ b/monitoring/iostats_context.cc @@ -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 << ", "; \ diff --git a/monitoring/perf_context.cc b/monitoring/perf_context.cc index 8d66731079..3e38d30f4d 100644 --- a/monitoring/perf_context.cc +++ b/monitoring/perf_context.cc @@ -5,6 +5,7 @@ // #include +#include #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(); + } + 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; diff --git a/src.mk b/src.mk index 604ecd22a4..196a8b305e 100644 --- a/src.mk +++ b/src.mk @@ -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 \ diff --git a/table/block_based/block_based_table_reader_sync_and_async.h b/table/block_based/block_based_table_reader_sync_and_async.h index 5b875eb27b..bf0e444b9e 100644 --- a/table/block_based/block_based_table_reader_sync_and_async.h +++ b/table/block_based/block_based_table_reader_sync_and_async.h @@ -331,9 +331,7 @@ DEFINE_SYNC_AND_ASYNC(TBlockIter*, BlockBasedTable::NewDataBlockIterator)( using IterBlocklike = std::conditional_t, 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()) { diff --git a/table/block_fetcher_sync_and_async.h b/table/block_fetcher_sync_and_async.h index 85a5bcf66b..6c1176e5cc 100644 --- a/table/block_fetcher_sync_and_async.h +++ b/table/block_fetcher_sync_and_async.h @@ -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(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); diff --git a/test_util/mock_time_env.h b/test_util/mock_time_env.h index 701a92d6e5..fc314ed9f7 100644 --- a/test_util/mock_time_env.h +++ b/test_util/mock_time_env.h @@ -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::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 current_time_us_{0}; + std::atomic current_cpu_time_ns_{0}; + std::atomic mock_cpu_time_{false}; static constexpr uint64_t kMicrosInSecond = 1000U * 1000U; }; diff --git a/tools/db_bench_tool.cc b/tools/db_bench_tool.cc index 42dd5119e4..22856a5517 100644 --- a/tools/db_bench_tool.cc +++ b/tools/db_bench_tool.cc @@ -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 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 caller_request_context_; folly::coro::Baton baton; }; + void PrepareCoroutineJobPerfContext(PerfContext* job_perf_context) { + SetPerfLevel(static_cast(FLAGS_perf_level)); + if (job_perf_context == nullptr) { + return; + } +#ifndef NPERF_CONTEXT + get_perf_context()->EnablePerLevelPerfContext(); +#endif + } + folly::coro::Task 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 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* 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 ReadRandomCoroutineJob( uint64_t seed, int64_t ops, std::atomic* total_ops, - std::atomic* total_found, std::atomic* total_bytes) { + std::atomic* total_found, std::atomic* total_bytes, + int job_id, std::vector* job_perf_contexts) { Random64 rng(seed); ReadOptions options = read_options_; std::unique_ptr 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 MultiGetAsyncCoroutineJob( DB* db, ColumnFamilyHandle* cfh, uint64_t seed, int64_t ops, std::atomic* total_ops, std::atomic* total_found, - std::atomic* total_bytes) { + std::atomic* total_bytes, int job_id, + std::vector* job_perf_contexts) { Random64 rng(seed); ReadOptions options = read_options_; std::vector> 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 total_ops{0}; std::atomic total_found{0}; std::atomic total_bytes{0}; + std::vector job_perf_contexts; + std::vector* 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 total_ops{0}; std::atomic total_found{0}; std::atomic total_bytes{0}; + std::vector job_perf_contexts; + std::vector* 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, diff --git a/util/coro_stats_util.cc b/util/coro_stats_util.cc new file mode 100644 index 0000000000..8e7b1c1f10 --- /dev/null +++ b/util/coro_stats_util.cc @@ -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 +#include +#include +#include + +#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( + context->getThreadCachedContextData()); +} +#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(std::move(stats_config), env); + request_data_ = data.get(); + guard_ = std::make_unique( + CoroutineStatsRequestDataTraits::kToken, std::move(data)); +} + +CoroutineStatsContextScope::~CoroutineStatsContextScope() = default; + +void CoroutineStatsContextScope::Finalize() const { + auto* data = static_cast(request_data_); + assert(data != nullptr); + data->Finalize(); +} + +} // namespace ROCKSDB_NAMESPACE + +#endif // defined(USE_COROUTINES) diff --git a/util/coro_stats_util.h b/util/coro_stats_util.h new file mode 100644 index 0000000000..278853f454 --- /dev/null +++ b/util/coro_stats_util.h @@ -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 + +#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 guard_; +}; + +} // namespace ROCKSDB_NAMESPACE + +#endif // USE_COROUTINES