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" }
@@ -83,7 +84,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 = { git = "https://github.com/ibraheemdev/get-size2", branch = "ibraheem/non-static-tracked", features = [
     "derive",
     "smallvec",
     "hashbrown",

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

@@ -203,7 +203,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()
     }
@@ -526,7 +526,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();
 

crates/ty_project/src/lib.rs

@@ -96,7 +96,7 @@ impl Override {
 }
 
 /// Resolves the settings for a given file.
-#[salsa::tracked(returns(ref), heap_size=get_size2::heap_size)]
+#[salsa::tracked(returns(ref), heap_size=ruff_memory_usage::heap_size)]
 pub(crate) fn file_settings(db: &dyn Db, file: File) -> FileSettings {
     let settings = db.project().settings(db);
 
@@ -155,7 +155,7 @@ pub(crate) fn file_settings(db: &dyn Db, file: File) -> FileSettings {
 /// This is to make Salsa happy because it requires that queries with only a single argument
 /// take a salsa-struct as argument, which isn't the case here. The `()` enables salsa's
 /// automatic interning for the arguments.
-#[salsa::tracked(heap_size=get_size2::heap_size)]
+#[salsa::tracked(heap_size=ruff_memory_usage::heap_size)]
 fn merge_overrides(db: &dyn Db, overrides: Vec<Arc<InnerOverrideOptions>>, _: ()) -> FileSettings {
     let mut overrides = overrides.into_iter().rev();
     let mut merged = (*overrides.next().unwrap()).clone();