Skip to content

feat: instrumented joinsets and other custom spawns#392

Merged
rcoh merged 4 commits into
dial9-rs:mainfrom
noxware:feat/301-instrumented-custom-spawn
May 12, 2026
Merged

feat: instrumented joinsets and other custom spawns#392
rcoh merged 4 commits into
dial9-rs:mainfrom
noxware:feat/301-instrumented-custom-spawn

Conversation

@noxware

@noxware noxware commented May 11, 2026

Copy link
Copy Markdown
Contributor

Current solution

Exposes a handle.spawn_with(fut, |f: TracedFuture| { ... }) to allow spawning instrumented/traced futures through custom spawn APIs (like JoinSet::spawn).

Example

let handle = TelemetryHandle::current();
let mut set: JoinSet<()> = JoinSet::new();
handle.spawn_with(fut, |f| set.spawn(f));

For reviewers

Most important changes are at recorder/mod.rs and traced.rs.

Other

Closes #301

@noxware noxware left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added some comments to help with the code review :)

/// ```
///
/// [`TracedFuture<F>`]: crate::telemetry::TracedFuture
pub fn spawn_with<F, S>(

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the new API, the root part of the PR.

Comment thread dial9-tokio-telemetry/src/traced.rs Outdated

pin_project! {
/// Future wrapper that captures wake events for a known Tokio task.
pub(crate) struct WakeTracked<F> {

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is what used to be Traced (kinda).

/// and uses it for task instrumentation. If the future is polled outside a
/// Tokio task context, it runs as a transparent passthrough without wake
/// tracking or task dumps.
pub struct TracedFuture<F> {

@noxware noxware May 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the concreate type of the spawn_with closure. It composes WakeTracked and TaskDumped (if enabled), without dynamic dispatching.

Note: It can also be called Traced<F> as the name is now free. Although we also have other explicit "traced" types like TracedRuntime. Opinons?

}

pin_project! {
#[project = TracedFutureStateProj]

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

An explicit name is needed to pattern match the enum.

#[pin]
inner: InstrumentedFuture<F>,
},
Empty,

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Empty is just used to transit the state safely.

@noxware noxware May 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This file contains the auxiliar changes to make spawn_with work by exposing the TracedFuture type. It has the most important changes.


type InstrumentedFuture<F> = WakeTracked<MaybeTaskDumped<F>>;

pin_project! {

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We rely on pin projection to avoid writting unsafe code.

@noxware noxware marked this pull request as ready for review May 11, 2026 20:36
@noxware noxware requested review from rcoh and yulnr May 11, 2026 20:37

@yulnr yulnr left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me! happy we could keep the instrumented flag an internal concern.
Leaving some minor comments

Comment thread dial9-tokio-telemetry/src/telemetry/recorder/mod.rs Outdated
Comment thread dial9-tokio-telemetry/tests/spawn_with.rs
Comment thread dial9-tokio-telemetry/src/traced.rs Outdated
Comment thread dial9-tokio-telemetry/src/telemetry/recorder/mod.rs Outdated
@noxware

noxware commented May 12, 2026

Copy link
Copy Markdown
Contributor Author

@yulnr It should be done now :)

@rcoh rcoh left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice!

@rcoh rcoh added this pull request to the merge queue May 12, 2026
Merged via the queue into dial9-rs:main with commit 4c820c2 May 12, 2026
20 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

make Traced future public

3 participants