fix(query): avoid cyclic package docs deadlock

Author: Myriad-Dreamin

crates/tinymist-query/src/analysis/global.rs

@@ -8,7 +8,7 @@ use std::{collections::HashSet, ops::Deref};
 use comemo::{Track, Tracked};
 use ecow::EcoString;
 use lsp_types::Url;
-use parking_lot::Mutex;
+use parking_lot::{Condvar, Mutex};
 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
 use rustc_hash::FxHashMap;
 use tinymist_analysis::docs::DocString;
@@ -821,15 +821,26 @@ impl SharedContext {
         }
     }
 
-    /// Gets the expression information of a source file.
-    pub(crate) fn expr_stage_by_id(self: &Arc<Self>, fid: TypstFileId) -> Option<ExprInfo> {
-        Some(self.expr_stage(&self.source_by_id(fid).ok()?))
+    /// Gets expression information for a recursive dependency without waiting
+    /// on another worker's cache initializer.
+    pub(crate) fn expr_stage_recursive_by_id(
+        self: &Arc<Self>,
+        fid: TypstFileId,
+    ) -> Option<ExprInfo> {
+        let source = self.source_by_id(fid).ok()?;
+        let mut route = ExprRoute::default();
+        Some(self.expr_stage_(&source, &mut route))
     }
 
     /// Gets the expression information of a source file.
     pub(crate) fn expr_stage(self: &Arc<Self>, source: &Source) -> ExprInfo {
         let mut route = ExprRoute::default();
-        self.expr_stage_(source, &mut route)
+        use crate::syntax::expr_of;
+
+        let guard = self.query_stat(source.id(), "expr_stage");
+        self.slot.expr_stage.compute(hash128(&source), |prev| {
+            expr_of(self.clone(), source.clone(), &mut route, guard, prev)
+        })
     }
 
     /// Gets the expression information of a source file.
@@ -839,17 +850,30 @@ impl SharedContext {
         route: &mut ExprRoute,
     ) -> ExprInfo {
         use crate::syntax::expr_of;
+
+        // Recursive expression traversals must not wait for another worker's
+        // initializer; their route-local result is intentionally unpublished.
+        if let Some(cached) = route.completed(&source.id()) {
+            return cached.clone();
+        }
+
+        if let Some(cached) = self.slot.expr_stage.get(&hash128(&source)) {
+            return cached;
+        }
+
         let guard = self.query_stat(source.id(), "expr_stage");
-        self.slot.expr_stage.compute(hash128(&source), |prev| {
-            expr_of(self.clone(), source.clone(), route, guard, prev)
-        })
+        expr_of(self.clone(), source.clone(), route, guard, None)
     }
 
     pub(crate) fn exports_of(
         self: &Arc<Self>,
         source: &Source,
         route: &mut ExprRoute,
     ) -> Option<Arc<LazyHash<LexicalScope>>> {
+        if let Some(ei) = route.completed(&source.id()) {
+            return Some(ei.exports.clone());
+        }
+
         if let Some(s) = route.get(&source.id()) {
             return s.clone();
         }
@@ -859,18 +883,9 @@ impl SharedContext {
 
     /// Gets the type check information of a source file.
     pub(crate) fn type_check(self: &Arc<Self>, source: &Source) -> Arc<TypeInfo> {
-        let mut route = TypeEnv::default();
-        self.type_check_(source, &mut route)
-    }
-
-    /// Gets the type check information of a source file.
-    pub(crate) fn type_check_(
-        self: &Arc<Self>,
-        source: &Source,
-        route: &mut TypeEnv,
-    ) -> Arc<TypeInfo> {
         use crate::analysis::type_check;
 
+        let mut route = TypeEnv::default();
         let ei = self.expr_stage(source);
         let guard = self.query_stat(source.id(), "type_check");
         self.slot.type_check.compute(hash128(&ei), |prev| {
@@ -880,10 +895,40 @@ impl SharedContext {
             }
 
             guard.miss();
-            type_check(self.clone(), ei, route)
+            type_check(self.clone(), ei, &mut route)
         })
     }
 
+    /// Gets type check information for a recursive dependency.
+    ///
+    /// A dependency already running on another worker is recomputed in this
+    /// traversal instead of waited on, because that worker may be waiting on
+    /// this traversal. Only the owner of the shared cache slot publishes its
+    /// result.
+    pub(crate) fn type_check_recursive(
+        self: &Arc<Self>,
+        ei: ExprInfo,
+        route: &mut TypeEnv,
+    ) -> Arc<TypeInfo> {
+        use crate::analysis::type_check;
+
+        if let Some(cached) = route.type_infos.get(&ei.fid) {
+            return cached.clone();
+        }
+
+        let guard = self.query_stat(ei.fid, "type_check");
+        if let Some(cached) = self.slot.type_check.get(&hash128(&ei)) {
+            route.type_infos.insert(ei.fid, cached.clone());
+            return cached;
+        }
+
+        guard.miss();
+        let fid = ei.fid;
+        let type_info = type_check(self.clone(), ei, route);
+        route.type_infos.insert(fid, type_info.clone());
+        type_info
+    }
+
     /// Gets the lint result of a source file.
     #[typst_macros::time(span = source.root().span())]
     pub(crate) fn lint(self: &Arc<Self>, source: &Source, issues: &KnownIssues) -> LintInfo {
@@ -1044,7 +1089,7 @@ impl SharedContext {
                 Some(DefDocs::Variable(docs))
             }
             DefKind::Module => {
-                let ei = self.expr_stage_by_id(def.decl.file_id()?)?;
+                let ei = self.expr_stage_recursive_by_id(def.decl.file_id()?)?;
                 Some(DefDocs::Module(TidyModuleDocs {
                     docs: ei.module_docstring.docs.clone().unwrap_or_default(),
                 }))
@@ -1249,7 +1294,7 @@ impl SharedContext {
 
         let preloader = Preloader {
             shared: self,
-            analyzed: Arc::default(),
+            analyzed: Arc::new(Mutex::new(HashSet::from([entry_point]))),
         };
 
         preloader.work(entry_point);
@@ -1258,13 +1303,88 @@ impl SharedContext {
 
 // Needed by recursive computation
 type DeferredCompute<T> = Arc<OnceLock<T>>;
+type IncrDeferredCompute<T> = Arc<DeferredSlot<T>>;
+
+#[derive(Default)]
+enum DeferredState<T> {
+    #[default]
+    Fresh,
+    Running,
+    Done(T),
+}
+
+struct DeferredSlot<T> {
+    state: Mutex<DeferredState<T>>,
+    ready: Condvar,
+}
+
+impl<T> Default for DeferredSlot<T> {
+    fn default() -> Self {
+        Self {
+            state: Mutex::new(DeferredState::Fresh),
+            ready: Condvar::new(),
+        }
+    }
+}
+
+impl<T: Clone> DeferredSlot<T> {
+    fn get(&self) -> Option<T> {
+        let state = self.state.lock();
+        match &*state {
+            DeferredState::Done(value) => Some(value.clone()),
+            DeferredState::Fresh | DeferredState::Running => None,
+        }
+    }
+
+    fn compute(&self, compute: impl FnOnce() -> T) -> T {
+        let mut compute = Some(compute);
+        loop {
+            let mut state = self.state.lock();
+            match &*state {
+                DeferredState::Done(value) => return value.clone(),
+                DeferredState::Fresh => {
+                    *state = DeferredState::Running;
+                    drop(state);
+                    return self.finish(compute.take().expect("compute closure used once"));
+                }
+                DeferredState::Running => self.ready.wait(&mut state),
+            }
+        }
+    }
+
+    fn finish(&self, compute: impl FnOnce() -> T) -> T {
+        let mut reset = ResetDeferredSlot {
+            slot: self,
+            armed: true,
+        };
+        let value = compute();
+        *self.state.lock() = DeferredState::Done(value.clone());
+        reset.armed = false;
+        self.ready.notify_all();
+        value
+    }
+}
+
+struct ResetDeferredSlot<'a, T> {
+    slot: &'a DeferredSlot<T>,
+    armed: bool,
+}
+
+impl<T> Drop for ResetDeferredSlot<'_, T> {
+    fn drop(&mut self) {
+        if self.armed {
+            *self.slot.state.lock() = DeferredState::Fresh;
+            self.slot.ready.notify_all();
+        }
+    }
+}
 
 #[derive(Clone)]
 struct IncrCacheMap<K, V> {
     revision: usize,
     global: Arc<Mutex<FxDashMap<K, (usize, V)>>>,
-    prev: Arc<Mutex<FxHashMap<K, DeferredCompute<V>>>>,
-    next: Arc<Mutex<FxHashMap<K, DeferredCompute<V>>>>,
+    prev: Arc<Mutex<FxHashMap<K, IncrDeferredCompute<V>>>>,
+    next: Arc<Mutex<FxHashMap<K, IncrDeferredCompute<V>>>>,
 }
 
 impl<K: Eq + Hash, V> Default for IncrCacheMap<K, V> {
@@ -1279,41 +1399,60 @@ impl<K: Eq + Hash, V> Default for IncrCacheMap<K, V> {
 }
 
 impl<K, V> IncrCacheMap<K, V> {
+    fn get(&self, key: &K) -> Option<V>
+    where
+        K: Eq + Hash,
+        V: Clone,
+    {
+        // Only reuse a completed slot from this revision. Global values are
+        // consumed by the blocking owner so recursive routes stay isolated.
+        let next = self.next.lock().get(key).cloned();
+        next.and_then(|slot| slot.get())
+    }
+
     fn compute(&self, key: K, compute: impl FnOnce(Option<V>) -> V) -> V
     where
         K: Clone + Eq + Hash,
         V: Clone,
     {
         let next = self.next.lock().entry(key.clone()).or_default().clone();
 
-        next.get_or_init(|| {
-            let prev = self.prev.lock().get(&key).cloned();
-            let prev = prev.and_then(|prev| prev.get().cloned());
-            let prev = prev.or_else(|| {
-                let global = self.global.lock();
-                global.get(&key).map(|global| global.1.clone())
-            });
-
-            let res = compute(prev);
+        next.compute(|| self.compute_and_publish(key, compute))
+    }
 
+    fn compute_and_publish(&self, key: K, compute: impl FnOnce(Option<V>) -> V) -> V
+    where
+        K: Clone + Eq + Hash,
+        V: Clone,
+    {
+        let prev = self.prev.lock().get(&key).cloned();
+        let prev = prev.and_then(|prev| prev.get());
+        let prev = prev.or_else(|| {
             let global = self.global.lock();
-            let entry = global.entry(key.clone());
-            use dashmap::mapref::entry::Entry;
-            match entry {
-                Entry::Occupied(mut entry) => {
-                    let (revision, _) = entry.get();
-                    if *revision < self.revision {
-                        entry.insert((self.revision, res.clone()));
-                    }
-                }
-                Entry::Vacant(entry) => {
+            global
+                .get(&key)
+                .filter(|global| global.0 <= self.revision)
+                .map(|global| global.1.clone())
+        });
+
+        let res = compute(prev);
+
+        let global = self.global.lock();
+        let entry = global.entry(key.clone());
+        use dashmap::mapref::entry::Entry;
+        match entry {
+            Entry::Occupied(mut entry) => {
+                let (revision, _) = entry.get();
+                if *revision < self.revision {
                     entry.insert((self.revision, res.clone()));
                 }
             }
+            Entry::Vacant(entry) => {
+                entry.insert((self.revision, res.clone()));
+            }
+        }
 
-            res
-        })
-        .clone()
+        res
     }
 
     fn crawl(&self, revision: usize) -> Self {
@@ -1326,6 +1465,68 @@ impl<K, V> IncrCacheMap<K, V> {
     }
 }
 
+#[cfg(test)]
+mod incr_cache_tests {
+    use std::sync::{Arc, Barrier, mpsc::sync_channel};
+
+    use super::IncrCacheMap;
+
+    #[test]
+    fn recursive_lookup_does_not_wait_on_running_slot() {
+        let cache = IncrCacheMap::<u8, u8>::default();
+        let started = Arc::new(Barrier::new(2));
+        let (release, wait) = sync_channel(0);
+
+        std::thread::scope(|scope| {
+            let owner_cache = cache.clone();
+            let owner_started = started.clone();
+            let owner = scope.spawn(move || {
+                owner_cache.compute(1, |_| {
+                    owner_started.wait();
+                    wait.recv().unwrap();
+                    1
+                })
+            });
+
+            started.wait();
+            assert_eq!(cache.get(&1), None);
+            release.send(()).unwrap();
+            assert_eq!(owner.join().unwrap(), 1);
+        });
+
+        assert_eq!(cache.get(&1), Some(1));
+    }
+
+    #[test]
+    fn panicked_owner_releases_deferred_slot() {
+        let cache = IncrCacheMap::<u8, u8>::default();
+        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+            cache.compute(1, |_| panic!("test panic"))
+        }));
+
+        assert!(result.is_err());
+        assert_eq!(cache.get(&1), None);
+        assert_eq!(cache.compute(1, |_| 2), 2);
+    }
+
+    #[test]
+    fn older_revision_does_not_read_newer_global_value() {
+        let cache = IncrCacheMap::<u8, u8>::default();
+        let newer = cache.crawl(2);
+        assert_eq!(newer.compute(1, |_| 20), 20);
+
+        let older = cache.crawl(1);
+        assert_eq!(
+            older.compute(1, |prev| {
+                assert_eq!(prev, None);
+                10
+            }),
+            10
+        );
+        assert_eq!(older.get(&1), Some(10));
+    }
+}
+
 #[derive(Clone)]
 struct CacheMap<T> {
     m: Arc<FxDashMap<u128, (u64, T)>>,