avoid overcounting memory usage of shared objects

Author: ibraheemdev

Cargo.lock

@@ -23,6 +23,7 @@ ruff_graph = { path = "crates/ruff_graph" }
 ruff_index = { path = "crates/ruff_index" }
 ruff_linter = { path = "crates/ruff_linter" }
 ruff_macros = { path = "crates/ruff_macros" }
+ruff_memory_usage = { path = "crates/ruff_memory_usage" }
 ruff_notebook = { path = "crates/ruff_notebook" }
 ruff_options_metadata = { path = "crates/ruff_options_metadata" }
 ruff_python_ast = { path = "crates/ruff_python_ast" }
@@ -84,7 +85,7 @@ etcetera = { version = "0.10.0" }
 fern = { version = "0.7.0" }
 filetime = { version = "0.2.23" }
 getrandom = { version = "0.3.1" }
-get-size2 = { version = "0.6.0", features = [
+get-size2 = { version = "0.6.2", features = [
     "derive",
     "smallvec",
     "hashbrown",
@@ -142,7 +143,7 @@ regex-automata = { version = "0.4.9" }
 rustc-hash = { version = "2.0.0" }
 rustc-stable-hash = { version = "0.1.2" }
 # When updating salsa, make sure to also update the revision in `fuzz/Cargo.toml`
-salsa = { git = "https://github.com/salsa-rs/salsa.git", rev = "d66fe331d546216132ace503512b94d5c68d2c50", default-features = false, features = [
+salsa = { git = "https://github.com/salsa-rs/salsa.git", rev = "b121ee46c4483ba74c19e933a3522bd548eb7343", default-features = false, features = [
     "compact_str",
     "macros",
     "salsa_unstable",

Cargo.toml

@@ -14,6 +14,7 @@ license = { workspace = true }
 ruff_annotate_snippets = { workspace = true }
 ruff_cache = { workspace = true, optional = true }
 ruff_diagnostics = { workspace = true }
+ruff_memory_usage = { workspace = true }
 ruff_notebook = { workspace = true }
 ruff_python_ast = { workspace = true, features = ["get-size"] }
 ruff_python_parser = { workspace = true }

crates/ruff_db/Cargo.toml

@@ -21,7 +21,7 @@ use crate::source::source_text;
 /// reflected in the changed AST offsets.
 /// The other reason is that Ruff's AST doesn't implement `Eq` which Salsa requires
 /// for determining if a query result is unchanged.
-#[salsa::tracked(returns(ref), no_eq, heap_size=get_size2::heap_size)]
+#[salsa::tracked(returns(ref), no_eq, heap_size=ruff_memory_usage::heap_size)]
 pub fn parsed_module(db: &dyn Db, file: File) -> ParsedModule {
     let _span = tracing::trace_span!("parsed_module", ?file).entered();
 

crates/ruff_db/src/parsed.rs

@@ -9,7 +9,7 @@ use crate::Db;
 use crate::files::{File, FilePath};
 
 /// Reads the source text of a python text file (must be valid UTF8) or notebook.
-#[salsa::tracked(heap_size=get_size2::heap_size)]
+#[salsa::tracked(heap_size=ruff_memory_usage::heap_size)]
 pub fn source_text(db: &dyn Db, file: File) -> SourceText {
     let path = file.path(db);
     let _span = tracing::trace_span!("source_text", file = %path).entered();
@@ -157,7 +157,7 @@ pub enum SourceTextError {
 }
 
 /// Computes the [`LineIndex`] for `file`.
-#[salsa::tracked(heap_size=get_size2::heap_size)]
+#[salsa::tracked(heap_size=ruff_memory_usage::heap_size)]
 pub fn line_index(db: &dyn Db, file: File) -> LineIndex {
     let _span = tracing::trace_span!("line_index", ?file).entered();
 

crates/ruff_db/src/source.rs

@@ -0,0 +1,17 @@
+[package]
+name = "ruff_memory_usage"
+version = "0.0.0"
+publish = false
+authors = { workspace = true }
+edition = { workspace = true }
+rust-version = { workspace = true }
+homepage = { workspace = true }
+documentation = { workspace = true }
+repository = { workspace = true }
+license = { workspace = true }
+
+[dependencies]
+get-size2 = { workspace = true }
+
+[lints]
+workspace = true

crates/ruff_memory_usage/Cargo.toml

@@ -0,0 +1,14 @@
+use std::sync::{LazyLock, Mutex};
+
+use get_size2::{GetSize, StandardTracker};
+
+/// Returns the memory usage of the provided object, using a global tracker to avoid
+/// double-counting shared objects.
+pub fn heap_size<T: GetSize>(value: &T) -> usize {
+    static TRACKER: LazyLock<Mutex<StandardTracker>> =
+        LazyLock::new(|| Mutex::new(StandardTracker::new()));
+
+    value
+        .get_heap_size_with_tracker(&mut *TRACKER.lock().unwrap())
+        .0
+}

crates/ruff_memory_usage/src/lib.rs

@@ -15,6 +15,7 @@ license.workspace = true
 ruff_cache = { workspace = true }
 ruff_db = { workspace = true, features = ["cache", "serde"] }
 ruff_macros = { workspace = true }
+ruff_memory_usage = { workspace = true }
 ruff_options_metadata = { workspace = true }
 ruff_python_ast = { workspace = true, features = ["serde"] }
 ruff_python_formatter = { workspace = true, optional = true }

crates/ty_project/Cargo.toml

@@ -127,10 +127,9 @@ impl ProjectDatabase {
     /// Returns a [`SalsaMemoryDump`] that can be use to dump Salsa memory usage information
     /// to the CLI after a typechecker run.
     pub fn salsa_memory_dump(&self) -> SalsaMemoryDump {
-        let salsa_db = self as &dyn salsa::Database;
-
-        let mut ingredients = salsa_db.structs_info();
-        let mut memos = salsa_db.queries_info().into_iter().collect::<Vec<_>>();
+        let memory_usage = <dyn salsa::Database>::memory_usage(self);
+        let mut ingredients = memory_usage.structs;
+        let mut memos = memory_usage.queries.into_iter().collect::<Vec<_>>();
 
         ingredients.sort_by_key(|ingredient| cmp::Reverse(ingredient.size_of_fields()));
         memos.sort_by_key(|(_, memo)| cmp::Reverse(memo.size_of_fields()));
@@ -140,12 +139,14 @@ impl ProjectDatabase {
         for ingredient in &ingredients {
             total_metadata += ingredient.size_of_metadata();
             total_fields += ingredient.size_of_fields();
+            total_fields += ingredient.heap_size_of_fields().unwrap_or(0);
         }
 
         let mut total_memo_fields = 0;
         let mut total_memo_metadata = 0;
         for (_, memo) in &memos {
             total_memo_fields += memo.size_of_fields();
+            total_memo_fields += memo.heap_size_of_fields().unwrap_or(0);
             total_memo_metadata += memo.size_of_metadata();
         }
 

crates/ty_project/src/db.rs

@@ -201,7 +201,7 @@ impl Project {
     /// This is a salsa query to prevent re-computing queries if other, unrelated
     /// settings change. For example, we don't want that changing the terminal settings
     /// invalidates any type checking queries.
-    #[salsa::tracked(returns(deref), heap_size=get_size2::heap_size)]
+    #[salsa::tracked(returns(deref), heap_size=ruff_memory_usage::heap_size)]
     pub fn rules(self, db: &dyn Db) -> Arc<RuleSelection> {
         self.settings(db).to_rules()
     }
@@ -524,7 +524,7 @@ impl Project {
     }
 }
 
-#[salsa::tracked(returns(ref), heap_size=get_size2::heap_size)]
+#[salsa::tracked(returns(ref), heap_size=ruff_memory_usage::heap_size)]
 pub(crate) fn check_file_impl(db: &dyn Db, file: File) -> Result<Box<[Diagnostic]>, Diagnostic> {
     let mut diagnostics: Vec<Diagnostic> = Vec::new();