diff --git a/.github/workflows/pr-jobs.yml b/.github/workflows/pr-jobs.yml index 3856228395..3c213e2e4c 100644 --- a/.github/workflows/pr-jobs.yml +++ b/.github/workflows/pr-jobs.yml @@ -426,13 +426,20 @@ jobs: steps: - uses: actions/checkout@v4.1.0 - uses: "./.github/actions/pre-steps" + - uses: "./.github/actions/cache-getdeps-downloads" + - uses: "./.github/actions/setup-folly" - uses: "./.github/actions/setup-ccache" with: cache-key-prefix: mini-crashtest + - uses: "./.github/actions/cache-folly" + id: cache-folly + - uses: "./.github/actions/build-folly" + with: + cache-hit: ${{ steps.cache-folly.outputs.cache-hit }} - run: | ulimit -S -n $(ulimit -Hn) # Raise open files soft limit to hard limit echo "Updated ulimit -Hn: $(ulimit -Hn) -Sn: $(ulimit -Sn)" - make V=1 -j8 CRASH_TEST_EXT_ARGS='--duration=${{ matrix.crash_duration }} --max_key=2500000' ${{ matrix.crash_test_target }} + USE_COROUTINES=1 USE_FOLLY=1 LIB_MODE=static make V=1 -j8 CRASH_TEST_EXT_ARGS='--duration=${{ matrix.crash_duration }} --max_key=2500000' ${{ matrix.crash_test_target }} - uses: "./.github/actions/teardown-ccache" if: always() - uses: "./.github/actions/post-steps" diff --git a/Makefile b/Makefile index 04e77ecdf1..2b1418f42a 100644 --- a/Makefile +++ b/Makefile @@ -589,9 +589,12 @@ TESTS += $(PLUGIN_TESTS) ifneq ($(filter check-headers, $(MAKECMDGOALS)),) # TODO: add/support JNI headers DEV_HEADER_DIRS := $(sort include/ $(dir $(ALL_SOURCES))) + FOLLY_DEPENDENT_HEADERS := \ + include/rocksdb/coro_db.h \ + include/rocksdb/utilities/coro_stackable_db.h # Some headers like in port/ are platform-specific - DEV_HEADERS_TO_CHECK := $(shell $(FIND) $(DEV_HEADER_DIRS) -type f -name '*.h' | grep -E -v 'port/|plugin/|range_tree/|secondary_index/') - PUBLIC_HEADERS_TO_CHECK := $(shell $(FIND) include/ -type f -name '*.h') + DEV_HEADERS_TO_CHECK := $(filter-out $(FOLLY_DEPENDENT_HEADERS), $(shell $(FIND) $(DEV_HEADER_DIRS) -type f -name '*.h' | grep -E -v 'port/|plugin/|range_tree/|secondary_index/')) + PUBLIC_HEADERS_TO_CHECK := $(filter-out $(FOLLY_DEPENDENT_HEADERS), $(shell $(FIND) include/ -type f -name '*.h')) else DEV_HEADERS_TO_CHECK := PUBLIC_HEADERS_TO_CHECK := diff --git a/db/db_basic_test.cc b/db/db_basic_test.cc index 52a6ac9891..b882895931 100644 --- a/db/db_basic_test.cc +++ b/db/db_basic_test.cc @@ -4020,7 +4020,8 @@ TEST_P(DBMultiGetAsyncIOTest, GetFromL0) { AssertMultiGetIOBatchSize(3, 1); } #else // ROCKSDB_IOURING_PRESENT - ASSERT_EQ(statistics()->getTickerCount(MULTIGET_COROUTINE_COUNT), 0); + ASSERT_EQ(statistics()->getTickerCount(MULTIGET_COROUTINE_COUNT), + use_coroutine_ ? 3 : 0); #endif // ROCKSDB_IOURING_PRESENT } @@ -4047,7 +4048,8 @@ TEST_P(DBMultiGetAsyncIOTest, GetFromL1) { ASSERT_EQ(statistics()->getTickerCount(MULTIGET_COROUTINE_COUNT), UseCoroutineRead() ? 6 : 3); #else // ROCKSDB_IOURING_PRESENT - ASSERT_EQ(statistics()->getTickerCount(MULTIGET_COROUTINE_COUNT), 0); + ASSERT_EQ(statistics()->getTickerCount(MULTIGET_COROUTINE_COUNT), + use_coroutine_ ? 6 : 0); #endif // ROCKSDB_IOURING_PRESENT } @@ -4200,7 +4202,8 @@ TEST_P(DBMultiGetAsyncIOTest, GetFromL2WithRangeOverlapL0L1) { ASSERT_EQ(statistics()->getTickerCount(MULTIGET_COROUTINE_COUNT), UseCoroutineRead() ? 3 : 2); #else // ROCKSDB_IOURING_PRESENT - ASSERT_EQ(statistics()->getTickerCount(MULTIGET_COROUTINE_COUNT), 0); + ASSERT_EQ(statistics()->getTickerCount(MULTIGET_COROUTINE_COUNT), + use_coroutine_ ? 3 : 0); #endif // ROCKSDB_IOURING_PRESENT } @@ -4263,7 +4266,7 @@ TEST_P(DBMultiGetAsyncIOTest, GetNoIOUring) { ASSERT_EQ(async_read_bytes.count, 0); const uint64_t coroutine_count = statistics()->getTickerCount(MULTIGET_COROUTINE_COUNT); - if (UseCoroutineRead()) { + if (use_coroutine_) { ASSERT_GT(coroutine_count, 0); } else { ASSERT_EQ(coroutine_count, 0); diff --git a/db/db_impl/compacted_db_impl_sync_and_async.h b/db/db_impl/compacted_db_impl_sync_and_async.h index 3d41505718..2fe3c27bcf 100644 --- a/db/db_impl/compacted_db_impl_sync_and_async.h +++ b/db/db_impl/compacted_db_impl_sync_and_async.h @@ -6,6 +6,10 @@ #include "util/coro_utils.h" +#if defined(USE_COROUTINES) && defined(WITH_COROUTINES) +#include "util/coro_stats_util.h" +#endif // USE_COROUTINES && WITH_COROUTINES + #if defined(WITHOUT_COROUTINES) || \ (defined(USE_COROUTINES) && defined(WITH_COROUTINES)) @@ -14,6 +18,10 @@ namespace ROCKSDB_NAMESPACE { DEFINE_SYNC_AND_ASYNC(Status, CompactedDBImpl::Get) (const ReadOptions& _read_options, ColumnFamilyHandle*, const Slice& key, PinnableSlice* value, std::string* timestamp) { +#ifdef WITH_COROUTINES + INSTALL_COROUTINE_STATS_CONTEXT_SCOPE( + immutable_db_options_.fs->GetReadExecutor(), immutable_db_options_.env); +#endif if (_read_options.io_activity != Env::IOActivity::kUnknown && _read_options.io_activity != Env::IOActivity::kGet) { CO_RETURN Status::InvalidArgument( @@ -102,6 +110,10 @@ DEFINE_SYNC_AND_ASYNC(void, CompactedDBImpl::MultiGet) ColumnFamilyHandle** /*column_families*/, const Slice* keys, PinnableSlice* values, std::string* timestamps, Status* statuses, const bool /*sorted_input*/) { +#ifdef WITH_COROUTINES + INSTALL_COROUTINE_STATS_CONTEXT_SCOPE( + immutable_db_options_.fs->GetReadExecutor(), immutable_db_options_.env); +#endif assert(user_comparator_); Status s; if (_read_options.io_activity != Env::IOActivity::kUnknown && diff --git a/db/db_impl/db_impl.cc b/db/db_impl/db_impl.cc index dfc72e7c0e..989e0b6e6b 100644 --- a/db/db_impl/db_impl.cc +++ b/db/db_impl/db_impl.cc @@ -15,13 +15,6 @@ #include #endif -#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 - #include #include #include @@ -137,6 +130,15 @@ #include "util/udt_util.h" #include "utilities/trace/replayer_impl.h" +#if USE_COROUTINES +#include "folly/Executor.h" +#include "folly/coro/Nothrow.h" +#include "folly/coro/Task.h" +#include "folly/executors/IOExecutor.h" +#include "folly/io/async/EventBase.h" +#include "rocksdb/coro_db.h" +#endif // USE_COROUTINES + namespace ROCKSDB_NAMESPACE { const std::string kDefaultColumnFamilyName("default"); @@ -7772,9 +7774,24 @@ const IOStatsContext* CurrentIOStatsContextForAsyncCallback( } #if USE_COROUTINES +void InstallCoroutineStatsConfigToTLS( + const CoroutineStatsConfig& stats_config) { +#ifndef NPERF_CONTEXT + if (stats_config.per_level_perf_context_enabled) { + get_perf_context()->EnablePerLevelPerfContext(); + } else { + get_perf_context()->per_level_perf_context_enabled = false; + } +#endif +#ifndef NIOSTATS_CONTEXT + get_iostats_context()->disable_iostats = stats_config.iostats_disabled; +#endif + SetPerfLevel(stats_config.perf_level); +} + std::optional CaptureCoroutineStatsConfigForCallback( bool stats_enabled) { - if (!stats_enabled) { + if (!stats_enabled || !IsCoroutineStatsEnabled()) { return std::nullopt; } return CaptureCoroutineStatsConfig(); @@ -7784,49 +7801,49 @@ std::optional CaptureCoroutineStatsConfigForCallback( } // namespace -void DBImpl::GetAsync(const ReadOptions& options, - ColumnFamilyHandle* column_family, const Slice& key, - PinnableSlice* value, std::string* timestamp, - Status& status, AsyncCallback& callback) { +void DB::GetAsync(const ReadOptions& options, ColumnFamilyHandle* column_family, + const Slice& key, PinnableSlice* value, + std::string* timestamp, Status& status, + AsyncCallback& callback) { const bool stats_enabled = callback.EnableStats(); #if USE_COROUTINES - auto* read_executor = immutable_db_options_.fs->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, 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; + CoroDB* coro_db = GetCoroDB(); + if (coro_db != nullptr) { + auto* read_executor = GetFileSystem()->GetReadExecutor(); + if (read_executor != nullptr) { + auto* read_event_base = read_executor->getEventBase(); + assert(read_event_base != nullptr); + auto stats_config = CaptureCoroutineStatsConfigForCallback(stats_enabled); + auto task = + [](std::optional task_stats_config, + CoroDB* task_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 { + if (task_stats_config.has_value()) { + InstallCoroutineStatsConfigToTLS(*task_stats_config); + } + task_status = co_await folly::coro::co_nothrow( + task_db->GetCoroutine(task_options, task_column_family, task_key, + task_value, task_timestamp)); + + const bool task_stats_enabled = task_stats_config.has_value(); + task_callback.OnComplete( + CurrentPerfContextForAsyncCallback(task_stats_enabled), + CurrentIOStatsContextForAsyncCallback(task_stats_enabled)); + }(std::move(stats_config), coro_db, 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 + if (stats_enabled) { - // Mirror behavior of coroutine version where each coroutine gets a clean - // context + // Match coroutine reads: callback stats describe only this request. ResetThreadLocalStatsForAsyncCallback(); } status = Get(options, column_family, key, value, timestamp); @@ -7834,51 +7851,52 @@ void DBImpl::GetAsync(const ReadOptions& options, 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) { +void DB::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_.fs->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; + CoroDB* coro_db = GetCoroDB(); + if (coro_db != nullptr) { + auto* read_executor = GetFileSystem()->GetReadExecutor(); + if (read_executor != nullptr) { + auto* read_event_base = read_executor->getEventBase(); + assert(read_event_base != nullptr); + auto stats_config = CaptureCoroutineStatsConfigForCallback(stats_enabled); + auto task = + [](std::optional task_stats_config, + CoroDB* task_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 { + if (task_stats_config.has_value()) { + InstallCoroutineStatsConfigToTLS(*task_stats_config); + } + co_await folly::coro::co_nothrow(task_db->MultiGetCoroutine( + task_options, task_num_keys, task_column_families, task_keys, + task_values, task_timestamps, task_statuses, task_sorted_input)); + + const bool task_stats_enabled = task_stats_config.has_value(); + task_callback.OnComplete( + CurrentPerfContextForAsyncCallback(task_stats_enabled), + CurrentIOStatsContextForAsyncCallback(task_stats_enabled)); + }(std::move(stats_config), coro_db, 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 + // Match coroutine reads: callback stats describe only this request. ResetThreadLocalStatsForAsyncCallback(); } MultiGet(options, num_keys, column_families, keys, values, timestamps, diff --git a/db/db_impl/db_impl.h b/db/db_impl/db_impl.h index 0f6c37f4d7..a9194039c7 100644 --- a/db/db_impl/db_impl.h +++ b/db/db_impl/db_impl.h @@ -75,6 +75,10 @@ #include "util/stop_watch.h" #include "util/thread_local.h" +#if USE_COROUTINES +#include "rocksdb/coro_db.h" +#endif + namespace ROCKSDB_NAMESPACE { class Arena; @@ -200,14 +204,22 @@ struct DBOpenLogRecordReadReporter : public log::Reader::Reporter { // // Since it's a very large class, the definition of the functions is // divided in several db_impl_*.cc files, besides db_impl.cc. -class DBImpl : public DB { +#if USE_COROUTINES +class DBImpl : public DB, + public CoroDB +#else +class DBImpl : public DB +#endif +{ public: DBImpl(const DBOptions& options, const std::string& dbname, const bool seq_per_batch = false, const bool batch_per_txn = true, bool read_only = false); - // No copying allowed + // No copying or moving allowed DBImpl(const DBImpl&) = delete; void operator=(const DBImpl&) = delete; + DBImpl(DBImpl&&) = delete; + void operator=(DBImpl&&) = delete; virtual ~DBImpl(); @@ -273,14 +285,10 @@ class DBImpl : public DB { bool HasAnyBlobDirectWriteColumnFamily(); using DB::Get; - using DB::GetAsync; - DECLARE_SYNC_ASYNC_AND_CALLBACK( - Status, Get, GetAsync, - (const ReadOptions& options, ColumnFamilyHandle* column_family, - const Slice& key, PinnableSlice* value, std::string* timestamp, - Status& status, AsyncCallback& callback), - const ReadOptions& _read_options, ColumnFamilyHandle* column_family, - const Slice& key, PinnableSlice* value, std::string* timestamp); + DECLARE_SYNC_AND_ASYNC_OVERRIDE(Status, Get, const ReadOptions& _read_options, + ColumnFamilyHandle* column_family, + const Slice& key, PinnableSlice* value, + std::string* timestamp); using DB::GetEntity; Status GetEntity(const ReadOptions& options, @@ -309,7 +317,6 @@ class DBImpl : public DB { } using DB::MultiGet; - using DB::MultiGetAsync; // This MultiGet is a batched version, which may be faster than calling Get // multiple times, especially if the keys have some spatial locality that // enables them to be queried in the same SST files/set of files. The larger @@ -317,16 +324,13 @@ class DBImpl : public DB { // The values and statuses parameters are arrays with number of elements // equal to keys.size(). This allows the storage for those to be alloacted // by the caller on the stack for small batches - DECLARE_SYNC_ASYNC_AND_CALLBACK( - void, MultiGet, 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 ReadOptions& _read_options, const size_t num_keys, - ColumnFamilyHandle** column_families, const Slice* keys, - PinnableSlice* values, std::string* timestamps, Status* statuses, - const bool sorted_input = false); + DECLARE_SYNC_AND_ASYNC_OVERRIDE(void, MultiGet, + const ReadOptions& _read_options, + const size_t num_keys, + ColumnFamilyHandle** column_families, + const Slice* keys, PinnableSlice* values, + std::string* timestamps, Status* statuses, + const bool sorted_input = false); void MultiGetWithCallback( const ReadOptions& _read_options, ColumnFamilyHandle* column_family, @@ -1498,6 +1502,10 @@ class DBImpl : public DB { bool seq_per_batch() const { return seq_per_batch_; } +#if USE_COROUTINES + CoroDB* GetCoroDB() override { return this; } +#endif + protected: const std::string dbname_; const bool read_only_; diff --git a/db/db_impl/db_impl_sync_and_async.h b/db/db_impl/db_impl_sync_and_async.h index bb2e1887b4..ed1e6c7436 100644 --- a/db/db_impl/db_impl_sync_and_async.h +++ b/db/db_impl/db_impl_sync_and_async.h @@ -7,8 +7,7 @@ #include "util/coro_utils.h" #if defined(USE_COROUTINES) && defined(WITH_COROUTINES) -#include "folly/executors/IOExecutor.h" -#include "folly/io/async/EventBase.h" +#include "util/coro_stats_util.h" #endif // USE_COROUTINES && WITH_COROUTINES #if defined(WITHOUT_COROUTINES) || \ @@ -26,6 +25,10 @@ DEFINE_SYNC_AND_ASYNC(Status, DBImpl::GetImpl) DEFINE_SYNC_AND_ASYNC(Status, DBImpl::Get) (const ReadOptions& _read_options, ColumnFamilyHandle* column_family, const Slice& key, PinnableSlice* value, std::string* timestamp) { +#ifdef WITH_COROUTINES + INSTALL_COROUTINE_STATS_CONTEXT_SCOPE( + immutable_db_options_.fs->GetReadExecutor(), immutable_db_options_.env); +#endif assert(value != nullptr); value->Reset(); @@ -41,18 +44,8 @@ DEFINE_SYNC_AND_ASYNC(Status, DBImpl::Get) read_options.io_activity = Env::IOActivity::kGet; } -#ifdef WITH_COROUTINES - auto* read_executor = immutable_db_options_.fs->GetReadExecutor(); - if (read_executor != nullptr) { - auto* read_event_base = read_executor->getEventBase(); - assert(read_event_base != nullptr); - 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); + CO_RETURN CO_AWAIT(GetImpl, read_options, column_family, key, value, + timestamp); } DEFINE_SYNC_AND_ASYNC(Status, DBImpl::GetImpl) @@ -808,6 +801,10 @@ DEFINE_SYNC_AND_ASYNC(void, DBImpl::MultiGet) (const ReadOptions& _read_options, const size_t num_keys, ColumnFamilyHandle** column_families, const Slice* keys, PinnableSlice* values, std::string* timestamps, Status* statuses, const bool sorted_input) { +#ifdef WITH_COROUTINES + INSTALL_COROUTINE_STATS_CONTEXT_SCOPE( + immutable_db_options_.fs->GetReadExecutor(), immutable_db_options_.env); +#endif if (_read_options.io_activity != Env::IOActivity::kUnknown && _read_options.io_activity != Env::IOActivity::kMultiGet) { Status s = Status::InvalidArgument( @@ -824,21 +821,10 @@ DEFINE_SYNC_AND_ASYNC(void, DBImpl::MultiGet) if (read_options.io_activity == Env::IOActivity::kUnknown) { read_options.io_activity = Env::IOActivity::kMultiGet; } -#ifdef WITH_COROUTINES - auto* read_executor = immutable_db_options_.fs->GetReadExecutor(); - if (read_executor != nullptr) { - auto* read_event_base = read_executor->getEventBase(); - assert(read_event_base != nullptr); - co_await folly::coro::co_nothrow(folly::coro::co_withExecutor( - folly::Executor::getKeepAliveToken(read_event_base), - MultiGetCommonCoroutine( - read_options, num_keys, column_families, keys, values, - /* columns */ nullptr, timestamps, statuses, sorted_input))); - co_return; - } -#endif - MultiGetCommon(read_options, num_keys, column_families, keys, values, - /* columns */ nullptr, timestamps, statuses, sorted_input); + CO_AWAIT(MultiGetCommon, read_options, num_keys, column_families, keys, + values, + /* columns */ nullptr, timestamps, statuses, sorted_input); + CO_RETURN; } } // namespace ROCKSDB_NAMESPACE diff --git a/db/db_test_util.cc b/db/db_test_util.cc index eb013249a6..696a032308 100644 --- a/db/db_test_util.cc +++ b/db/db_test_util.cc @@ -24,7 +24,14 @@ #include "util/random.h" #if USE_COROUTINES +#include + +#include "folly/Executor.h" #include "folly/coro/BlockingWait.h" +#include "folly/coro/Task.h" +#include "folly/executors/IOExecutor.h" +#include "folly/io/async/EventBase.h" +#include "rocksdb/file_system.h" #endif // USE_COROUTINES namespace ROCKSDB_NAMESPACE { @@ -36,6 +43,21 @@ int64_t MaybeCurrentTime(Env* env) { return time; } +#if USE_COROUTINES +template +T BlockingWaitOnReadExecutor(DB* db, folly::coro::Task task) { + auto* read_executor = db->GetFileSystem()->GetReadExecutor(); + if (read_executor == nullptr) { + return folly::coro::blockingWait(std::move(task)); + } + + auto* read_event_base = read_executor->getEventBase(); + assert(read_event_base != nullptr); + return folly::coro::blockingWait(folly::coro::co_withExecutor( + folly::Executor::getKeepAliveToken(read_event_base), std::move(task))); +} +#endif // USE_COROUTINES + } // anonymous namespace // Special Env used to delay background operations @@ -872,9 +894,11 @@ std::string DBTestBase::Get(const std::string& k, const Snapshot* snapshot, #if USE_COROUTINES if (use_coroutine) { PinnableSlice pinnable_value(&result); - s = folly::coro::blockingWait(dbfull()->GetCoroutine( - options, dbfull()->DefaultColumnFamily(), k, &pinnable_value, - /*timestamp=*/nullptr)); + s = BlockingWaitOnReadExecutor( + dbfull(), + dbfull()->GetCoroutine(options, dbfull()->DefaultColumnFamily(), k, + &pinnable_value, + /*timestamp=*/nullptr)); if (s.ok() && pinnable_value.IsPinned()) { result.assign(pinnable_value.data(), pinnable_value.size()); } @@ -903,8 +927,10 @@ std::string DBTestBase::Get(int cf, const std::string& k, #if USE_COROUTINES if (use_coroutine) { PinnableSlice pinnable_value(&result); - s = folly::coro::blockingWait(dbfull()->GetCoroutine( - options, handles_[cf], k, &pinnable_value, /*timestamp=*/nullptr)); + s = BlockingWaitOnReadExecutor( + dbfull(), + dbfull()->GetCoroutine(options, handles_[cf], k, &pinnable_value, + /*timestamp=*/nullptr)); if (s.ok() && pinnable_value.IsPinned()) { result.assign(pinnable_value.data(), pinnable_value.size()); } @@ -956,9 +982,11 @@ std::vector DBTestBase::MultiGet(std::vector cfs, s.resize(cfs.size()); #if USE_COROUTINES if (use_coroutine) { - folly::coro::blockingWait(dbfull()->MultiGetCoroutine( - options, cfs.size(), handles.data(), keys.data(), pin_values.data(), - /*timestamps=*/nullptr, s.data(), /*sorted_input=*/false)); + BlockingWaitOnReadExecutor( + dbfull(), dbfull()->MultiGetCoroutine( + options, cfs.size(), handles.data(), keys.data(), + pin_values.data(), /*timestamps=*/nullptr, s.data(), + /*sorted_input=*/false)); } else #else (void)use_coroutine; @@ -1005,9 +1033,12 @@ std::vector DBTestBase::MultiGet(const std::vector& k, if (use_coroutine) { std::vector cfs(keys.size(), dbfull()->DefaultColumnFamily()); - folly::coro::blockingWait(dbfull()->MultiGetCoroutine( - options, keys.size(), cfs.data(), keys.data(), pin_values.data(), - /*timestamps=*/nullptr, statuses.data(), /*sorted_input=*/false)); + BlockingWaitOnReadExecutor( + dbfull(), + dbfull()->MultiGetCoroutine(options, keys.size(), cfs.data(), + keys.data(), pin_values.data(), + /*timestamps=*/nullptr, statuses.data(), + /*sorted_input=*/false)); } else #else (void)use_coroutine; @@ -1037,7 +1068,8 @@ Status DBTestBase::Get(const std::string& k, PinnableSlice* v, options.verify_checksums = true; #if USE_COROUTINES if (use_coroutine) { - return folly::coro::blockingWait( + return BlockingWaitOnReadExecutor( + dbfull(), dbfull()->GetCoroutine(options, dbfull()->DefaultColumnFamily(), k, v, /*timestamp=*/nullptr)); } diff --git a/db/perf_context_test.cc b/db/perf_context_test.cc index 7e6844aecc..a0e94918a5 100644 --- a/db/perf_context_test.cc +++ b/db/perf_context_test.cc @@ -112,44 +112,45 @@ TEST_F(PerfContextTest, CoroutineStatsContextScopeCollectsStats) { 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); + 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; } - 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); @@ -199,32 +200,33 @@ TEST_F(PerfContextTest, CoroutineStatsContextsRemainIsolatedWhenInterleaved) { 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()); + { + 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); + 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; + 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; + 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; }; @@ -250,17 +252,18 @@ TEST_F(PerfContextTest, CoroutineStatsContextScopeCollectsGetCpuNanos) { 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); + 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); } - clock->MockSleepForCPUNanos(7); - scope.Finalize(); EXPECT_EQ(29, get_perf_context()->get_cpu_nanos); } } diff --git a/include/rocksdb/coro_db.h b/include/rocksdb/coro_db.h new file mode 100644 index 0000000000..daf4a0a6e5 --- /dev/null +++ b/include/rocksdb/coro_db.h @@ -0,0 +1,78 @@ +// 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 + +#include +#include + +#include "folly/coro/Task.h" +#include "rocksdb/db.h" + +namespace ROCKSDB_NAMESPACE { + +// EXPERIMENTAL native coroutine read interface. +// +// Using this interface requires Folly to be available and RocksDB to be built +// with USE_COROUTINES=1. +// +// These are the coroutine counterparts to DB::GetAsync() and +// DB::MultiGetAsync(). See those APIs for the shared read, input/output +// lifetime, and concurrency contract. Task completion takes the place of +// OnComplete for those requirements. +// +// The returned tasks are lazy: no read begins until a task is awaited or +// started. +// +// IMPORTANT: RocksDB assumes each request runs entirely on one thread. Schedule +// coroutine reads on the read IOExecutor's EventBase and do not migrate them. +// +// STATS: +// Unlike the callback APIs, coroutine stats are returned through TLS. Each task +// publishes its request-local PerfContext and IOStatsContext to TLS on its +// execution thread before completing. Read them after the await resumes on the +// same EventBase. The TLS values are defined only on the thread where the +// RocksDB operation completes. Consume or copy them immediately after the +// RocksDB await returns; no coroutine suspension may occur in between: +// +// auto* coro_db = db->GetCoroDB(); +// auto* read_executor = db->GetFileSystem()->GetReadExecutor(); +// auto* read_event_base = +// read_executor == nullptr ? nullptr : read_executor->getEventBase(); +// if (coro_db == nullptr || read_event_base == nullptr) { +// co_return db->Get(options, column_family, key, value, timestamp); +// } +// +// auto read = [&]() -> folly::coro::Task { +// Status status = co_await coro_db->GetCoroutine( +// options, column_family, key, value, timestamp); +// const PerfContext* perf_context = get_perf_context(); +// const IOStatsContext* iostats_context = get_iostats_context(); +// // Consume or copy the stats here, before any further co_await. +// co_return status; +// }; +// co_return co_await folly::coro::co_withExecutor( +// folly::Executor::getKeepAliveToken(read_event_base), read()); +class CoroDB { + public: + virtual ~CoroDB() = default; + + // Reads key and completes after populating value and the optional timestamp. + // The returned Status describes the operation and its outputs. + virtual folly::coro::Task GetCoroutine( + const ReadOptions& options, ColumnFamilyHandle* column_family, + const Slice& key, PinnableSlice* value, std::string* timestamp) = 0; + + // Reads num_keys keys and completes after populating all per-key statuses, + // values, and optional timestamps. There is no aggregate Status. + virtual folly::coro::Task MultiGetCoroutine( + const ReadOptions& options, size_t num_keys, + ColumnFamilyHandle** column_families, const Slice* keys, + PinnableSlice* values, std::string* timestamps, Status* statuses, + bool sorted_input) = 0; +}; + +} // namespace ROCKSDB_NAMESPACE diff --git a/include/rocksdb/db.h b/include/rocksdb/db.h index 955c8c9bd7..5ed8dd41d7 100644 --- a/include/rocksdb/db.h +++ b/include/rocksdb/db.h @@ -52,6 +52,7 @@ struct TableProperties; struct WriteOptions; struct WaitForCompactOptions; class Env; +class CoroDB; class EventListener; class FileSystem; class LazyWideColumns; @@ -1164,18 +1165,21 @@ class DB { // EXPERIMENTAL // - // RocksDB async read variants of Get and MultiGet. Implementations may - // complete asynchronously. The default implementation runs synchronously and - // invokes the callback inline before returning. + // RocksDB async read variants of Get and MultiGet. The default implementation + // uses the native coroutine capability when exposed and a FileSystem read + // executor is available, and otherwise invokes the callback synchronously + // before returning. // - // Full async execution requires RocksDB to be built with coroutine support, - // a configured FileSystem read executor, and an underlying FileSystem that - // implements FSRandomAccessFile::SubmitReadAsync(). Without those - // requirements, implementations delegate to the synchronous Get/MultiGet path - // and invoke the callback inline before returning. When full async execution - // is available, RocksDB can run an internal read coroutine on the read - // executor, suspend while async file IO is outstanding, and resume when the - // filesystem signals completion. The callback is invoked after the requested + // Full async file IO requires RocksDB to be built with coroutine support, a + // DB implementation that exposes the native coroutine capability, a + // configured FileSystem read executor, and an underlying FileSystem that + // implements FSRandomAccessFile::SubmitReadAsync(). When the native + // coroutine capability and a read executor are available, RocksDB can run an + // internal read coroutine on the read executor, suspend while async file IO + // is outstanding, and resume when the filesystem signals completion. Without + // a native coroutine capability or read executor, the default implementation + // delegates to the synchronous Get/MultiGet path and invokes the callback + // inline before returning. The callback is invoked after the requested // statuses and output buffers have been populated. Applications can set // `DBOptions::read_io_executor_threads` before opening the DB to configure // executor parallelism for their workload. @@ -1209,11 +1213,7 @@ class DB { virtual void GetAsync(const ReadOptions& options, ColumnFamilyHandle* column_family, const Slice& key, PinnableSlice* value, std::string* timestamp, - Status& status, AsyncCallback& callback) { - status = Get(options, column_family, key, value, timestamp); - callback.OnComplete(/*callback_perf_context=*/nullptr, - /*callback_iostats_context=*/nullptr); - } + Status& status, AsyncCallback& callback); virtual void GetAsync(const ReadOptions& options, ColumnFamilyHandle* column_family, const Slice& key, @@ -1287,12 +1287,7 @@ class DB { ColumnFamilyHandle** column_families, const Slice* keys, PinnableSlice* values, std::string* timestamps, Status* statuses, - const bool sorted_input, AsyncCallback& callback) { - MultiGet(options, num_keys, column_families, keys, values, timestamps, - statuses, sorted_input); - callback.OnComplete(/*callback_perf_context=*/nullptr, - /*callback_iostats_context=*/nullptr); - } + const bool sorted_input, AsyncCallback& callback); virtual void MultiGetAsync(const ReadOptions& options, const size_t num_keys, ColumnFamilyHandle** column_families, @@ -2607,6 +2602,10 @@ class DB { virtual Status TryCatchUpWithPrimary() { return Status::NotSupported("Supported only by secondary instance"); } + + // Returns the non-owning native coroutine interface, or nullptr when this DB + // does not support it. The returned pointer must not outlive this DB. + virtual CoroDB* GetCoroDB() { return nullptr; } }; struct WriteStallStatsMapKeys { diff --git a/include/rocksdb/file_system.h b/include/rocksdb/file_system.h index f7222cf53c..4b4c4ea9a7 100644 --- a/include/rocksdb/file_system.h +++ b/include/rocksdb/file_system.h @@ -463,7 +463,8 @@ class FileSystem : public Customizable { // // Used in the RocksDB coroutine query interface. Coroutine-based read // requests will run on an EventBase from this executor. Returns nullptr when - // coroutine read execution is not supported. + // coroutine read execution is not supported. When called from one of its + // EventBases, the executor's getEventBase() must return that EventBase. virtual folly::IOExecutor* GetReadExecutor() { return nullptr; } // Increase the maximum number of threads used by the FileSystem-owned read diff --git a/tools/db_bench_tool.cc b/tools/db_bench_tool.cc index ff7bb84a55..b29f364036 100644 --- a/tools/db_bench_tool.cc +++ b/tools/db_bench_tool.cc @@ -111,13 +111,12 @@ #endif #if USE_COROUTINES -#include "folly/coro/Baton.h" -#include "folly/coro/CurrentExecutor.h" +#include "folly/coro/Nothrow.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" +#include "rocksdb/coro_db.h" #endif // USE_COROUTINES using GFLAGS_NAMESPACE::ParseCommandLineFlags; @@ -204,7 +203,8 @@ DEFINE_string( "\treadtocache -- 1 thread reading database sequentially\n" "\treadreverse -- read N times in reverse order\n" "\treadrandom -- read N times in random order\n" - "\treadrandomcoroutine -- read N times in random order using async Get\n" + "\treadrandomcoroutine -- read N times in random order using " + "GetCoroutine\n" "\treadmissing -- read N missing keys in random order\n" "\treadwhilewriting -- 1 writer, N threads doing random " "reads\n" @@ -4254,8 +4254,8 @@ class Benchmark { } else if (name == "readrandomcoroutine") { fprintf(stderr, "read executor pool size (threads) = %d\n", FLAGS_threads > 0 ? FLAGS_threads : 1); - // A single db_bench driver thread issues concurrent async Get calls - // onto the configured read executor. + // A single db_bench driver thread starts concurrent GetCoroutine tasks + // on the configured read executor. num_threads = 1; method = &Benchmark::ReadRandomCoroutine; } else if (name == "readrandomfast") { @@ -4269,8 +4269,8 @@ class Benchmark { entries_per_batch_); fprintf(stderr, "read executor pool size (threads) = %d\n", FLAGS_threads > 0 ? FLAGS_threads : 1); - // A single db_bench driver thread issues concurrent async MultiGet - // calls onto the configured read executor. + // A single db_bench driver thread starts concurrent MultiGetCoroutine + // tasks on the configured read executor. num_threads = 1; method = &Benchmark::MultiReadRandomCoroutine; } else if (name == "readrandomentity") { @@ -7592,10 +7592,8 @@ class Benchmark { #if USE_COROUTINES folly::IOExecutor* GetReadExecutorForCoroutineBenchmark( - const char* benchmark_name, int pool_threads) { - Env* env = - open_options_.env != nullptr ? open_options_.env : Env::Default(); - auto fs = env->GetFileSystem(); + DB* db, const char* benchmark_name, int pool_threads) { + auto* fs = db->GetFileSystem(); if (fs != nullptr) { fs->SetReadIOExecutorThreads(pool_threads); } @@ -7610,29 +7608,6 @@ class Benchmark { return read_executor; } - struct AsyncAwaitCallback : public DB::AsyncCallback { - 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) { @@ -7643,34 +7618,14 @@ class Benchmark { #endif } - folly::coro::Task AwaitGetAsync(DB* db, const ReadOptions& options, - ColumnFamilyHandle* cfh, - const Slice& key, - PinnableSlice* value, - 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, 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 MergeCoroutineJobPerfContext(PerfContext* job_perf_context) { +#ifndef NPERF_CONTEXT + if (job_perf_context != nullptr) { + job_perf_context->Merge(*get_perf_context()); + } +#else + (void)job_perf_context; +#endif } void SetCoroutineJobPerfContextMessage( @@ -7685,7 +7640,8 @@ class Benchmark { } folly::coro::Task ReadRandomCoroutineJob( - uint64_t seed, int64_t ops, std::atomic* total_ops, + const std::vector& db_choices, uint64_t seed, + int64_t ops, std::atomic* total_ops, std::atomic* total_found, std::atomic* total_bytes, int job_id, std::vector* job_perf_contexts) { Random64 rng(seed); @@ -7703,8 +7659,13 @@ class Benchmark { int64_t job_found = 0; int64_t job_bytes = 0; PerfContext job_perf_context; + PerfContext* const job_perf_context_ptr = + job_perf_contexts != nullptr ? &job_perf_context : nullptr; + PrepareCoroutineJobPerfContext(job_perf_context_ptr); for (int64_t done = 0; done < ops; ++done) { - DBWithColumnFamilies* db_with_cfh = SelectDBWithCfh(rng.Next()); + const uint64_t db_rand = rng.Next(); + DBWithColumnFamilies* db_with_cfh = + db_choices[db_rand % db_choices.size()]; const int64_t key_rand = GetRandomKey(&rng); GenerateKeyFromInt(key_rand, FLAGS_num, &key); @@ -7724,15 +7685,22 @@ class Benchmark { cfh = db_with_cfh->db->DefaultColumnFamily(); } - Status s = co_await AwaitGetAsync( - db_with_cfh->db, options, cfh, key, &pinnable_val, ts_ptr, - job_perf_contexts != nullptr ? &job_perf_context : nullptr); + CoroDB* coro_db = db_with_cfh->db->GetCoroDB(); + if (coro_db == nullptr) { + fprintf(stderr, + "readrandomcoroutine requires a DB with native coroutine " + "reads\n"); + abort(); + } + Status s = co_await folly::coro::co_nothrow( + coro_db->GetCoroutine(options, cfh, key, &pinnable_val, ts_ptr)); + MergeCoroutineJobPerfContext(job_perf_context_ptr); ++job_ops; if (s.ok()) { ++job_found; job_bytes += key.size() + pinnable_val.size() + user_timestamp_size_; } else if (!s.IsNotFound()) { - fprintf(stderr, "GetAsync returned an error: %s\n", + fprintf(stderr, "GetCoroutine returned an error: %s\n", s.ToString().c_str()); abort(); } @@ -7747,11 +7715,11 @@ class Benchmark { } // One coroutine "job": mirrors a single multireadrandom thread -- a loop of - // batched async MultiGet reads -- but as a coroutine, so it yields its + // batched MultiGetCoroutine reads -- but as a coroutine, so it yields its // executor thread while its IO is in flight. Per-job buffers and RNG mean // jobs share no mutable state; results are folded into the shared counters. - folly::coro::Task MultiGetAsyncCoroutineJob( - DB* db, ColumnFamilyHandle* cfh, uint64_t seed, int64_t ops, + folly::coro::Task MultiGetCoroutineJob( + CoroDB* coro_db, ColumnFamilyHandle* cfh, uint64_t seed, int64_t ops, std::atomic* total_ops, std::atomic* total_found, std::atomic* total_bytes, int job_id, std::vector* job_perf_contexts) { @@ -7771,23 +7739,26 @@ class Benchmark { int64_t job_found = 0; int64_t job_bytes = 0; PerfContext job_perf_context; + PerfContext* const job_perf_context_ptr = + job_perf_contexts != nullptr ? &job_perf_context : nullptr; + PrepareCoroutineJobPerfContext(job_perf_context_ptr); 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(), - job_perf_contexts != nullptr ? &job_perf_context : nullptr); + co_await folly::coro::co_nothrow(coro_db->MultiGetCoroutine( + options, entries_per_batch_, cfs.data(), keys.data(), values.get(), + /*timestamps=*/nullptr, statuses.data(), /*sorted_input=*/false)); + MergeCoroutineJobPerfContext(job_perf_context_ptr); job_ops += entries_per_batch_; for (int64_t i = 0; i < entries_per_batch_; ++i) { if (statuses[i].ok()) { job_bytes += keys[i].size() + values[i].size(); ++job_found; } else if (!statuses[i].IsNotFound()) { - fprintf(stderr, "MultiGetAsync returned an error: %s\n", + fprintf(stderr, "MultiGetCoroutine returned an error: %s\n", statuses[i].ToString().c_str()); abort(); } @@ -7804,7 +7775,7 @@ class Benchmark { #endif // USE_COROUTINES // Coroutine variant of ReadRandom. A single db_bench driver thread launches - // (-coro_jobs_per_thread * -threads) async Get jobs onto the FileSystem + // (-coro_jobs_per_thread * -threads) GetCoroutine jobs onto the FileSystem // read executor and waits for them all. void ReadRandomCoroutine(ThreadState* thread) { #if USE_COROUTINES @@ -7812,9 +7783,16 @@ class Benchmark { const int jobs_per_thread = FLAGS_coro_jobs_per_thread > 0 ? FLAGS_coro_jobs_per_thread : 1; const int num_jobs = jobs_per_thread * pool_threads; - auto* read_executor = GetReadExecutorForCoroutineBenchmark( - "readrandomcoroutine", pool_threads); DB* db = SelectDB(thread); + const size_t db_count = db_.db != nullptr ? 1 : multi_dbs_.size(); + std::vector db_choices; + db_choices.reserve(db_count); + for (size_t i = 0; i < db_count; ++i) { + db_choices.push_back(SelectDBWithCfh(i)); + } + assert(!db_choices.empty()); + auto* read_executor = GetReadExecutorForCoroutineBenchmark( + db, "readrandomcoroutine", pool_threads); std::atomic total_ops{0}; std::atomic total_found{0}; @@ -7844,8 +7822,8 @@ class Benchmark { futures.push_back( folly::coro::co_withExecutor( folly::Executor::getKeepAliveToken(read_executor->getEventBase()), - ReadRandomCoroutineJob(thread->rand.Next(), ops, &total_ops, - &total_found, &total_bytes, j, + ReadRandomCoroutineJob(db_choices, thread->rand.Next(), ops, + &total_ops, &total_found, &total_bytes, j, job_perf_contexts_ptr)) .start()); } @@ -7872,7 +7850,7 @@ class Benchmark { } // Coroutine variant of MultiReadRandom. A single db_bench driver thread - // launches (-coro_jobs_per_thread * -threads) async MultiGet jobs onto the + // launches (-coro_jobs_per_thread * -threads) MultiGetCoroutine jobs onto the // FileSystem read executor and waits for them all. Because each job suspends // while its IO is in flight, jobs over-subscribe the executor and keep its // threads busy. @@ -7882,9 +7860,16 @@ class Benchmark { const int jobs_per_thread = FLAGS_coro_jobs_per_thread > 0 ? FLAGS_coro_jobs_per_thread : 1; const int num_jobs = jobs_per_thread * pool_threads; - auto* read_executor = GetReadExecutorForCoroutineBenchmark( - "multireadrandomcoroutine", pool_threads); DB* db = SelectDB(thread); + auto* read_executor = GetReadExecutorForCoroutineBenchmark( + db, "multireadrandomcoroutine", pool_threads); + CoroDB* coro_db = db->GetCoroDB(); + if (coro_db == nullptr) { + fprintf(stderr, + "multireadrandomcoroutine requires a DB with native coroutine " + "reads\n"); + ErrorExit(); + } ColumnFamilyHandle* cfh = db->DefaultColumnFamily(); std::atomic total_ops{0}; @@ -7919,9 +7904,9 @@ class Benchmark { futures.push_back( folly::coro::co_withExecutor( folly::Executor::getKeepAliveToken(read_executor->getEventBase()), - MultiGetAsyncCoroutineJob(db, cfh, thread->rand.Next(), ops, - &total_ops, &total_found, &total_bytes, - j, job_perf_contexts_ptr)) + MultiGetCoroutineJob(coro_db, cfh, thread->rand.Next(), ops, + &total_ops, &total_found, &total_bytes, j, + job_perf_contexts_ptr)) .start()); } for (auto& f : futures) { diff --git a/tools/db_crashtest_test.py b/tools/db_crashtest_test.py index d434bd9493..3db3ce1783 100644 --- a/tools/db_crashtest_test.py +++ b/tools/db_crashtest_test.py @@ -111,6 +111,20 @@ class DBCrashTestTest(unittest.TestCase): self.assertEqual("halt_on_error=1", env[_TSAN_OPTIONS_ENV_VAR]) + def test_generated_command_uses_async_db_api(self): + db_crashtest = self.load_db_crashtest() + params = self.build_params( + db_crashtest.default_params, + {"use_async_db_api": 1}, + ) + + command, _ = db_crashtest.gen_cmd(params, []) + + self.assertIn("--use_async_db_api=1", command) + self.assertFalse( + any(arg.startswith("--use_coro_db_api=") for arg in command) + ) + def test_get_ev_parent_dir_preserves_existing_contents(self): os.makedirs(self.expected_dir) marker = os.path.join(self.expected_dir, "marker") diff --git a/unreleased_history/public_api_changes/native_coroutine_reads.md b/unreleased_history/public_api_changes/native_coroutine_reads.md new file mode 100644 index 0000000000..9b2fd09748 --- /dev/null +++ b/unreleased_history/public_api_changes/native_coroutine_reads.md @@ -0,0 +1 @@ +Added the experimental `CoroDB` capability, exposed by `DB::GetCoroDB()`, with lazy native coroutine variants of `Get` and `MultiGet`. The existing `DB::GetAsync()` and `DB::MultiGetAsync()` APIs dispatch through it. diff --git a/util/coro_stats_util.cc b/util/coro_stats_util.cc index 8e7b1c1f10..cd7de52366 100644 --- a/util/coro_stats_util.cc +++ b/util/coro_stats_util.cc @@ -86,11 +86,19 @@ class CoroutineStatsRequestData : public folly::RequestData { PerfLevel CapturedPerfLevel() const { return perf_level_; } - void Finalize() { - assert(GetCoroutineStatsData() == this); - assert(!finalized_); - finalized_ = true; - StopGetCpuTimer(); + 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 } private: @@ -122,9 +130,6 @@ class CoroutineStatsRequestData : public folly::RequestData { void StartGetCpuTimer() { #ifndef NPERF_CONTEXT - if (finalized_) { - return; - } if (!PerfLevelAtLeast(PerfLevel::kEnableTimeAndCPUTimeExceptForMutex)) { return; } @@ -134,21 +139,6 @@ class CoroutineStatsRequestData : public folly::RequestData { #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 @@ -157,7 +147,6 @@ class CoroutineStatsRequestData : public folly::RequestData { #ifndef NIOSTATS_CONTEXT IOStatsContext iostats_context_; #endif - bool finalized_ = false; #ifndef NDEBUG uint64_t owner_thread_id_ = 0; #endif @@ -189,9 +178,22 @@ CoroutineStatsConfig CaptureCoroutineStatsConfig() { return stats_config; } +bool IsCoroutineStatsEnabled() { +#ifndef NPERF_CONTEXT + if (GetPerfLevel() != PerfLevel::kDisable) { + return true; + } +#endif +#ifndef NIOSTATS_CONTEXT + if (!get_iostats_context()->disable_iostats) { + return true; + } +#endif + return false; +} + CoroutineStatsContextScope::CoroutineStatsContextScope( - CoroutineStatsConfig stats_config, Env* env) - : request_data_(nullptr) { + CoroutineStatsConfig stats_config, Env* env) { assert(GetCoroutineStatsData() == nullptr); // NO nesting allowed auto data = std::make_unique(std::move(stats_config), env); @@ -200,12 +202,23 @@ CoroutineStatsContextScope::CoroutineStatsContextScope( CoroutineStatsRequestDataTraits::kToken, std::move(data)); } -CoroutineStatsContextScope::~CoroutineStatsContextScope() = default; - -void CoroutineStatsContextScope::Finalize() const { - auto* data = static_cast(request_data_); - assert(data != nullptr); - data->Finalize(); +CoroutineStatsContextScope::~CoroutineStatsContextScope() { + auto* request_data = static_cast(request_data_); + assert(GetCoroutineStatsData() == request_data); + request_data->StopGetCpuTimer(); +#ifndef NPERF_CONTEXT + PerfContext request_perf_context = std::move(*get_perf_context()); +#endif +#ifndef NIOSTATS_CONTEXT + IOStatsContext request_iostats_context = std::move(*get_iostats_context()); +#endif + guard_.reset(); +#ifndef NPERF_CONTEXT + *get_perf_context() = std::move(request_perf_context); +#endif +#ifndef NIOSTATS_CONTEXT + *get_iostats_context() = std::move(request_iostats_context); +#endif } } // namespace ROCKSDB_NAMESPACE diff --git a/util/coro_stats_util.h b/util/coro_stats_util.h index 278853f454..c92eee01ec 100644 --- a/util/coro_stats_util.h +++ b/util/coro_stats_util.h @@ -7,7 +7,15 @@ #pragma once #if defined(USE_COROUTINES) +#include #include +#include + +#ifndef NDEBUG +#include "folly/coro/CurrentExecutor.h" +#include "folly/executors/IOExecutor.h" +#include "folly/io/async/EventBase.h" +#endif #include "rocksdb/perf_level.h" #include "rocksdb/rocksdb_namespace.h" @@ -28,12 +36,13 @@ struct CoroutineStatsConfig { }; CoroutineStatsConfig CaptureCoroutineStatsConfig(); +bool IsCoroutineStatsEnabled(); // 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. +// Collected stats are published to thread-local storage on destruction. class CoroutineStatsContextScope { public: explicit CoroutineStatsContextScope(CoroutineStatsConfig stats_config, @@ -46,13 +55,30 @@ class CoroutineStatsContextScope { CoroutineStatsContextScope(CoroutineStatsContextScope&&) = delete; CoroutineStatsContextScope& operator=(CoroutineStatsContextScope&&) = delete; - void Finalize() const; - private: - folly::RequestData* request_data_; + folly::RequestData* request_data_ = nullptr; std::unique_ptr guard_; }; +#ifndef NDEBUG +#define INSTALL_COROUTINE_STATS_CONTEXT_SCOPE(read_executor_arg, env_arg) \ + auto* const coroutine_stats_read_executor = (read_executor_arg); \ + if (coroutine_stats_read_executor != nullptr) { \ + assert(co_await folly::coro::co_current_executor == \ + coroutine_stats_read_executor->getEventBase()); \ + } \ + std::optional stats_scope; \ + if (IsCoroutineStatsEnabled()) { \ + stats_scope.emplace(CaptureCoroutineStatsConfig(), (env_arg)); \ + } +#else +#define INSTALL_COROUTINE_STATS_CONTEXT_SCOPE(read_executor_arg, env_arg) \ + std::optional stats_scope; \ + if (IsCoroutineStatsEnabled()) { \ + stats_scope.emplace(CaptureCoroutineStatsConfig(), (env_arg)); \ + } +#endif + } // namespace ROCKSDB_NAMESPACE #endif // USE_COROUTINES diff --git a/util/coro_utils.h b/util/coro_utils.h index 3eb21054c7..d1c5ceacd0 100644 --- a/util/coro_utils.h +++ b/util/coro_utils.h @@ -46,14 +46,6 @@ namespace ROCKSDB_NAMESPACE { __template_decl__ folly::coro::Task<__ret_type__> __func_name__##Coroutine( \ __VA_ARGS__) const; -#define DECLARE_SYNC_ASYNC_AND_CALLBACK(__ret_type__, __func_name__, \ - __callback_func_name__, \ - __callback_args__, ...) \ - __ret_type__ __func_name__(__VA_ARGS__) override; \ - virtual folly::coro::Task<__ret_type__> __func_name__##Coroutine( \ - __VA_ARGS__); \ - void __callback_func_name__ __callback_args__ final override; - constexpr bool using_coroutines() { return true; } #else // !USE_COROUTINES @@ -75,12 +67,6 @@ constexpr bool using_coroutines() { return true; } __func_name__, ...) \ __template_decl__ __ret_type__ __func_name__(__VA_ARGS__) const; -#define DECLARE_SYNC_ASYNC_AND_CALLBACK(__ret_type__, __func_name__, \ - __callback_func_name__, \ - __callback_args__, ...) \ - __ret_type__ __func_name__(__VA_ARGS__) override; \ - void __callback_func_name__ __callback_args__ final override; - constexpr bool using_coroutines() { return false; } #endif // USE_COROUTINES #endif // UTIL_CORO_UTILS_H_