From e2b0cf9934c888301f8f14f3c1d236ac1ff41366 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Thu, 28 Mar 2024 15:34:32 +0800 Subject: [PATCH] allow registering anonymous trees (#13) --- src/context.rs | 23 ++++++++++---------- src/lib.rs | 2 +- src/registry.rs | 57 +++++++++++++++++++++++++++++++++++++++++++------ 3 files changed, 63 insertions(+), 19 deletions(-) diff --git a/src/context.rs b/src/context.rs index 2e628e9..d1c8efb 100644 --- a/src/context.rs +++ b/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 } diff --git a/src/lib.rs b/src/lib.rs index 39e4c93..7ea9ef0 100644 --- a/src/lib.rs +++ b/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}; diff --git a/src/registry.rs b/src/registry.rs index e5586c2..9c0b02f 100644 --- a/src/registry.rs +++ b/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::() + } } type Contexts = RwLock>>; @@ -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) -> TreeRoot { - let context = Arc::new(TreeContext::new(root_span.into(), self.config().verbose)); + fn register_inner(&self, key: impl Key, context: Arc) -> 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) -> 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) -> 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 { + 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::(); 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); } }