Isolate coroutine statistics across async reads

[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

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 296.3s.

[github-actions]

✅ 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