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

Add native coroutine read dispatch to RocksDB (#15073)

Browse Source 
Summary:
## Summary

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

The callback-based async read APIs make coroutine-native implementations and StackableDB composition awkward because wrappers must bridge callbacks back into tasks. This change introduces `CoroDB` as a side capability exposing lazy `GetCoroutine` and `MultiGetCoroutine` tasks. `DBImpl` and `CompactedDBImpl` implement this capability directly, while `DB::GetAsync` and `DB::MultiGetAsync` remain available as compatibility entry points with synchronous fallback.

The native interface should also be more efficient because callers directly await the DB coroutine. It avoids the callback/Baton adapter, callback context handoff, and explicit rescheduling required for each request. Native tasks must run on the DB filesystem read executor's EventBase; debug builds verify this affinity.

Public API documentation, a release note, and native-coroutine support in `db_bench` and `db_stress` are included.

## Dispatch

  DB::GetAsync
    -> GetCoroDB() on the concrete DB
    -> virtual CoroDB::GetCoroutine(...)
       -> DBImpl::GetCoroutine or CompactedDBImpl::GetCoroutine
    -> AsyncCallback::OnComplete(...)

Virtual dispatch selects the coroutine implementation belonging to the concrete DB while preserving the existing callback API.

## Statistics

Coroutine reads preserve RocksDB's TLS statistics behavior while keeping concurrent requests isolated. Each top-level task copies only the caller's statistics configuration—`PerfLevel`, per-level perf context enablement, and IO statistics enablement—into fresh request-local `PerfContext` and `IOStatsContext` objects. Existing counters are not inherited. If both perf and IO statistics are disabled, no request statistics context is installed.

Folly `RequestContext` preserves these contexts across suspension and resumption on the read executor's EventBase. At completion, RocksDB publishes the request's collected values to TLS, allowing the awaiting caller to use `get_perf_context()` and `get_iostats_context()` as after a synchronous read. Callers must consume or copy them before leaving that EventBase. Statistics scopes are installed only by top-level coroutine reads and cannot be nested.

## Benchmarks

Compared the previous `GetAsync` callback bridge with native `GetCoroutine`. Both variants used optimized builds, the same compacted DB, a fixed seed, eight executor threads, ten jobs per thread, and 2.56M reads per invocation. Results are the mean of three independent invocations using the default RocksDB block cache. A sequential read warmed the OS page cache before the non-direct-read measurements.

| Workload | `GetAsync` bridge | Native `GetCoroutine` | Change |
| Direct I/O | 237,710 ops/s | 243,588 ops/s | +2.47% |
| Warm OS page cache | 783,859 ops/s | 891,531 ops/s | +13.74% |

The direct-I/O difference is within observed run-to-run variation. When reads are served from the OS page cache and API overhead is a larger portion of latency, the native interface improves throughput by 13.74%. Every invocation completed all requested reads.

```bash
db_bench \
  --benchmarks=readrandomcoroutine \
  --db=$CORO_BENCH_DB \
  --use_existing_db=true \
  --num=2000000 \
  --value_size=1024 \
  --compression_type=none \
  --threads=8 \
  --coro_jobs_per_thread=10 \
  --reads=320000 \
  --use_direct_reads=<true|false> \
  --perf_level=1 \
  --statistics=true \
  --seed=20260807
```

Reviewed By: xingbowang

Differential Revision: D115232511

fbshipit-source-id: 30b48da2a798094d9080178f5430aae26e6c5954
This commit is contained in:
Josh Kang
2026-08-10 09:43:05 -07:00
committed by meta-codesync[bot]
parent 1dfa729793
commit 88752298ff
18 changed files with 557 additions and 382 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
.github/workflows/pr-jobs.yml vendored
View File

@@ -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"

7
Makefile
View File

@@ -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 :=

11
db/db_basic_test.cc
View File

@@ -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);

12
db/db_impl/compacted_db_impl_sync_and_async.h
View File

@@ -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 &&

190
db/db_impl/db_impl.cc
View File

@@ -15,13 +15,6 @@
#include <alloca.h>
#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 <cinttypes>
#include <cstdio>
#include <deque>
@@ -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<CoroutineStatsConfig> CaptureCoroutineStatsConfigForCallback(
bool stats_enabled) {
if (!stats_enabled) {
if (!stats_enabled || !IsCoroutineStatsEnabled()) {
return std::nullopt;
}
return CaptureCoroutineStatsConfig();
@@ -7784,49 +7801,49 @@ std::optional<CoroutineStatsConfig> 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<CoroutineStatsConfig> task_stats_config, DBImpl* db,
ReadOptions task_options, ColumnFamilyHandle* task_column_family,
Slice task_key, PinnableSlice* task_value,
std::string* task_timestamp, Status& task_status,
AsyncCallback& task_callback) mutable -> folly::coro::Task<void> {
std::optional<CoroutineStatsContextScope> coroutine_stats_scope;
if (task_stats_config.has_value()) {
coroutine_stats_scope.emplace(std::move(*task_stats_config),
db->immutable_db_options_.env);
}
task_status = co_await folly::coro::co_nothrow(
db->GetCoroutine(task_options, task_column_family, task_key,
task_value, task_timestamp));
if (coroutine_stats_scope.has_value()) {
coroutine_stats_scope->Finalize();
}
task_callback.OnComplete(
CurrentPerfContextForAsyncCallback(coroutine_stats_scope.has_value()),
CurrentIOStatsContextForAsyncCallback(
coroutine_stats_scope.has_value()));
}(CaptureCoroutineStatsConfigForCallback(stats_enabled), this, options,
column_family, key, value,
timestamp, status, callback);
folly::coro::co_withExecutor(
folly::Executor::getKeepAliveToken(read_event_base), std::move(task))
.start();
return;
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<CoroutineStatsConfig> 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<void> {
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<CoroutineStatsConfig> task_stats_config,
DBImpl* db, ReadOptions task_options, size_t task_num_keys,
ColumnFamilyHandle** task_column_families,
const Slice* task_keys, PinnableSlice* task_values,
std::string* task_timestamps, Status* task_statuses,
bool task_sorted_input, AsyncCallback& task_callback) mutable
-> folly::coro::Task<void> {
std::optional<CoroutineStatsContextScope> coroutine_stats_scope;
if (task_stats_config.has_value()) {
coroutine_stats_scope.emplace(std::move(*task_stats_config),
db->immutable_db_options_.env);
}
co_await folly::coro::co_nothrow(db->MultiGetCoroutine(
task_options, task_num_keys, task_column_families, task_keys,
task_values, task_timestamps, task_statuses, task_sorted_input));
if (coroutine_stats_scope.has_value()) {
coroutine_stats_scope->Finalize();
}
task_callback.OnComplete(
CurrentPerfContextForAsyncCallback(coroutine_stats_scope.has_value()),
CurrentIOStatsContextForAsyncCallback(
coroutine_stats_scope.has_value()));
}(CaptureCoroutineStatsConfigForCallback(stats_enabled), this, options,
num_keys, column_families, keys, values, timestamps, statuses,
sorted_input, callback);
folly::coro::co_withExecutor(
folly::Executor::getKeepAliveToken(read_event_base), std::move(task))
.start();
return;
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<CoroutineStatsConfig> 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<void> {
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,

50
db/db_impl/db_impl.h
View File

@@ -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_;

44
db/db_impl/db_impl_sync_and_async.h
View File

@@ -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

56
db/db_test_util.cc
View File

@@ -24,7 +24,14 @@
#include "util/random.h"
#if USE_COROUTINES
#include <utility>
#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 <typename T>
T BlockingWaitOnReadExecutor(DB* db, folly::coro::Task<T> 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<std::string> DBTestBase::MultiGet(std::vector<int> 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<std::string> DBTestBase::MultiGet(const std::vector<std::string>& k,
if (use_coroutine) {
std::vector<ColumnFamilyHandle*> 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));
}

135
db/perf_context_test.cc
View File

@@ -112,44 +112,45 @@ TEST_F(PerfContextTest, CoroutineStatsContextScopeCollectsStats) {
std::move(stats_config)]() mutable -> folly::coro::Task<void> {
EXPECT_TRUE(event_base->isInEventBaseThread());
CoroutineStatsContextScope scope(std::move(stats_config),
Env::Default());
EXPECT_EQ(PerfLevel::kEnableTimeAndCPUTimeExceptForMutex,
GetPerfLevel());
EXPECT_EQ(0, get_perf_context()->block_read_count);
EXPECT_TRUE(get_perf_context()->per_level_perf_context_enabled);
get_perf_context()->block_read_count += 3;
(*get_perf_context()->level_to_perf_context)[2].block_cache_hit_count +=
5;
EXPECT_EQ(0, get_iostats_context()->bytes_read);
get_iostats_context()->bytes_read += 7;
get_iostats_context()->cpu_read_nanos += 11;
get_iostats_context()
->file_io_stats_by_temperature.hot_file_read_count += 1;
{
PERF_TIMER_GUARD_WITH_CLOCK(block_read_time, clock);
clock->SleepForMicroseconds(7);
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);
}
}

78
include/rocksdb/coro_db.h Normal file
View File

@@ -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 <cstddef>
#include <string>
#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 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<Status> 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<void> 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

43
include/rocksdb/db.h
View File

@@ -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 {

3
include/rocksdb/file_system.h
View File

@@ -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

161
tools/db_bench_tool.cc
View File

@@ -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<folly::RequestContext> caller_request_context)
: job_perf_context_(job_perf_context),
caller_request_context_(std::move(caller_request_context)) {}
bool EnableStats() const override { return job_perf_context_ != nullptr; }
void OnComplete(const PerfContext* callback_perf_context,
const IOStatsContext* /*iostats_context*/) override {
if (job_perf_context_ != nullptr && callback_perf_context != nullptr) {
job_perf_context_->Merge(*callback_perf_context);
}
folly::RequestContextScopeGuard context_scope(caller_request_context_);
baton.post();
}
PerfContext* job_perf_context_;
std::shared_ptr<folly::RequestContext> caller_request_context_;
folly::coro::Baton baton;
};
void PrepareCoroutineJobPerfContext(PerfContext* job_perf_context) {
SetPerfLevel(static_cast<PerfLevel>(FLAGS_perf_level));
if (job_perf_context == nullptr) {
@@ -7643,34 +7618,14 @@ class Benchmark {
#endif
}
folly::coro::Task<Status> 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<void> 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<void> ReadRandomCoroutineJob(
uint64_t seed, int64_t ops, std::atomic<int64_t>* total_ops,
const std::vector<DBWithColumnFamilies*>& db_choices, uint64_t seed,
int64_t ops, std::atomic<int64_t>* total_ops,
std::atomic<int64_t>* total_found, std::atomic<int64_t>* total_bytes,
int job_id, std::vector<std::string>* job_perf_contexts) {
Random64 rng(seed);
@@ -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<void> MultiGetAsyncCoroutineJob(
DB* db, ColumnFamilyHandle* cfh, uint64_t seed, int64_t ops,
folly::coro::Task<void> MultiGetCoroutineJob(
CoroDB* coro_db, ColumnFamilyHandle* cfh, uint64_t seed, int64_t ops,
std::atomic<int64_t>* total_ops, std::atomic<int64_t>* total_found,
std::atomic<int64_t>* total_bytes, int job_id,
std::vector<std::string>* 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<DBWithColumnFamilies*> 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<int64_t> total_ops{0};
std::atomic<int64_t> 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<int64_t> 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) {

14
tools/db_crashtest_test.py
View File

@@ -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")

1
unreleased_history/public_api_changes/native_coroutine_reads.md Normal file
View File

@@ -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.

77
util/coro_stats_util.cc
View File

@@ -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<CoroutineStatsRequestData>(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<CoroutineStatsRequestData*>(request_data_);
assert(data != nullptr);
data->Finalize();
CoroutineStatsContextScope::~CoroutineStatsContextScope() {
auto* request_data = static_cast<CoroutineStatsRequestData*>(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

34
util/coro_stats_util.h
View File

@@ -7,7 +7,15 @@
#pragma once
#if defined(USE_COROUTINES)
#include <cassert>
#include <memory>
#include <optional>
#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<folly::ShallowCopyRequestContextScopeGuard> 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<CoroutineStatsContextScope> stats_scope; \
if (IsCoroutineStatsEnabled()) { \
stats_scope.emplace(CaptureCoroutineStatsConfig(), (env_arg)); \
}
#else
#define INSTALL_COROUTINE_STATS_CONTEXT_SCOPE(read_executor_arg, env_arg) \
std::optional<CoroutineStatsContextScope> stats_scope; \
if (IsCoroutineStatsEnabled()) { \
stats_scope.emplace(CaptureCoroutineStatsConfig(), (env_arg)); \
}
#endif
} // namespace ROCKSDB_NAMESPACE
#endif // USE_COROUTINES

14
util/coro_utils.h
View File

@@ -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_