1 files changed, 56 insertions, 17 deletions

diff --git a/incria/src/lazy_mapping.rs b/incria/src/lazy_mapping.rs
index 27bd76e..f206bdd 100644
--- a/incria/src/lazy_mapping.rs
+++ b/incria/src/lazy_mapping.rs
@@ -1,9 +1,11 @@
 use std::{
-    collections::HashMap, fmt::Debug, future::Future, hash::Hash, mem::transmute, sync::Arc,
+    any::type_name, collections::HashMap, fmt::Debug, future::Future, hash::Hash, mem::transmute,
+    sync::Arc,
 };
 
 use blink_alloc::{Blink, SyncBlinkAlloc};
 use tokio::sync::{Mutex, Notify, RwLock, RwLockReadGuard};
+use tracing::{debug_span, instrument, trace_span, Instrument};
 
 use crate::{
     deps::{self, NodeId},
@@ -56,7 +58,7 @@ impl<K: Clone + Eq + Hash + Send + Sync, V: Send + Sync, B: LazyMappingBackend<K
 
 impl<K, V, B> Mapper for LazyMapper<K, V, B>
 where
-    K: Clone + Eq + Hash + Send + Sync + Debug,
+    K: Clone + Eq + Hash + Send + Sync,
     V: Send + Sync + 'static,
     B: LazyMappingBackend<K, V>,
 {
@@ -64,35 +66,64 @@ where
     type Value = V;
     type Wrapper<'a> = &'a V where V: 'a;
 
+    #[instrument(level = "debug", skip_all, fields(key_type = type_name::<K>(), value_type = type_name::<V>()))]
     async fn get<'a>(&'a self, key: &Self::Key) -> Self::Wrapper<'a> {
         // Attempt to reuse or evict the existing value
         let mut reuse_dep_id = None;
         {
-            let finished = self.calculated.read().await;
-            if let Some((val, node)) = finished.get(key) {
-                let dirty = deps::is_dirty(*node) || self.backend.is_dirty(key).await;
+            let finished = self
+                .calculated
+                .read()
+                .instrument(trace_span!("calculated.read"))
+                .await;
+            if let Some((val, node)) = finished.get(key).copied() {
+                drop(finished);
+
+                let dirty = deps::is_dirty(node)
+                    || self
+                        .backend
+                        .is_dirty(key)
+                        .instrument(trace_span!("backend.is_dirty"))
+                        .await;
                 if !dirty {
-                    deps::add_dep(*node);
+                    deps::add_dep(node);
                     return val;
                 } else {
-                    reuse_dep_id = Some(*node);
-                    drop(finished);
+                    reuse_dep_id = Some(node);
                     // Dirty, so we'll recompute below but we should remove it now
-                    if self.calculated.write().await.remove(key).is_none() {
+                    let removed = async { self.calculated.write().await.remove(key).is_none() }
+                        .instrument(trace_span!("remove dirty value"))
+                        .await;
+                    if removed {
                         // Someone else already noticed it was dirty and removed it before us, so we need to deal with that
                         todo!("dirty value removed between us noticing and us doing something")
                     }
                 }
             }
-        }
+        };
+
+        let barrier = self
+            .waiting
+            .lock()
+            .instrument(trace_span!("waiting.lock"))
+            .await
+            .get(key)
+            .cloned();
 
-        let barrier = self.waiting.lock().await.get(key).cloned();
         if let Some(barrier) = barrier {
             // Waiting for completion
-            barrier.notified().await;
-
-            let (value, node_id) =
-                *RwLockReadGuard::map(self.calculated.read().await, |hm| hm.get(key).unwrap());
+            barrier
+                .notified()
+                .instrument(trace_span!("completion_wait"))
+                .await;
+
+            let (value, node_id) = *RwLockReadGuard::map(
+                self.calculated
+                    .read()
+                    .instrument(trace_span!("calculated.read"))
+                    .await,
+                |hm| hm.get(key).unwrap(),
+            );
 
             deps::add_dep(node_id);
 
@@ -102,6 +133,7 @@ where
             let notify = Arc::new(Notify::new());
             self.waiting
                 .lock()
+                .instrument(trace_span!("waiting.insert"))
                 .await
                 .insert(key.clone(), notify.clone());
 
@@ -112,15 +144,22 @@ where
                 deps::next_node_id()
             };
 
-            let val = deps::with_node_id(dep, self.backend.compute(key.clone())).await;
+            let val = deps::with_node_id(dep, self.backend.compute(key.clone()))
+                .instrument(debug_span!("backend.compute"))
+                .await;
             let val_ref: &'static V = unsafe { transmute(Blink::new_in(&self.values).put(val)) };
             deps::add_dep(dep);
 
             self.calculated
                 .write()
+                .instrument(trace_span!("calculated.insert"))
                 .await
                 .insert(key.clone(), (val_ref, dep));
-            self.waiting.lock().await.remove(key);
+            self.waiting
+                .lock()
+                .instrument(trace_span!("waiting.remove"))
+                .await
+                .remove(key);
 
             notify.notify_waiters();