diff --git a/rust-toolchain b/rust-toolchain deleted file mode 100644 index 74c280f..0000000 --- a/rust-toolchain +++ /dev/null @@ -1 +0,0 @@ -1.83 diff --git a/rust-toolchain.toml b/rust-toolchain.toml new file mode 100644 index 0000000..88d2c65 --- /dev/null +++ b/rust-toolchain.toml @@ -0,0 +1,3 @@ +[toolchain] +channel = "1.84" +profile = "default" diff --git a/src/lib.rs b/src/lib.rs index b221dff..f59e03c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -84,11 +84,11 @@ mod spawn; pub use context::{current_tree, Tree}; pub use future::Instrumented; pub use global::init_global_registry; -pub use registry::{AnyKey, Config, ConfigBuilder, ConfigBuilderError, Key, Registry}; +pub use registry::{AnyKey, Config, ConfigBuilder, ConfigBuilderError, Key, Registry, ToRootSpan}; pub use root::TreeRoot; pub use span::{Span, SpanExt}; #[cfg(feature = "tokio")] -pub use spawn::{spawn, spawn_anonymous}; +pub use spawn::{spawn, spawn_anonymous, spawn_root}; #[doc(hidden)] pub mod __private { diff --git a/src/registry.rs b/src/registry.rs index e103670..b90e819 100644 --- a/src/registry.rs +++ b/src/registry.rs @@ -52,6 +52,13 @@ impl Key for T where T: Hash + Eq + Debug + Send + Sync + 'static {} trait ObjKey: DynHash + DynEq + Debug + Send + Sync + 'static {} impl ObjKey for T where T: DynHash + DynEq + Debug + Send + Sync + 'static {} +/// A trait for types that can be converted to a [`Span`] that can be used as the root of an +/// await-tree. +pub trait ToRootSpan { + /// Convert the type to a [`Span`] that can be used as the root of an await-tree. + fn to_root_span(&self) -> Span; +} + /// Key type for anonymous await-trees. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] struct AnonymousKey(ContextId); @@ -211,7 +218,8 @@ impl Registry { } } - /// Register with given key. Returns a [`TreeRoot`] that can be used to instrument a future. + /// Register with given key and root span. Returns a [`TreeRoot`] that can be used to instrument + /// a future. /// /// If the key already exists, a new [`TreeRoot`] is returned and the reference to the old /// [`TreeRoot`] is dropped. @@ -220,6 +228,15 @@ impl Registry { self.register_inner(key, context) } + /// Derive the root span from the given key and register with it. + /// + /// This is a convenience method for `self.register(key, key.to_root_span())`. See + /// [`Registry::register`] for more details. + pub fn register_root(&self, key: impl Key + ToRootSpan) -> TreeRoot { + let root_span = key.to_root_span(); + self.register(key, root_span) + } + /// Register an anonymous await-tree without specifying a key. Returns a [`TreeRoot`] that can /// be used to instrument a future. /// diff --git a/src/spawn.rs b/src/spawn.rs index a4ae42c..5b17a78 100644 --- a/src/spawn.rs +++ b/src/spawn.rs @@ -19,7 +19,7 @@ use std::future::Future; use tokio::task::JoinHandle; -use crate::{Key, Registry, Span}; +use crate::{Key, Registry, Span, ToRootSpan}; /// Spawns a new asynchronous task instrumented with the given root [`Span`], returning a /// [`JoinHandle`] for it. @@ -39,6 +39,20 @@ where } } +/// Spawns a new asynchronous task instrumented with the root span derived from the given key, +/// returning a [`JoinHandle`] for it. +/// +/// This is a convenience function for `spawn(key, key.to_root_span(), future)`. See [`spawn`] for +/// more details. +pub fn spawn_root(key: impl Key + ToRootSpan, future: T) -> JoinHandle +where + T: Future + Send + 'static, + T::Output: Send + 'static, +{ + let root_span = key.to_root_span(); + spawn(key, root_span, future) +} + /// Spawns a new asynchronous task instrumented with the given root [`Span`], returning a /// [`JoinHandle`] for it. ///