add option to track heap memory usage of memos

Author: ibraheemdev

Cargo.toml

@@ -63,6 +63,7 @@ expect-test = "1.5.0"
 rustversion = "1.0"
 test-log = { version = "0.2.11", features = ["trace"] }
 trybuild = "1.0"
+get-size2 = { version = "0.4.1", features = ["derive"] }
 
 [target.'cfg(all(not(target_os = "windows"), not(target_os = "openbsd"), any(target_arch = "x86_64", target_arch = "aarch64", target_arch = "powerpc64")))'.dev-dependencies]
 tikv-jemallocator = "0.6.0"

components/salsa-macro-rules/src/setup_tracked_fn.rs

@@ -55,6 +55,9 @@ macro_rules! setup_tracked_fn {
         // If true, the input needs an interner (because it has >1 argument).
         needs_interner: $needs_interner:tt,
 
+        // The function used to implement `C::heap_size`.
+        heap_size_fn: $($heap_size_fn:path)?,
+
         // LRU capacity (a literal, maybe 0)
         lru: $lru:tt,
 
@@ -196,6 +199,12 @@ macro_rules! setup_tracked_fn {
 
                 $($values_equal)+
 
+                $(
+                    fn heap_size(value: &Self::Output<'_>) -> usize {
+                        $heap_size_fn(value)
+                    }
+                )?
+
                 fn execute<$db_lt>($db: &$db_lt Self::DbView, ($($input_id),*): ($($interned_input_ty),*)) -> Self::Output<$db_lt> {
                     $($assert_return_type_is_update)*
 

components/salsa-macros/src/accumulator.rs

@@ -45,6 +45,7 @@ impl AllowedOptions for Accumulator {
     const CONSTRUCTOR_NAME: bool = false;
     const ID: bool = false;
     const REVISIONS: bool = false;
+    const HEAP_SIZE: bool = false;
 }
 
 struct StructMacro {

components/salsa-macros/src/input.rs

@@ -64,6 +64,8 @@ impl crate::options::AllowedOptions for InputStruct {
     const ID: bool = false;
 
     const REVISIONS: bool = false;
+
+    const HEAP_SIZE: bool = false;
 }
 
 impl SalsaStructAllowedOptions for InputStruct {

components/salsa-macros/src/interned.rs

@@ -64,6 +64,8 @@ impl crate::options::AllowedOptions for InternedStruct {
     const ID: bool = true;
 
     const REVISIONS: bool = true;
+
+    const HEAP_SIZE: bool = false;
 }
 
 impl SalsaStructAllowedOptions for InternedStruct {

components/salsa-macros/src/options.rs

@@ -99,6 +99,12 @@ pub(crate) struct Options<A: AllowedOptions> {
     /// This is stored as a `syn::Expr` to support `usize::MAX`.
     pub revisions: Option<syn::Expr>,
 
+    /// The `heap_size = <path>` option can be used to track heap memory usage of memoized
+    /// values.
+    ///
+    /// If this is `Some`, the value is the provided `heap_size` function.
+    pub heap_size_fn: Option<syn::Path>,
+
     /// Remember the `A` parameter, which plays no role after parsing.
     phantom: PhantomData<A>,
 }
@@ -123,6 +129,7 @@ impl<A: AllowedOptions> Default for Options<A> {
             singleton: Default::default(),
             id: Default::default(),
             revisions: Default::default(),
+            heap_size_fn: Default::default(),
         }
     }
 }
@@ -145,6 +152,7 @@ pub(crate) trait AllowedOptions {
     const CONSTRUCTOR_NAME: bool;
     const ID: bool;
     const REVISIONS: bool;
+    const HEAP_SIZE: bool;
 }
 
 type Equals = syn::Token![=];
@@ -392,6 +400,22 @@ impl<A: AllowedOptions> syn::parse::Parse for Options<A> {
                         "`revisions` option not allowed here",
                     ));
                 }
+            } else if ident == "heap_size" {
+                if A::HEAP_SIZE {
+                    let _eq = Equals::parse(input)?;
+                    let path = syn::Path::parse(input)?;
+                    if let Some(old) = options.heap_size_fn.replace(path) {
+                        return Err(syn::Error::new(
+                            old.span(),
+                            "option `heap_size` provided twice",
+                        ));
+                    }
+                } else {
+                    return Err(syn::Error::new(
+                        ident.span(),
+                        "`heap_size` option not allowed here",
+                    ));
+                }
             } else {
                 return Err(syn::Error::new(
                     ident.span(),

components/salsa-macros/src/tracked_fn.rs

@@ -57,6 +57,8 @@ impl crate::options::AllowedOptions for TrackedFn {
     const ID: bool = false;
 
     const REVISIONS: bool = false;
+
+    const HEAP_SIZE: bool = true;
 }
 
 struct Macro {
@@ -97,6 +99,7 @@ impl Macro {
             self.cycle_recovery()?;
         let is_specifiable = self.args.specify.is_some();
         let requires_update = self.args.non_update_return_type.is_none();
+        let heap_size_fn = self.args.heap_size_fn.iter();
         let eq = if let Some(token) = &self.args.no_eq {
             if self.args.cycle_fn.is_some() {
                 return Err(syn::Error::new_spanned(
@@ -217,6 +220,7 @@ impl Macro {
                 is_specifiable: #is_specifiable,
                 values_equal: {#eq},
                 needs_interner: #needs_interner,
+                heap_size_fn: #(#heap_size_fn)*,
                 lru: #lru,
                 return_mode: #return_mode,
                 assert_return_type_is_update: { #assert_return_type_is_update },

components/salsa-macros/src/tracked_struct.rs

@@ -60,6 +60,8 @@ impl crate::options::AllowedOptions for TrackedStruct {
     const ID: bool = false;
 
     const REVISIONS: bool = false;
+
+    const HEAP_SIZE: bool = false;
 }
 
 impl SalsaStructAllowedOptions for TrackedStruct {

src/database.rs

@@ -135,7 +135,10 @@ impl dyn Database {
 }
 
 #[cfg(feature = "salsa_unstable")]
-pub use memory_usage::{IngredientInfo, SlotInfo};
+pub use memory_usage::IngredientInfo;
+
+#[cfg(feature = "salsa_unstable")]
+pub(crate) use memory_usage::{MemoInfo, SlotInfo};
 
 #[cfg(feature = "salsa_unstable")]
 mod memory_usage {
@@ -171,8 +174,8 @@ mod memory_usage {
         /// Returns information about any memoized Salsa queries.
         ///
         /// The returned map holds memory usage information for memoized values of a given query, keyed
-        /// by its `(input, output)` type names.
-        pub fn queries_info(&self) -> HashMap<(&'static str, &'static str), IngredientInfo> {
+        /// by the query function name.
+        pub fn queries_info(&self) -> HashMap<&'static str, IngredientInfo> {
             let mut queries = HashMap::new();
 
             for input_ingredient in self.zalsa().ingredients() {
@@ -181,17 +184,15 @@ mod memory_usage {
                 };
 
                 for input in input_info {
-                    for output in input.memos {
-                        let info = queries
-                            .entry((input.debug_name, output.debug_name))
-                            .or_insert(IngredientInfo {
-                                debug_name: output.debug_name,
-                                ..Default::default()
-                            });
+                    for memo in input.memos {
+                        let info = queries.entry(memo.debug_name).or_insert(IngredientInfo {
+                            debug_name: memo.output.debug_name,
+                            ..Default::default()
+                        });
 
                         info.count += 1;
-                        info.size_of_fields += output.size_of_fields;
-                        info.size_of_metadata += output.size_of_metadata;
+                        info.size_of_fields += memo.output.size_of_fields;
+                        info.size_of_metadata += memo.output.size_of_metadata;
                     }
                 }
             }
@@ -236,6 +237,12 @@ mod memory_usage {
         pub(crate) debug_name: &'static str,
         pub(crate) size_of_metadata: usize,
         pub(crate) size_of_fields: usize,
-        pub(crate) memos: Vec<SlotInfo>,
+        pub(crate) memos: Vec<MemoInfo>,
+    }
+
+    /// Memory usage information about a particular memo.
+    pub struct MemoInfo {
+        pub(crate) debug_name: &'static str,
+        pub(crate) output: SlotInfo,
     }
 }

src/function.rs

@@ -36,7 +36,7 @@ mod memo;
 mod specify;
 mod sync;
 
-pub type Memo<C> = memo::Memo<<C as Configuration>::Output<'static>>;
+pub type Memo<C> = memo::Memo<'static, C>;
 
 pub trait Configuration: Any {
     const DEBUG_NAME: &'static str;
@@ -72,6 +72,11 @@ pub trait Configuration: Any {
     /// This is a no-op if the input to the function is a salsa struct.
     fn id_to_input(db: &Self::DbView, key: Id) -> Self::Input<'_>;
 
+    /// Returns the size of any heap allocations in the output value, in bytes.
+    fn heap_size(_value: &Self::Output<'_>) -> usize {
+        0
+    }
+
     /// Invoked when we need to compute the value for the given key, either because we've never
     /// computed it before or because the old one relied on inputs that have changed.
     ///
@@ -181,8 +186,8 @@ where
     /// only cleared with `&mut self`.
     unsafe fn extend_memo_lifetime<'this>(
         &'this self,
-        memo: &memo::Memo<C::Output<'this>>,
-    ) -> &'this memo::Memo<C::Output<'this>> {
+        memo: &memo::Memo<'this, C>,
+    ) -> &'this memo::Memo<'this, C> {
         // SAFETY: the caller must guarantee that the memo will not be released before `&self`
         unsafe { std::mem::transmute(memo) }
     }
@@ -191,9 +196,9 @@ where
         &'db self,
         zalsa: &'db Zalsa,
         id: Id,
-        mut memo: memo::Memo<C::Output<'db>>,
+        mut memo: memo::Memo<'db, C>,
         memo_ingredient_index: MemoIngredientIndex,
-    ) -> &'db memo::Memo<C::Output<'db>> {
+    ) -> &'db memo::Memo<'db, C> {
         if let Some(tracked_struct_ids) = memo.revisions.tracked_struct_ids_mut() {
             tracked_struct_ids.shrink_to_fit();
         }