improve usability of registry (#12)

benches/basic.rs

@@ -69,13 +69,15 @@ async fn test_baseline() {
 }
 
 async fn spawn_many(size: usize) {
-    let mut root = Registry::new(Config::default());
+    let registry = Registry::new(Config::default());
     let mut handles = vec![];
     for i in 0..size {
         let task = async {
             tokio::time::sleep(Duration::from_millis(10)).await;
         };
-        handles.push(tokio::spawn(root.register(i, "new_task").instrument(task)));
+        handles.push(tokio::spawn(
+            registry.register(i, "new_task").instrument(task),
+        ));
     }
     futures::future::try_join_all(handles)
         .await
@@ -102,9 +104,9 @@ fn bench_basic(c: &mut Criterion) {
     c.bench_function("basic", |b| {
         b.to_async(runtime()).iter(|| async {
             let config = ConfigBuilder::default().verbose(false).build().unwrap();
-            let mut mgr = Registry::new(config);
+            let registry = Registry::new(config);
 
-            let root = mgr.register(233, "root");
+            let root = registry.register(233, "root");
             root.instrument(test()).await;
         })
     });
@@ -114,9 +116,9 @@ fn bench_basic_baseline(c: &mut Criterion) {
     c.bench_function("basic_baseline", |b| {
         b.to_async(runtime()).iter(|| async {
             let config = ConfigBuilder::default().verbose(false).build().unwrap();
-            let mut mgr = Registry::new(config);
+            let registry = Registry::new(config);
 
-            let root = mgr.register(233, "root");
+            let root = registry.register(233, "root");
             black_box(root);
             test_baseline().await
         })

examples/basic.rs

@@ -43,12 +43,12 @@ async fn foo() {
 
 #[tokio::main]
 async fn main() {
-    let mut registry = Registry::new(Config::default());
+    let registry = Registry::new(Config::default());
     let root = registry.register((), "foo");
     tokio::spawn(root.instrument(foo()));
 
     sleep(Duration::from_secs(1)).await;
-    let tree = registry.get(&()).unwrap().to_string();
+    let tree = registry.get(()).unwrap().to_string();
 
     // foo [1.006s]
     //   bar [1.006s]

examples/detach.rs

@@ -47,7 +47,7 @@ async fn work(rx: Receiver<()>) {
 
 #[tokio::main]
 async fn main() {
-    let mut registry = Registry::new(Config::default());
+    let registry = Registry::new(Config::default());
     let root = registry.register((), "work");
     let (tx, rx) = oneshot::channel();
     tokio::spawn(root.instrument(work(rx)));

examples/multiple.rs

@@ -31,8 +31,8 @@ async fn foo() {
 
 #[tokio::main]
 async fn main() {
-    let mut registry = Registry::new(Config::default());
-    for i in 0..3 {
+    let registry = Registry::new(Config::default());
+    for i in 0_i32..3 {
         let root = registry.register(i, format!("actor {i}"));
         tokio::spawn(root.instrument(work(i)));
     }
@@ -50,7 +50,11 @@ async fn main() {
     // actor 2 [1.007s]
     //   actor work 2 [1.007s]
     //     pending [1.007s]
-    for (_, tree) in registry.iter().sorted_by_key(|(i, _)| *i) {
+    for (_, tree) in registry
+        .collect::<i32>()
+        .into_iter()
+        .sorted_by_key(|(i, _)| *i)
+    {
         println!("{tree}");
     }
 }

examples/verbose.rs

@@ -27,7 +27,7 @@ async fn foo() {
 
 async fn work(verbose: bool) -> String {
     let config = ConfigBuilder::default().verbose(verbose).build().unwrap();
-    let mut registry = Registry::new(config);
+    let registry = Registry::new(config);
     let root = registry.register((), "foo");
     tokio::spawn(root.instrument(foo()));
 

src/lib.rs

@@ -12,16 +12,21 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+//! Instrument await-tree for actor-based applications.
+
+#![forbid(missing_docs)]
+
 use std::future::Future;
 
 mod context;
 mod future;
+mod obj_utils;
 mod registry;
 
 pub use context::{current_tree, TreeContext};
 use flexstr::SharedStr;
 pub use future::Instrumented;
-pub use registry::{Config, ConfigBuilder, ConfigBuilderError, Registry, TreeRoot};
+pub use registry::{AnyKey, Config, ConfigBuilder, ConfigBuilderError, Key, Registry, TreeRoot};
 
 /// A cheaply cloneable span in the await-tree.
 #[derive(Debug, Clone, Default, PartialEq, Eq, PartialOrd, Ord)]

src/obj_utils.rs

@@ -0,0 +1,74 @@
+// Copyright 2023 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Utilities for using `Any` as the key of a `HashMap`.
+
+// Adopted from:
+// https://github.com/bevyengine/bevy/blob/56bcbb097552b45e3ff48c48947ed8ee4e2c24b1/crates/bevy_utils/src/label.rs
+
+use std::any::Any;
+use std::hash::{Hash, Hasher};
+
+/// An object safe version of [`Eq`]. This trait is automatically implemented
+/// for any `'static` type that implements `Eq`.
+pub(crate) trait DynEq: Any {
+    /// Casts the type to `dyn Any`.
+    fn as_any(&self) -> &dyn Any;
+
+    /// This method tests for `self` and `other` values to be equal.
+    ///
+    /// Implementers should avoid returning `true` when the underlying types are
+    /// not the same.
+    fn dyn_eq(&self, other: &dyn DynEq) -> bool;
+}
+
+impl<T> DynEq for T
+where
+    T: Any + Eq,
+{
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn dyn_eq(&self, other: &dyn DynEq) -> bool {
+        if let Some(other) = other.as_any().downcast_ref::<T>() {
+            return self == other;
+        }
+        false
+    }
+}
+
+/// An object safe version of [`Hash`]. This trait is automatically implemented
+/// for any `'static` type that implements `Hash`.
+pub(crate) trait DynHash: DynEq {
+    /// Casts the type to `dyn Any`.
+    fn as_dyn_eq(&self) -> &dyn DynEq;
+
+    /// Feeds this value into the given [`Hasher`].
+    fn dyn_hash(&self, state: &mut dyn Hasher);
+}
+
+impl<T> DynHash for T
+where
+    T: DynEq + Hash,
+{
+    fn as_dyn_eq(&self) -> &dyn DynEq {
+        self
+    }
+
+    fn dyn_hash(&self, mut state: &mut dyn Hasher) {
+        T::hash(self, &mut state);
+        self.type_id().hash(&mut state);
+    }
+}