Add snapshotting mechanism:

Cargo.toml

@@ -30,7 +30,7 @@ console_error_panic_hook = { version = "0.1.7" }
 console_log = { version = "1.0.0" }
 countme = { version = "3.0.1" }
 criterion = { version = "0.5.1", default-features = false }
-crossbeam-channel = { version = "0.5.12" }
+crossbeam = { version = "0.8.4" }
 dashmap = { version = "5.5.3" }
 dirs = { version = "5.0.0" }
 drop_bomb = { version = "0.1.5" }

crates/red_knot/Cargo.toml

@@ -22,7 +22,7 @@ ruff_notebook = { path = "../ruff_notebook" }
 anyhow = { workspace = true }
 bitflags = { workspace = true }
 ctrlc = "3.4.4"
-crossbeam-channel = { workspace = true }
+crossbeam = { workspace = true }
 dashmap = { workspace = true }
 hashbrown = { workspace = true }
 indexmap = { workspace = true }

crates/red_knot/src/cache.rs

@@ -2,6 +2,7 @@ use std::fmt::Formatter;
 use std::hash::Hash;
 use std::sync::atomic::{AtomicUsize, Ordering};
 
+use crate::db::QueryResult;
 use dashmap::mapref::entry::Entry;
 
 use crate::FxDashMap;
@@ -27,11 +28,11 @@ where
         }
     }
 
-    pub fn get<F>(&self, key: &K, compute: F) -> V
+    pub fn get<F>(&self, key: &K, compute: F) -> QueryResult<V>
     where
-        F: FnOnce(&K) -> V,
+        F: FnOnce(&K) -> QueryResult<V>,
     {
-        match self.map.entry(key.clone()) {
+        Ok(match self.map.entry(key.clone()) {
             Entry::Occupied(cached) => {
                 self.statistics.hit();
 
@@ -40,11 +41,11 @@ where
             Entry::Vacant(vacant) => {
                 self.statistics.miss();
 
-                let value = compute(key);
+                let value = compute(key)?;
                 vacant.insert(value.clone());
                 value
             }
-        }
+        })
     }
 
     pub fn set(&mut self, key: K, value: V) {

crates/red_knot/src/cancellation.rs

@@ -1,6 +1,6 @@
 use std::sync::{Arc, Condvar, Mutex};
 
-#[derive(Debug, Default)]
+#[derive(Debug, Clone, Default)]
 pub struct CancellationTokenSource {
     signal: Arc<(Mutex<bool>, Condvar)>,
 }

crates/red_knot/src/db.rs

@@ -1,3 +1,7 @@
+mod jars;
+mod query;
+mod storage;
+
 use std::path::Path;
 use std::sync::Arc;
 
@@ -9,32 +13,35 @@ use crate::source::{Source, SourceStorage};
 use crate::symbols::{SymbolId, SymbolTable, SymbolTablesStorage};
 use crate::types::{Type, TypeStore};
 
+pub use jars::HasJar;
+pub use query::{QueryError, QueryResult};
+
 pub trait SourceDb {
     // queries
     fn file_id(&self, path: &std::path::Path) -> FileId;
 
     fn file_path(&self, file_id: FileId) -> Arc<std::path::Path>;
 
-    fn source(&self, file_id: FileId) -> Source;
+    fn source(&self, file_id: FileId) -> QueryResult<Source>;
 
-    fn parse(&self, file_id: FileId) -> Parsed;
+    fn parse(&self, file_id: FileId) -> QueryResult<Parsed>;
 
-    fn lint_syntax(&self, file_id: FileId) -> Diagnostics;
+    fn lint_syntax(&self, file_id: FileId) -> QueryResult<Diagnostics>;
 }
 
 pub trait SemanticDb: SourceDb {
     // queries
-    fn resolve_module(&self, name: ModuleName) -> Option<Module>;
+    fn resolve_module(&self, name: ModuleName) -> QueryResult<Option<Module>>;
 
-    fn file_to_module(&self, file_id: FileId) -> Option<Module>;
+    fn file_to_module(&self, file_id: FileId) -> QueryResult<Option<Module>>;
 
-    fn path_to_module(&self, path: &Path) -> Option<Module>;
+    fn path_to_module(&self, path: &Path) -> QueryResult<Option<Module>>;
 
-    fn symbol_table(&self, file_id: FileId) -> Arc<SymbolTable>;
+    fn symbol_table(&self, file_id: FileId) -> QueryResult<Arc<SymbolTable>>;
 
-    fn eval_symbol(&self, file_id: FileId, symbol_id: SymbolId) -> Type;
+    fn eval_symbol(&self, file_id: FileId, symbol_id: SymbolId) -> QueryResult<Type>;
 
-    fn lint_semantic(&self, file_id: FileId) -> Diagnostics;
+    fn lint_semantic(&self, file_id: FileId) -> QueryResult<Diagnostics>;
 
     // mutations
 
@@ -60,26 +67,6 @@ pub struct SemanticJar {
     pub lint_semantic: LintSemanticStorage,
 }
 
-/// Gives access to a specific jar in the database.
-///
-/// Nope, the terminology isn't borrowed from Java but from Salsa <https://salsa-rs.github.io/salsa/>,
-/// which is an analogy to storing the salsa in different jars.
-///
-/// The basic idea is that each crate can define its own jar and the jars can be combined to a single
-/// database in the top level crate. Each crate also defines its own `Database` trait. The combination of
-/// `Database` trait and the jar allows to write queries in isolation without having to know how they get composed at the upper levels.
-///
-/// Salsa further defines a `HasIngredient` trait which slices the jar to a specific storage (e.g. a specific cache).
-/// We don't need this just jet because we write our queries by hand. We may want a similar trait if we decide
-/// to use a macro to generate the queries.
-pub trait HasJar<T> {
-    /// Gives a read-only reference to the jar.
-    fn jar(&self) -> &T;
-
-    /// Gives a mutable reference to the jar.
-    fn jar_mut(&mut self) -> &mut T;
-}
-
 #[cfg(test)]
 pub(crate) mod tests {
     use std::path::Path;

crates/red_knot/src/db/jars.rs

@@ -0,0 +1,33 @@
+use crate::db::QueryResult;
+
+/// Gives access to a specific jar in the database.
+///
+/// Nope, the terminology isn't borrowed from Java but from Salsa <https://salsa-rs.github.io/salsa/>,
+/// which is an analogy to storing the salsa in different jars.
+///
+/// The basic idea is that each crate can define its own jar and the jars can be combined to a single
+/// database in the top level crate. Each crate also defines its own `Database` trait. The combination of
+/// `Database` trait and the jar allows to write queries in isolation without having to know how they get composed at the upper levels.
+///
+/// Salsa further defines a `HasIngredient` trait which slices the jar to a specific storage (e.g. a specific cache).
+/// We don't need this just jet because we write our queries by hand. We may want a similar trait if we decide
+/// to use a macro to generate the queries.
+pub trait HasJar<T> {
+    /// Gives a read-only reference to the jar.
+    fn jar(&self) -> QueryResult<&T>;
+
+    fn jar_by_pass_cancellation(&self) -> &T;
+
+    /// Gives a mutable reference to the jar.
+    fn jar_mut(&mut self) -> &mut T;
+}
+
+pub trait HasJars {
+    type Jars;
+
+    fn jars(&self) -> &QueryResult<Self::Jars>;
+
+    fn jars_unwrap(&self) -> &Self::Jars;
+
+    fn jars_mut(&mut self) -> &mut Self::Jars;
+}

crates/red_knot/src/db/query.rs

@@ -0,0 +1,6 @@
+#[derive(Debug, Clone, Copy)]
+pub enum QueryError {
+    Cancelled,
+}
+
+pub type QueryResult<T> = Result<T, QueryError>;

crates/red_knot/src/db/storage.rs

@@ -0,0 +1,61 @@
+use crate::cancellation::{CancellationToken, CancellationTokenSource};
+use crate::db::jars::HasJars;
+use crate::db::query::{QueryError, QueryResult};
+use crossbeam::sync::WaitGroup;
+use std::sync::Arc;
+
+pub struct JarStorage<T>
+where
+    T: HasJars,
+{
+    db: T,
+}
+
+#[derive(Clone, Debug)]
+pub struct SharedStorage<T>
+where
+    T: HasJars,
+{
+    // It's important that the wait group is declared after `jars` to ensure that `jars` is dropped first.
+    // See https://doc.rust-lang.org/reference/destructors.html
+    jars: Arc<T::Jars>,
+
+    /// Used to count the references to `jars`. Allows implementing [`jars_mut`] without requiring to clone `jars`.
+    jars_references: WaitGroup,
+
+    cancellation_token_source: CancellationTokenSource,
+}
+
+impl<T> SharedStorage<T>
+where
+    T: HasJars,
+{
+    pub(super) fn jars(&self) -> QueryResult<&T::Jars> {
+        self.err_if_cancelled()?;
+        Ok(&self.jars)
+    }
+
+    pub(super) fn jars_mut(&mut self) -> &mut T::Jars {
+        // Cancel all pending queries.
+        self.cancellation_token_source.cancel();
+
+        let existing_wait = std::mem::take(&mut self.jars_references);
+        existing_wait.wait();
+        self.cancellation_token_source = CancellationTokenSource::new();
+
+        // Now all other references to `self.jars` should have been released. We can now safely return a mutable reference
+        // to the Arc's content.
+        let jars =
+            Arc::get_mut(&mut self.jars).expect("All references to jars should have been released");
+
+        jars
+    }
+
+    pub(super) fn err_if_cancelled(&self) -> QueryResult<()> {
+        if self.cancellation_token_source.is_cancelled() {
+            Err(QueryError::Cancelled)
+        } else {
+            Ok(())
+        }
+    }
+}