support deriving root span from key

rust-toolchain.toml

@@ -0,0 +1,3 @@
+[toolchain]
+channel = "1.84"
+profile = "default"

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 {

src/registry.rs

@@ -52,6 +52,13 @@ impl<T> Key for T where T: Hash + Eq + Debug + Send + Sync + 'static {}
 trait ObjKey: DynHash + DynEq + Debug + Send + Sync + 'static {}
 impl<T> 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.
     ///

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<T>(key: impl Key + ToRootSpan, future: T) -> JoinHandle<T::Output>
+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.
 ///