8040: Allow impl FnMut() in taskdumps instead of just fn()#94
8040: Allow impl FnMut() in taskdumps instead of just fn()#94martin-augment wants to merge 4 commits intomasterfrom
Conversation
WalkthroughThe changes refactor the task tracing system to support mutable closures in callbacks instead of static function pointers, and to pass trace data through function parameters rather than relying on thread-local storage. The ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the task tracing infrastructure to support FnMut closures instead of requiring function pointers. By updating trace_with and internal context management to handle dynamic closures via NonNull, the implementation eliminates the need for several thread-local variables in both the runtime and tests. Feedback indicates that the documentation and examples for trace_with should be updated to reflect this new capability and the more idiomatic usage of closures for state capture.
| /// # } | ||
| /// ``` | ||
| pub fn trace_with<F, R>(f: F, trace_leaf: fn(&TraceMeta)) -> R | ||
| pub fn trace_with<FN, FT, R>(f: FN, mut trace_leaf: FT) -> R |
There was a problem hiding this comment.
The documentation and example for trace_with (lines 170-211) are now outdated. They still state that trace_leaf must be a function pointer (fn) and demonstrate using thread-locals for communication. Since the API now supports FnMut, the documentation should be updated to reflect that closures can be used to capture state directly, which is more idiomatic and simplifies the usage.
There was a problem hiding this comment.
value:useful; category:documentation; feedback: The Gemini AI reviewer is correct! The way the function works has changed - it no more needs the two extra thread locals and its signature changed a bit. Its documentation should be updated to show how to use it now. Prevents confusion in the readers about its usage
Code ReviewThis PR is a solid improvement — eliminating the auxiliary thread-locals ( Stale documentation (required fix)The docstring for
This is now wrong on two counts:
The embedded example still shows the old thread-local pattern ( let mut count = 0u32;
Trace::root(std::future::poll_fn(|cx| {
trace_with(|| { let _ = fut.as_mut().poll(cx); }, |_meta| count += 1);
Poll::Ready(())
})).await;
assert!(count > 0);Safety analysis of the
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tokio/src/runtime/task/trace/mod.rs`:
- Around line 212-215: Update the rustdoc for trace_with to reflect the new
callback model: replace references to a stand-alone fn callback and any
thread-local communication workaround with a description that trace_with accepts
a closure (FT: FnMut(&TraceMeta)) which is invoked directly with a &TraceMeta
for the current trace, and that callers should use the provided TraceMeta
argument rather than thread-local state; ensure the doc comments around the
trace_with function mention the exact signature (trace_with and TraceMeta) and
explain the intended usage and ownership/borrowing semantics of the &TraceMeta
parameter.
- Around line 147-150: The current is_tracing() calls
try_with_current_trace_leaf_fn which clears the trace slot when invoking a
callback, causing is_tracing() to return false when called from inside an active
leaf callback; update is_tracing() so it checks presence without
consuming/clearing the slot (e.g., call a non-consuming accessor such as a
ref-based helper like try_with_current_trace_leaf_ref_fn or add a new
try_peek_current_trace_leaf() that returns Some(()) without clearing) so
Handle::is_tracing() correctly returns true during active callbacks; keep the
public is_tracing() signature but replace the call to
try_with_current_trace_leaf_fn with the non-consuming check.
In `@tokio/src/runtime/task/trace/trace_impl.rs`:
- Around line 27-42: The code currently always calls
trace.push_backtrace(frames) even when TraceMeta::root_addr is None, producing
an empty backtrace; fix by only recording the backtrace when meta.root_addr is
present — move the trace.push_backtrace(frames) call into the if let
Some(root_addr) = meta.root_addr { ... } block (the block that uses
backtrace::trace and builds frames based on meta.trace_leaf_addr and ptr::eq),
so you only push frames when a root_addr was used to collect them.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 27c934f4-c4c8-4246-b004-0e386d8e462e
📒 Files selected for processing (3)
tokio/src/runtime/task/trace/mod.rstokio/src/runtime/task/trace/trace_impl.rstokio/tests/task_trace_self.rs
| /// Produces `true` if the current task is being traced; otherwise false. | ||
| pub(crate) fn is_tracing() -> bool { | ||
| Self::try_with_current_trace_leaf_fn(|maybe_trace_leaf| maybe_trace_leaf.get().is_some()) | ||
| .unwrap_or(false) | ||
| Self::try_with_current_trace_leaf_fn(|_| ()).is_some() | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n -C3 'Context::is_tracing\s*\(' --type rustRepository: martin-augment/tokio
Length of output: 468
🏁 Script executed:
#!/bin/bash
# Search for all calls to Handle::is_tracing()
rg -n 'Handle::is_tracing\s*\(\)' --type rustRepository: martin-augment/tokio
Length of output: 125
🏁 Script executed:
#!/bin/bash
# Search for all references to is_tracing in the codebase
rg -n 'is_tracing' --type rust | head -30Repository: martin-augment/tokio
Length of output: 362
🏁 Script executed:
#!/bin/bash
# Examine the try_with_current_trace_leaf_fn implementation to confirm slot clearing
rg -n -A 10 'try_with_current_trace_leaf_fn' tokio/src/runtime/task/trace/mod.rs | head -40Repository: martin-augment/tokio
Length of output: 1409
🏁 Script executed:
#!/bin/bash
# Check the context of the is_tracing() call in the test
sed -n '100,130p' tokio/tests/dump.rsRepository: martin-augment/tokio
Length of output: 835
🏁 Script executed:
#!/bin/bash
# Confirm the trace_with implementation timing - when is the callback invoked relative to task polling
cat -n tokio/src/runtime/task/trace/mod.rs | sed -n '290,330p'Repository: martin-augment/tokio
Length of output: 1799
Audit Context::is_tracing() behavior with active callbacks.
The slot-clearing in try_with_current_trace_leaf_fn means is_tracing() now returns false when called from within an active leaf callback. The test at tokio/tests/dump.rs:116 (complete_during_trace()) relies on Handle::is_tracing() returning true during tracing to avoid deadlocks; this test will now fail to detect active tracing and futures will remain pending instead of completing as designed.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tokio/src/runtime/task/trace/mod.rs` around lines 147 - 150, The current
is_tracing() calls try_with_current_trace_leaf_fn which clears the trace slot
when invoking a callback, causing is_tracing() to return false when called from
inside an active leaf callback; update is_tracing() so it checks presence
without consuming/clearing the slot (e.g., call a non-consuming accessor such as
a ref-based helper like try_with_current_trace_leaf_ref_fn or add a new
try_peek_current_trace_leaf() that returns Some(()) without clearing) so
Handle::is_tracing() correctly returns true during active callbacks; keep the
public is_tracing() signature but replace the call to
try_with_current_trace_leaf_fn with the non-consuming check.
| pub fn trace_with<FN, FT, R>(f: FN, mut trace_leaf: FT) -> R | ||
| where | ||
| F: FnOnce() -> R, | ||
| FN: FnOnce() -> R, | ||
| FT: FnMut(&TraceMeta), |
There was a problem hiding this comment.
Update the rustdoc for the new callback model.
The docs above trace_with still describe a fn callback plus thread-local communication. After this signature change, the public API docs are teaching the workaround this PR is removing.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tokio/src/runtime/task/trace/mod.rs` around lines 212 - 215, Update the
rustdoc for trace_with to reflect the new callback model: replace references to
a stand-alone fn callback and any thread-local communication workaround with a
description that trace_with accepts a closure (FT: FnMut(&TraceMeta)) which is
invoked directly with a &TraceMeta for the current trace, and that callers
should use the provided TraceMeta argument rather than thread-local state;
ensure the doc comments around the trace_with function mention the exact
signature (trace_with and TraceMeta) and explain the intended usage and
ownership/borrowing semantics of the &TraceMeta parameter.
There was a problem hiding this comment.
value:useful; category:documentation; feedback: The CodeRabbit AI reviewer is correct! The way the function works has changed - it no more needs the two extra thread locals and its signature changed a bit. Its documentation should be updated to show how to use it now. Prevents confusion in the readers about its usage
| if let Some(root_addr) = meta.root_addr { | ||
| backtrace::trace(|frame| { | ||
| let below_root = !ptr::eq(frame.symbol_address(), root_addr); | ||
|
|
||
| if ptr::eq(frame.symbol_address(), meta.trace_leaf_addr) { | ||
| above_leaf = true; | ||
| } | ||
| if above_leaf && below_root { | ||
| frames.push(frame.to_owned().into()); | ||
| } | ||
|
|
||
| below_root | ||
| }); | ||
| if ptr::eq(frame.symbol_address(), meta.trace_leaf_addr) { | ||
| above_leaf = true; | ||
| } | ||
| collector.push_backtrace(frames); | ||
| state.collector.set(Some(collector)); | ||
| } | ||
| }); | ||
|
|
||
| below_root | ||
| }); | ||
| } | ||
| trace.push_backtrace(frames); |
There was a problem hiding this comment.
Skip recording a trace when root_addr is missing.
TraceMeta::root_addr is optional, but this version always appends frames afterward. That turns “no trace captured” into “one empty backtrace captured”, which can surface phantom entries for callers that reach trace_leaf outside a Trace::root boundary.
Proposed fix
pub(crate) fn trace_leaf(meta: &TraceMeta, trace: &mut Trace) {
let mut frames: Vec<backtrace::BacktraceFrame> = vec![];
let mut above_leaf = false;
if let Some(root_addr) = meta.root_addr {
backtrace::trace(|frame| {
let below_root = !ptr::eq(frame.symbol_address(), root_addr);
if above_leaf && below_root {
frames.push(frame.to_owned().into());
}
if ptr::eq(frame.symbol_address(), meta.trace_leaf_addr) {
above_leaf = true;
}
below_root
});
+
+ trace.push_backtrace(frames);
}
- trace.push_backtrace(frames);
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if let Some(root_addr) = meta.root_addr { | |
| backtrace::trace(|frame| { | |
| let below_root = !ptr::eq(frame.symbol_address(), root_addr); | |
| if ptr::eq(frame.symbol_address(), meta.trace_leaf_addr) { | |
| above_leaf = true; | |
| } | |
| if above_leaf && below_root { | |
| frames.push(frame.to_owned().into()); | |
| } | |
| below_root | |
| }); | |
| if ptr::eq(frame.symbol_address(), meta.trace_leaf_addr) { | |
| above_leaf = true; | |
| } | |
| collector.push_backtrace(frames); | |
| state.collector.set(Some(collector)); | |
| } | |
| }); | |
| below_root | |
| }); | |
| } | |
| trace.push_backtrace(frames); | |
| pub(crate) fn trace_leaf(meta: &TraceMeta, trace: &mut Trace) { | |
| let mut frames: Vec<backtrace::BacktraceFrame> = vec![]; | |
| let mut above_leaf = false; | |
| if let Some(root_addr) = meta.root_addr { | |
| backtrace::trace(|frame| { | |
| let below_root = !ptr::eq(frame.symbol_address(), root_addr); | |
| if above_leaf && below_root { | |
| frames.push(frame.to_owned().into()); | |
| } | |
| if ptr::eq(frame.symbol_address(), meta.trace_leaf_addr) { | |
| above_leaf = true; | |
| } | |
| below_root | |
| }); | |
| trace.push_backtrace(frames); | |
| } | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tokio/src/runtime/task/trace/trace_impl.rs` around lines 27 - 42, The code
currently always calls trace.push_backtrace(frames) even when
TraceMeta::root_addr is None, producing an empty backtrace; fix by only
recording the backtrace when meta.root_addr is present — move the
trace.push_backtrace(frames) call into the if let Some(root_addr) =
meta.root_addr { ... } block (the block that uses backtrace::trace and builds
frames based on meta.trace_leaf_addr and ptr::eq), so you only push frames when
a root_addr was used to collect them.
There was a problem hiding this comment.
value:useful; category:bug; feedback: The CodeRabbit AI reviewer is correct! There is no need to push the empty frames if there is no root_addr. Prevents useless allocation of the frames Vec
🤖 Augment PR SummarySummary: This PR expands taskdump tracing to accept Changes:
Technical Notes: The new callback storage relies on temporarily extending the callback’s lifetime while it is installed in Tokio’s thread-local trace context, and on clearing the slot while invoking the callback to enforce exclusivity. 🤖 Was this summary useful? React with 👍 or 👎 |
![augmentcode[bot]](https://avatars.githubusercontent.com/in/1027498?s=60&v=4){kind=small}
| /// # } | ||
| /// ``` | ||
| pub fn trace_with<F, R>(f: F, trace_leaf: fn(&TraceMeta)) -> R | ||
| pub fn trace_with<FN, FT, R>(f: FN, mut trace_leaf: FT) -> R |
There was a problem hiding this comment.
The doc comment above trace_with still says the trace-leaf callback must be a fn stored in TLS via a Cell, but trace_with now accepts FnMut closures stored via a raw pointer; consider updating the docs/example to reflect the new API and its safety invariants.
Severity: low
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
There was a problem hiding this comment.
value:useful; category:documentation; feedback: The Augment AI reviewer is correct! The way the function works has changed - it no more needs the two extra thread locals and its signature changed a bit. Its documentation should be updated to show how to use it now. Prevents confusion in the readers about its usage
value:useful; category:documentation; feedback: The Claude AI reviewer is correct! The way the function works has changed - it no more needs the two extra thread locals and its signature changed a bit. Its documentation should be updated to show how to use it now. Prevents confusion in the readers about its usage |
2 participants
8040: To review by AI