Isolate coroutine statistics across async reads (#15142)

[joshkang97]

Summary:

Coroutine statistics use TLS, but coroutine reads can interleave on the same executor thread. Previously, when a stats-enabled request suspended, it saved its counters without disabling the executor's TLS configuration. A stats-disabled request running next could inherit those enabled settings and collect statistics unexpectedly. Stats setup also lived in individual DB implementations, so wrapper early returns and future stackable DB implementations could bypass it.

Move statistics ownership to the public CoroDB and callback-based async read boundaries, before virtual dispatch. Each operation captures its caller's configuration, enabled requests preserve their counters across suspensions, and every suspension or completion leaves executor TLS disabled. Because the scope wraps stackable DB dispatch, current and future CoroStackableDB implementations automatically inherit the correct behavior without adding their own stats reset or scope.

For consistency between coroutine and callback-based reads, the experimental callback API now uses TLS for both configuring and consuming statistics. AsyncCallback::OnComplete() no longer receives context arguments; RocksDB publishes the completed per-operation counters to TLS before invoking it.

Before

time -------------------------------------------------------------->

stats-on A:   onSet(enable A) -- work -- onUnset(save A) ...... onSet(A)
executor TLS: [A enabled] ------------> [still enabled] ------> [A enabled]
stats-off B:                                  run
                                               ^ inherits A's enabled TLS

After

time -------------------------------------------------------------->

stats-on A:   onSet(enable A) -- work -- onUnset(save A, disable) ... onSet(A)
executor TLS: [A enabled] ------------> [disabled] ----------------> [A enabled]
stats-off B:                                  run
                                               ^ remains stats-disabled

Reviewed By: xingbowang

Differential Revision: D117017089

[meta-codesync]

@joshkang97 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D117017089.

[github-actions]

✅ clang-tidy: No findings on changed lines

Completed in 266.9s.

[github-actions]

Claude Code Review - OBSOLETE

Superseded by a newer AI review. Expand to see the original review.

✅ Claude Code Review

Auto-triggered after CI passed — reviewing commit a5c537f

---

Summary

Well-designed fix for a real TLS stats isolation bug in coroutine/async reads. The approach of moving stats ownership to the CoroDB boundary (before virtual dispatch) is sound and eliminates duplicated stats setup across DB implementations. The breaking AsyncCallback API change is justified given the experimental status.

High-severity findings (0):
No high-severity findings.

Full review (click to expand)

Findings

🟡 MEDIUM

M1. CoroutineStatsConfig default member initializers enable stats unexpectedly -- util/coro_stats_util.h:33

• Issue: The default-constructed CoroutineStatsConfig has perf_level = PerfLevel::kEnableCount and iostats_disabled = false, which means IsCoroutineStatsEnabled() returns true. Any code path that inadvertently uses a default-constructed config (e.g., forgetting to capture from TLS) would silently enable stats collection and allocate EnabledCoroutineStatsRequestData.
• Root cause: The defaults were designed to match the TLS defaults (perf_level TLS default is kEnableCount per monitoring/perf_level.cc:13), but this makes the "zero-value" config surprising.
• Suggested fix: Consider changing defaults to the disabled state (perf_level = PerfLevel::kDisable, iostats_disabled = true) so that a forgotten-capture scenario fails safe (no stats, no allocation) rather than fails open. This would require updating CaptureCoroutineStatsConfig() -- which already explicitly sets all fields from TLS -- so no behavior change there.

M2. explicit removed from CoroutineStatsContextScope constructor -- util/coro_stats_util.h

• Issue: The diff removes explicit from the 2-parameter constructor. While C++ allows implicit conversion via braced-init-list for non-explicit multi-parameter constructors, this loosens the API surface unnecessarily.
• Suggested fix: Keep explicit on the constructor.

M3. PrepareCoroutineJobPerfContext sets iostats_disabled = true unconditionally -- tools/db_bench_tool.cc:7846

• Issue: The new PrepareCoroutineJobPerfContext unconditionally sets get_iostats_context()->disable_iostats = true. The old version did not touch iostats at all. This means db_bench coroutine reads will never collect IO stats, even if a user wanted them. Since MergeCoroutineJobPerfContext only merges PerfContext (not IOStatsContext), this is likely intentional but is an undocumented behavioral change.
• Suggested fix: Add a brief comment explaining that iostats are intentionally disabled because db_bench only merges PerfContext.

M4. AsyncReadStatsScope always disables stats on exit, even for stats-disabled callers -- db/db_impl/db_impl.cc

• Issue: In the sync fallback path of GetAsync/MultiGetAsync, AsyncReadStatsScope destructor always calls DisableThreadLocalStatsForAsyncRead() which sets perf_level = kDisable and iostats_disabled = true. Previously, stats were only reset/disabled if EnableStats() returned true. This means a caller that had stats enabled for other purposes (e.g., tracking sync reads) will find their TLS stats disabled after calling GetAsync even if they never cared about async stats.
• Root cause: This is an intentional design choice documented in the new API ("Async reads reset the calling thread's configuration to disabled"), but it's a behavioral regression for callers that previously called GetAsync with EnableStats() = false and expected TLS to remain untouched.
• Suggested fix: The API documentation covers this, but consider whether the sync-fallback path (no coroutine support) should also unconditionally disable. The argument for doing so is consistency, which is reasonable.

🟢 LOW / NIT

L1. Duplicate DisableCoroutineStatsInTLS() calls in CoroutineStatsContextScope destructor -- util/coro_stats_util.cc

• Issue: When guard_ is non-null, the destructor path is: save TLS stats, guard_.reset() (which triggers onUnset -> SaveThreadLocalStats + DisableCoroutineStatsInTLS), restore stats to TLS, then DisableCoroutineStatsInTLS again. The disable in onUnset operates on already-moved-out TLS state (harmless), and the final explicit disable is the one that matters. The double-disable is correct but slightly confusing.
• Suggested fix: Consider adding a brief comment explaining why the double-disable is intentional (onUnset can't know we'll restore afterward).

L2. InstallCoroutineStatsConfigToTLS and DisableCoroutineStatsInTLS are not in the anonymous namespace -- util/coro_stats_util.cc

• Issue: These functions are defined outside the anonymous namespace in coro_stats_util.cc but are not declared in the header. They have external linkage but no header declaration, meaning they could accidentally be called from other translation units via extern declaration.
• Suggested fix: Either move them into the anonymous namespace or add declarations to the header if they're intended to be used externally.

L3. ManualExecutor vs IOThreadPoolExecutor in test -- db/perf_context_test.cc

• Issue: The test switches from IOThreadPoolExecutor to ManualExecutor. ManualExecutor is better for determinism (no actual threading), but it changes the execution model. The test no longer validates cross-thread behavior.
• Suggested fix: This tradeoff (determinism vs realism) is acceptable for a unit test. The important property tested is interleaving, which ManualExecutor still provides through collectAll.

Cross-Component Analysis

Context	Affected by this PR?	Stats isolation correct?	Notes
WritePreparedTxnDB	YES (INSTALL macro removed)	YES	Stats now managed at CoroDB::CoGet boundary before dispatch
CompactedDBImpl	YES (INSTALL macro removed)	YES	Same as above
DBImpl	YES (INSTALL macro removed)	YES	Same as above
CoroStackableDB	Indirectly	YES	GetCoroutine/MultiGetCoroutine called from within CoroDB::CoGet which manages stats scope
ReadOnly/Secondary DB	NO	N/A	Don't implement CoroDB
GetAsync sync fallback	YES (AsyncReadStatsScope)	YES	Always disables stats on exit
db_bench coroutine reads	YES	YES	Per-iteration PrepareCoroutineJobPerfContext + MergeCoroutineJobPerfContext works correctly
db_stress	YES (OnComplete signature)	YES	Updated, doesn't use stats

Assumption stress test results:

1. "Every suspension/completion leaves TLS disabled" -- Verified. onUnset() calls DisableCoroutineStatsInTLS(). Scope destructor calls it. Short-circuit constructor calls it. AsyncReadStatsScope destructor calls it.

2. "Stats-disabled requests don't accumulate stats" -- Verified. When IsCoroutineStatsEnabled(stats_config) returns false, no EnabledCoroutineStatsRequestData is created, no folly request context is installed, and TLS is disabled.

3. "Stats data survives DisableCoroutineStatsInTLS for reading" -- Verified. DisableCoroutineStatsInTLS only sets perf_level and iostats_disabled flag. The actual counter values in get_perf_context() and get_iostats_context() are untouched. MergeCoroutineJobPerfContext reads the counters, not the flags.

Positive Observations

• Clean elimination of the INSTALL_COROUTINE_STATS_CONTEXT_SCOPE macro, which was duplicated across 4 files and required each implementation to independently manage stats.
• The short-circuit path for disabled configs avoids unnecessary folly::RequestData allocation.
• The test correctly validates the new invariant that TLS is disabled after blockingWait returns.
• The EnabledCoroutineStatsRequestData naming clearly distinguishes it from the disabled path.
• Release notes are appropriately concise.

---

ℹ️ About this response

Generated by Claude Code.
Review methodology: claude_md/ci_review_prompt.md

Limitations:

• Claude may miss context from files not in the diff
• Large PRs may be truncated
• Always apply human judgment to AI suggestions

Commands:

• /claude-review [context] — Request a code review
• /claude-query <question> — Ask about the PR or codebase

[github-actions]

Claude Code Review - OBSOLETE

Superseded by a newer AI review. Expand to see the original review.

✅ Claude Code Review

Auto-triggered after CI passed — reviewing commit 6e29b5f

---

Summary

Well-designed fix for a real TLS stats isolation bug in coroutine reads. The approach of moving stats ownership to the CoroDB boundary is architecturally sound and eliminates the class of bugs where stackable DB wrappers could bypass stats setup. The breaking AsyncCallback API change is justified by the experimental status.

High-severity findings (0):
No high-severity findings.

Full review (click to expand)

Findings

🟡 MEDIUM

M1. CoroutineStatsConfig default constructor creates stats-enabled config -- util/coro_stats_util.h:32

• Issue: CoroutineStatsConfig defaults to perf_level = kEnableCount and iostats_disabled = false, meaning a default-constructed config is stats-ENABLED. If any code path constructs a CoroutineStatsConfig without going through CaptureCoroutineStatsConfig(), it will appear stats-enabled and incur the overhead of EnabledCoroutineStatsRequestData and folly RequestContext setup.
• Root cause: The defaults predate this PR and match the TLS defaults (perf_level starts at kEnableCount, disable_iostats starts at false). The struct was designed to snapshot TLS, not to be constructed directly.
• Suggested fix: Consider adding a comment to CoroutineStatsConfig documenting that the defaults mirror TLS initialization and the struct should normally be populated via CaptureCoroutineStatsConfig(). Alternatively, a named static factory like CoroutineStatsConfig::Disabled() could prevent accidental misuse.

M2. AsyncCallbackStatsStartEmptyOnSyncFallback test relies on re-enabling stats between async calls -- db/perf_context_test.cc:377

• Issue: The test correctly adds SetPerfLevel(PerfLevel::kEnableCount) and disable_iostats = false before the second async call (MultiGetAsync), since the first GetAsync now disables TLS stats on completion. However, the test does NOT verify that TLS stats are actually disabled after the first GetAsync returns. Adding EXPECT_EQ(PerfLevel::kDisable, GetPerfLevel()) after the first callback would strengthen the test and validate the core invariant.
• Root cause: Test covers the functional change but misses verifying the "disable on exit" contract for the sync fallback path.
• Suggested fix: Add assertions after each async call returns to verify PerfLevel is kDisable and iostats are disabled.

M3. db_bench now unconditionally disables IO stats for coroutine reads -- tools/db_bench_tool.cc:7851

• Issue: PrepareCoroutineJobPerfContext now sets disable_iostats = true unconditionally. The old code did not touch iostats, so coroutine reads could previously collect IO statistics. This is a behavioral change for db_bench users who relied on IO stats from readrandomcoroutine or multireadrandomcoroutine.
• Root cause: The new design requires callers to explicitly configure stats before each async/coroutine read. Since MergeCoroutineJobPerfContext only merges PerfContext (not IOStatsContext), disabling IO stats is functionally consistent. But it is an undocumented behavioral change.
• Suggested fix: Add a brief code comment explaining why iostats are disabled (not merged by db_bench coroutine jobs). If IO stats are desired in the future, MergeCoroutineJobPerfContext would need to merge IOStatsContext too.

🟢 LOW / NIT

L1. onSet() ordering change is safe but subtle -- util/coro_stats_util.cc:96

• Issue: The old onSet() called SetPerfLevel() before LoadThreadLocalStats(). The new code calls LoadThreadLocalStats() then InstallCoroutineStatsConfigToTLS(). Since PerfLevel is a separate TLS variable (not part of PerfContext), and InstallCoroutineStatsConfigToTLS also sets per_level_perf_context_enabled (which IS in PerfContext, potentially overwriting what LoadThreadLocalStats just loaded), the ordering is still correct: LoadThreadLocalStats restores the counters, then InstallCoroutineStatsConfigToTLS restores the configuration flags. The per_level_perf_context_enabled field is loaded by LoadThreadLocalStats (as part of the full PerfContext move) and then potentially overwritten by InstallCoroutineStatsConfigToTLS. This is safe because stats_config_.per_level_perf_context_enabled should match what was saved, but a comment would help future readers.
• Suggested fix: Consider adding a brief comment in onSet() explaining that InstallCoroutineStatsConfigToTLS re-applies the authoritative config after LoadThreadLocalStats restores counters.

L2. explicit keyword removed from CoroutineStatsContextScope constructor -- util/coro_stats_util.h

• Issue: The old constructor was explicit CoroutineStatsContextScope(CoroutineStatsConfig, Env*). The diff removes explicit. Since this is a two-parameter constructor, explicit is not strictly necessary, but keeping it is a common RocksDB convention for non-trivial constructors.
• Suggested fix: Keep explicit for consistency with RocksDB style.

L3. Destructor DisableCoroutineStatsInTLS() call after publishing stats is correct but non-obvious -- util/coro_stats_util.cc:260

• Issue: In the enabled path, guard_.reset() triggers onUnset() which calls DisableCoroutineStatsInTLS(). Then after restoring stats to TLS, the destructor calls DisableCoroutineStatsInTLS() again. The second call is necessary because *get_perf_context() = std::move(request_perf_context) restores fields like per_level_perf_context_enabled, and *get_iostats_context() restores disable_iostats. The final DisableCoroutineStatsInTLS() re-disables after publishing.
• Suggested fix: Add a comment like "Re-disable collection after publishing request stats to TLS".

L4. Missing test for coroutine-path GetAsync stats delivery -- db/perf_context_test.cc

• Issue: AsyncCallbackStatsStartEmptyOnSyncFallback only tests the sync fallback path (no read executor). There is no test verifying that GetAsync with a coroutine-capable DB correctly publishes stats to TLS before OnComplete() runs.
• Suggested fix: Consider adding a test for the coroutine path of GetAsync.

Cross-Component Analysis

Context	Affected?	Analysis
WritePreparedTxnDB	Yes	INSTALL_COROUTINE_STATS_CONTEXT_SCOPE removed from all 4 methods. Stats now correctly handled by CoroDB::CoGet wrapper before virtual dispatch. Safe.
CompactedDBImpl	Yes	Same removal, same correct wrapping. Safe.
CoroStackableDB	Indirectly	GetCoroutine just forwards to inner DB. Stats scope in CoGet wraps the entire dispatch chain. Safe.
db_stress	Yes	BlockingAsyncCallback::OnComplete() signature updated. No stats usage. Safe.
db_bench	Yes	PrepareCoroutineJobPerfContext moved inside loop, now disables iostats. Correct for per-operation capture pattern.
Sync DB::Get callers	No	Stats TLS only modified by async/coroutine paths.
ReadOnly DB	No	No coroutine support.
User-defined timestamps	No	Orthogonal to stats mechanism.

Positive Observations

• Moving stats ownership to the CoroDB boundary is a clean architectural improvement that eliminates the need for every DB implementation to independently manage stats, reducing error surface for future stackable DB implementations.
• The disabled-config fast path in CoroutineStatsContextScope (skipping folly RequestContext setup) avoids overhead for stats-disabled requests.
• The AsyncReadStatsScope RAII class for the sync fallback is clean and prevents resource leaks.
• The test improvements (ManualExecutor, disabled-stats task, post-completion TLS verification) are well-crafted.
• The release notes are appropriately concise.

---

ℹ️ About this response

Generated by Claude Code.
Review methodology: claude_md/ci_review_prompt.md

Limitations:

• Claude may miss context from files not in the diff
• Large PRs may be truncated
• Always apply human judgment to AI suggestions

Commands:

• /claude-review [context] — Request a code review
• /claude-query <question> — Ask about the PR or codebase

[github-actions]

✅ Claude Code Review

Auto-triggered after CI passed — reviewing commit ea3b2ec

---

Summary

Well-structured fix for a real correctness bug in coroutine stats isolation. The "capture-and-disable" pattern is clean, the macro elimination simplifies maintenance, and the API change unifies stats consumption. No high-severity issues found.

High-severity findings (0):

No high-severity findings.

Full review (click to expand)

Findings

🟡 MEDIUM

M1. CoroutineStatsConfig default member initializers are stats-enabled — util/coro_stats_util.h:33

• Issue: CoroutineStatsConfig defaults to perf_level = kEnableCount, iostats_disabled = false — i.e., stats enabled. A default-constructed config passed to CoroutineStatsContextScope would allocate EnabledCoroutineStatsRequestData and install a full request-context scope even when the caller intended no stats. The test creates a disabled config manually, showing the defaults are misleading.
• Root cause: The struct's defaults predate this PR and were designed for the old capture-then-check pattern. Now that CoroutineStatsContextScope directly consults IsCoroutineStatsEnabled(stats_config), the defaults matter more.
• Suggested fix: Consider changing defaults to perf_level = PerfLevel::kDisable and iostats_disabled = true (i.e., disabled by default). All real usage goes through CaptureCoroutineStatsConfig() which overwrites the defaults, so this would be a safe change and would make default-construction yield a disabled config.

M2. AsyncReadStatsScope unconditionally disables TLS stats for sync-fallback callers — db/db_impl/db_impl.cc:7826

• Issue: AsyncReadStatsScope destructor always calls DisableThreadLocalStatsForAsyncRead(). In the sync fallback path (no coroutine executor), this means calling GetAsync() or MultiGetAsync() now always leaves TLS stats disabled, even for callers who never opted into callback stats. Previously, TLS was untouched when EnableStats() returned false. This is a behavioral change for any code that: (1) sets PerfLevel before calling GetAsync, (2) doesn't set PerfLevel after GetAsync returns, and (3) expects TLS stats to remain enabled for subsequent sync operations.
• Root cause: The new design intentionally always disables TLS to prevent cross-request inheritance. This is documented in the updated db.h comments.
• Suggested fix: This is intentional and documented. Verify no internal callers rely on the old behavior. The db_bench changes correctly re-set stats per-iteration. The db.h docs say "Callers must set the desired configuration for each async read" and "Each operation leaves TLS statistics collection disabled."

M3. onSet ordering change may cause brief inconsistency — util/coro_stats_util.cc:98

• Issue: Old onSet: SetPerfLevel(...) then LoadThreadLocalStats(). New onSet: LoadThreadLocalStats() then InstallCoroutineStatsConfigToTLS(...). With the old order, PerfLevel was set before the stats context was loaded, so any PerfLevel-gated macro during the move would see the correct level. With the new order, during LoadThreadLocalStats(), TLS PerfLevel is whatever the previous request left (which is now kDisable). Since LoadThreadLocalStats only does std::move assignments (no PerfLevel-gated operations), this is harmless.
• Root cause: The reordering was needed because InstallCoroutineStatsConfigToTLS is now a separate function that also sets per_level_perf_context_enabled. The old code relied on SetPerfLevel being separate from loading stats.
• Suggested fix: No action needed — the move assignments in LoadThreadLocalStats do not consult PerfLevel.

🟢 LOW / NIT

L1. (void)stats_config suppression in IsCoroutineStatsEnabled — util/coro_stats_util.cc:201

• Issue: The function has (void)stats_config; at the top to suppress unused-variable warnings when both NPERF_CONTEXT and NIOSTATS_CONTEXT are defined. This is a standard pattern but slightly unusual placement before the #ifndef blocks.
• Suggested fix: Move the (void)stats_config to the end of the function, just before return false, or remove it and annotate the parameter with [[maybe_unused]].

L2. DisablePerLevelPerfContext() added in db_bench — tools/db_bench_tool.cc:7843

• Issue: PrepareCoroutineJobPerfContext now explicitly calls DisablePerLevelPerfContext() when job_perf_context_ptr == nullptr. Previously it just returned early. This is correct for the new per-iteration pattern but slightly changes the behavior for non-perf-context benchmarks.
• Suggested fix: No action needed — this matches the new "each call re-sets full config" pattern.

L3. Test uses ManualExecutor instead of IOThreadPoolExecutor — db/perf_context_test.cc:188

• Issue: The test CoroutineStatsContextsRemainIsolatedWhenInterleaved switched from IOThreadPoolExecutor to ManualExecutor. This gives better determinism for interleaving but loses the real-threading aspect. The test now runs entirely on the calling thread.
• Suggested fix: This is fine for testing the stats isolation logic. Consider adding a separate stress test variant with real threading.

L4. CoroutineStatsContextScope no longer explicit — util/coro_stats_util.h:49

• Issue: The constructor changed from explicit to non-explicit. This allows implicit construction from {stats_config, env} which is generally undesirable for two-argument constructors.
• Suggested fix: Keep explicit on the constructor.

Cross-Component Analysis

Context	Affected?	Analysis
WritePreparedTxnDB	YES	INSTALL_COROUTINE_STATS_CONTEXT_SCOPE removed. Stats now handled at CoroDB::CoGet layer before dispatching to WritePreparedTxnDB::GetCoroutine. Correct — the stackable DB chain is wrapped by the scope.
CompactedDBImpl	YES	Same as above. Stats scope wraps the entire GetCoroutine call chain.
CoroStackableDB	OK	Forwards to inner CoroDB. Stats scope at CoroDB::CoGet wraps the full stackable chain.
Sync fallback (no executor)	YES	CoroutineStatsContextScope wraps sync db->Get() call. TLS stats disabled after scope exits. Correct.
db_bench	YES	PrepareCoroutineJobPerfContext moved inside loop. Correctly re-sets stats before each CoGet since CoGet now disables TLS after each call.
db_stress	YES	BlockingAsyncCallback::OnComplete signature updated. No stats usage, so the change is transparent.

Positive Observations

1. Clean architectural improvement: Moving stats ownership from scattered DB implementations to the CoroDB boundary is the right layering. This eliminates the need for each DB impl to remember to add the stats macro.
2. Good bug fix: The core bug (stats-disabled requests inheriting enabled TLS state) was real and could cause incorrect metrics in production.
3. Allocation avoidance: When stats are disabled, no EnabledCoroutineStatsRequestData or ShallowCopyRequestContextScopeGuard is allocated. This is a performance improvement for the common "no stats" path.
4. Test improvement: Using ManualExecutor with collectAll and adding a disabled-stats task directly tests the reported bug scenario.
5. Consistent disable-on-exit: The "always disable TLS on exit" invariant is simple and prevents the class of bugs where enabled state leaks between requests.

---

ℹ️ About this response

Generated by Claude Code.
Review methodology: claude_md/ci_review_prompt.md

Limitations:

• Claude may miss context from files not in the diff
• Large PRs may be truncated
• Always apply human judgment to AI suggestions

Commands:

• /claude-review [context] — Request a code review
• /claude-query <question> — Ask about the PR or codebase

[meta-codesync]

This pull request has been merged in ac0f4ae.