allow registering anonymous trees (#13)

src/context.rs

@@ -44,7 +44,11 @@ impl SpanNode {
 
 /// The id of an await-tree context. We will check the id recorded in the instrumented future
 /// against the current task-local context before trying to update the tree.
-pub(crate) type ContextId = u64;
+
+// Also used as the key for anonymous trees in the registry.
+// Intentionally made private to prevent users from reusing the same id when registering a new tree.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub(crate) struct ContextId(u64);
 
 /// An await-tree for a task.
 #[derive(Debug, Clone)]
@@ -201,7 +205,7 @@ impl Tree {
 
 /// The task-local await-tree context.
 #[derive(Debug)]
-pub struct TreeContext {
+pub(crate) struct TreeContext {
     /// The id of the context.
     id: ContextId,
 
@@ -222,7 +226,7 @@ impl TreeContext {
         let root = arena.new_node(SpanNode::new(root_span));
 
         Self {
-            id,
+            id: ContextId(id),
             verbose,
             tree: Tree {
                 arena,
@@ -233,6 +237,11 @@ impl TreeContext {
         }
     }
 
+    /// Get the context id.
+    pub(crate) fn id(&self) -> ContextId {
+        self.id
+    }
+
     /// Returns the locked guard of the tree.
     pub(crate) fn tree(&self) -> MutexGuard<'_, Tree> {
         self.tree.lock()
@@ -244,14 +253,6 @@ impl TreeContext {
     }
 }
 
-/// Public interfaces.
-impl TreeContext {
-    /// Get the context id.
-    pub fn id(&self) -> ContextId {
-        self.id
-    }
-}
-
 tokio::task_local! {
     pub(crate) static CONTEXT: Arc<TreeContext>
 }

src/lib.rs

@@ -23,7 +23,7 @@ mod future;
 mod obj_utils;
 mod registry;
 
-pub use context::{current_tree, TreeContext};
+pub use context::current_tree;
 use flexstr::SharedStr;
 pub use future::Instrumented;
 pub use registry::{AnyKey, Config, ConfigBuilder, ConfigBuilderError, Key, Registry, TreeRoot};

src/registry.rs

@@ -22,7 +22,7 @@ use derive_builder::Builder;
 use parking_lot::RwLock;
 use weak_table::WeakValueHashMap;
 
-use crate::context::{Tree, TreeContext, CONTEXT};
+use crate::context::{ContextId, Tree, TreeContext, CONTEXT};
 use crate::obj_utils::{DynEq, DynHash};
 use crate::Span;
 
@@ -94,6 +94,11 @@ impl AnyKey {
     pub fn as_any(&self) -> &dyn Any {
         self.0.as_ref().as_any()
     }
+
+    /// Returns whether this key corresponds to an anonymous await-tree.
+    pub fn is_anonymous(&self) -> bool {
+        self.as_any().is::<ContextId>()
+    }
 }
 
 type Contexts = RwLock<WeakValueHashMap<AnyKey, Weak<TreeContext>>>;
@@ -138,12 +143,7 @@ impl Registry {
         )
     }
 
-    /// Register with given key. 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.
-    pub fn register(&self, key: impl Key, root_span: impl Into<Span>) -> TreeRoot {
-        let context = Arc::new(TreeContext::new(root_span.into(), self.config().verbose));
+    fn register_inner(&self, key: impl Key, context: Arc<TreeContext>) -> TreeRoot {
         self.contexts()
             .write()
             .insert(AnyKey::new(key), Arc::clone(&context));
@@ -154,6 +154,25 @@ impl Registry {
         }
     }
 
+    /// Register with given key. 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.
+    pub fn register(&self, key: impl Key, root_span: impl Into<Span>) -> TreeRoot {
+        let context = Arc::new(TreeContext::new(root_span.into(), self.config().verbose));
+        self.register_inner(key, context)
+    }
+
+    /// Register an anonymous await-tree without specifying a key. Returns a [`TreeRoot`] that can
+    /// be used to instrument a future.
+    ///
+    /// Anonymous await-trees are not able to be retrieved through the [`Registry::get`] method. Use
+    /// [`Registry::collect_anonymous`] or [`Registry::collect_all`] to collect them.
+    pub fn register_anonymous(&self, root_span: impl Into<Span>) -> TreeRoot {
+        let context = Arc::new(TreeContext::new(root_span.into(), self.config().verbose));
+        self.register_inner(context.id(), context) // use the private id as the key
+    }
+
     /// Get a clone of the await-tree with given key.
     ///
     /// Returns `None` if the key does not exist or the tree root has been dropped.
@@ -183,6 +202,21 @@ impl Registry {
             .collect()
     }
 
+    /// Collect the snapshots of all await-trees registered with [`Registry::register_anonymous`].
+    pub fn collect_anonymous(&self) -> Vec<Tree> {
+        self.contexts()
+            .read()
+            .iter()
+            .filter_map(|(k, v)| {
+                if k.is_anonymous() {
+                    Some(v.tree().clone())
+                } else {
+                    None
+                }
+            })
+            .collect()
+    }
+
     /// Collect the snapshots of all await-trees regardless of the key type.
     pub fn collect_all(&self) -> Vec<(AnyKey, Tree)> {
         self.contexts()
@@ -211,6 +245,9 @@ mod tests {
         let _unit = registry.register((), "()");
         let _unit_replaced = registry.register((), "[]");
 
+        let _anon = registry.register_anonymous("anon");
+        let _anon = registry.register_anonymous("anon");
+
         let i32s = registry.collect::<i32>();
         assert_eq!(i32s.len(), 3);
 
@@ -219,5 +256,11 @@ mod tests {
 
         let units = registry.collect::<()>();
         assert_eq!(units.len(), 1);
+
+        let anons = registry.collect_anonymous();
+        assert_eq!(anons.len(), 2);
+
+        let all = registry.collect_all();
+        assert_eq!(all.len(), 8);
     }
 }