Add API to dump memory usage (#916)

* add memory usage information hooks

* gate memory usage API under `salsa_unstable` feature

* use snapshot tests for memory usage API

Author: ibraheemdev

src/accumulator/accumulated_map.rs

@@ -1,6 +1,6 @@
 use std::ops;
 
-use rustc_hash::FxHashMap;
+use rustc_hash::FxBuildHasher;
 
 use crate::accumulator::accumulated::Accumulated;
 use crate::accumulator::{Accumulator, AnyAccumulated};
@@ -9,7 +9,7 @@ use crate::IngredientIndex;
 
 #[derive(Default)]
 pub struct AccumulatedMap {
-    map: FxHashMap<IngredientIndex, Box<dyn AnyAccumulated>>,
+    map: hashbrown::HashMap<IngredientIndex, Box<dyn AnyAccumulated>, FxBuildHasher>,
 }
 
 impl std::fmt::Debug for AccumulatedMap {
@@ -50,6 +50,10 @@ impl AccumulatedMap {
     pub fn clear(&mut self) {
         self.map.clear()
     }
+
+    pub fn allocation_size(&self) -> usize {
+        self.map.allocation_size()
+    }
 }
 
 /// Tracks whether any input read during a query's execution has any accumulated values.

src/cycle.rs

@@ -210,6 +210,11 @@ impl CycleHeads {
             }
         }
     }
+
+    #[cfg(feature = "salsa_unstable")]
+    pub(crate) fn allocation_size(&self) -> usize {
+        std::mem::size_of_val(self.0.as_slice())
+    }
 }
 
 impl IntoIterator for CycleHeads {

src/database.rs

@@ -133,3 +133,109 @@ impl dyn Database {
         views.downcaster_for().downcast(self)
     }
 }
+
+#[cfg(feature = "salsa_unstable")]
+pub use memory_usage::{IngredientInfo, SlotInfo};
+
+#[cfg(feature = "salsa_unstable")]
+mod memory_usage {
+    use crate::Database;
+    use hashbrown::HashMap;
+
+    impl dyn Database {
+        /// Returns information about any Salsa structs.
+        pub fn structs_info(&self) -> Vec<IngredientInfo> {
+            self.zalsa()
+                .ingredients()
+                .filter_map(|ingredient| {
+                    let mut size_of_fields = 0;
+                    let mut size_of_metadata = 0;
+                    let mut instances = 0;
+
+                    for slot in ingredient.memory_usage(self)? {
+                        instances += 1;
+                        size_of_fields += slot.size_of_fields;
+                        size_of_metadata += slot.size_of_metadata;
+                    }
+
+                    Some(IngredientInfo {
+                        count: instances,
+                        size_of_fields,
+                        size_of_metadata,
+                        debug_name: ingredient.debug_name(),
+                    })
+                })
+                .collect()
+        }
+
+        /// 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> {
+            let mut queries = HashMap::new();
+
+            for input_ingredient in self.zalsa().ingredients() {
+                let Some(input_info) = input_ingredient.memory_usage(self) else {
+                    continue;
+                };
+
+                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()
+                            });
+
+                        info.count += 1;
+                        info.size_of_fields += output.size_of_fields;
+                        info.size_of_metadata += output.size_of_metadata;
+                    }
+                }
+            }
+
+            queries
+        }
+    }
+
+    /// Information about instances of a particular Salsa ingredient.
+    #[derive(Default, Debug, PartialEq, Eq, PartialOrd, Ord)]
+    pub struct IngredientInfo {
+        debug_name: &'static str,
+        count: usize,
+        size_of_metadata: usize,
+        size_of_fields: usize,
+    }
+
+    impl IngredientInfo {
+        /// Returns the debug name of the ingredient.
+        pub fn debug_name(&self) -> &'static str {
+            self.debug_name
+        }
+
+        /// Returns the total size of the fields of any instances of this ingredient, in bytes.
+        pub fn size_of_fields(&self) -> usize {
+            self.size_of_fields
+        }
+
+        /// Returns the total size of Salsa metadata of any instances of this ingredient, in bytes.
+        pub fn size_of_metadata(&self) -> usize {
+            self.size_of_metadata
+        }
+
+        /// Returns the number of instances of this ingredient.
+        pub fn count(&self) -> usize {
+            self.count
+        }
+    }
+
+    /// Memory usage information about a particular instance of struct, input or output.
+    pub struct SlotInfo {
+        pub(crate) debug_name: &'static str,
+        pub(crate) size_of_metadata: usize,
+        pub(crate) size_of_fields: usize,
+        pub(crate) memos: Vec<SlotInfo>,
+    }
+}

src/function/memo.rs

@@ -316,6 +316,18 @@ impl<V: Send + Sync + Any> crate::table::memo::Memo for Memo<V> {
     fn origin(&self) -> QueryOriginRef<'_> {
         self.revisions.origin.as_ref()
     }
+
+    #[cfg(feature = "salsa_unstable")]
+    fn memory_usage(&self) -> crate::SlotInfo {
+        let size_of = std::mem::size_of::<Memo<V>>() + self.revisions.allocation_size();
+
+        crate::SlotInfo {
+            size_of_metadata: size_of - std::mem::size_of::<V>(),
+            debug_name: std::any::type_name::<V>(),
+            size_of_fields: std::mem::size_of::<V>(),
+            memos: Vec::new(),
+        }
+    }
 }
 
 pub(super) enum TryClaimHeadsResult<'me> {

src/ingredient.rs

@@ -181,6 +181,13 @@ pub trait Ingredient: Any + std::fmt::Debug + Send + Sync {
         let _ = (db, key_index);
         (None, InputAccumulatedValues::Empty)
     }
+
+    /// Returns memory usage information about any instances of the ingredient,
+    /// if applicable.
+    #[cfg(feature = "salsa_unstable")]
+    fn memory_usage(&self, _db: &dyn Database) -> Option<Vec<crate::SlotInfo>> {
+        None
+    }
 }
 
 impl dyn Ingredient {

src/input.rs

@@ -241,6 +241,18 @@ impl<C: Configuration> Ingredient for IngredientImpl<C> {
     fn memo_table_types(&self) -> Arc<MemoTableTypes> {
         self.memo_table_types.clone()
     }
+
+    /// Returns memory usage information about any inputs.
+    #[cfg(feature = "salsa_unstable")]
+    fn memory_usage(&self, db: &dyn Database) -> Option<Vec<crate::SlotInfo>> {
+        let memory_usage = self
+            .entries(db)
+            // SAFETY: The memo table belongs to a value that we allocated, so it
+            // has the correct type.
+            .map(|value| unsafe { value.memory_usage(&self.memo_table_types) })
+            .collect();
+        Some(memory_usage)
+    }
 }
 
 impl<C: Configuration> std::fmt::Debug for IngredientImpl<C> {
@@ -284,6 +296,24 @@ where
     pub fn fields(&self) -> &C::Fields {
         &self.fields
     }
+
+    /// Returns memory usage information about the input.
+    ///
+    /// # Safety
+    ///
+    /// The `MemoTable` must belong to a `Value` of the correct type.
+    #[cfg(feature = "salsa_unstable")]
+    unsafe fn memory_usage(&self, memo_table_types: &MemoTableTypes) -> crate::SlotInfo {
+        // SAFETY: The caller guarantees this is the correct types table.
+        let memos = unsafe { memo_table_types.attach_memos(&self.memos) };
+
+        crate::SlotInfo {
+            debug_name: C::DEBUG_NAME,
+            size_of_metadata: std::mem::size_of::<Self>() - std::mem::size_of::<C::Fields>(),
+            size_of_fields: std::mem::size_of::<C::Fields>(),
+            memos: memos.memory_usage(),
+        }
+    }
 }
 
 pub trait HasBuilder {

src/interned.rs

@@ -190,6 +190,28 @@ where
         // ensures that they are not reused while being accessed.
         unsafe { &*self.fields.get() }
     }
+
+    /// Returns memory usage information about the interned value.
+    ///
+    /// # Safety
+    ///
+    /// The `MemoTable` must belong to a `Value` of the correct type. Additionally, the
+    /// lock must be held for the shard containing the value.
+    #[cfg(all(not(feature = "shuttle"), feature = "salsa_unstable"))]
+    unsafe fn memory_usage(&self, memo_table_types: &MemoTableTypes) -> crate::SlotInfo {
+        // SAFETY: The caller guarantees we hold the lock for the shard containing the value, so we
+        // have at-least read-only access to the value's memos.
+        let memos = unsafe { &*self.memos.get() };
+        // SAFETY: The caller guarantees this is the correct types table.
+        let memos = unsafe { memo_table_types.attach_memos(memos) };
+
+        crate::SlotInfo {
+            debug_name: C::DEBUG_NAME,
+            size_of_metadata: std::mem::size_of::<Self>() - std::mem::size_of::<C::Fields<'_>>(),
+            size_of_fields: std::mem::size_of::<C::Fields<'_>>(),
+            memos: memos.memory_usage(),
+        }
+    }
 }
 
 impl<C: Configuration> Default for JarImpl<C> {
@@ -680,7 +702,7 @@ where
     //
     // # Safety
     //
-    // The lock must be held.
+    // The lock must be held for the shard containing the value.
     unsafe fn value_hash<'db>(&'db self, id: Id, zalsa: &'db Zalsa) -> u64 {
         // This closure is only called if the table is resized. So while it's expensive
         // to lookup all values, it will only happen rarely.
@@ -694,7 +716,7 @@ where
     //
     // # Safety
     //
-    // The lock must be held.
+    // The lock must be held for the shard containing the value.
     unsafe fn value_eq<'db, Key>(
         id: Id,
         key: &Key,
@@ -830,6 +852,31 @@ where
     fn memo_table_types(&self) -> Arc<MemoTableTypes> {
         self.memo_table_types.clone()
     }
+
+    /// Returns memory usage information about any interned values.
+    #[cfg(all(not(feature = "shuttle"), feature = "salsa_unstable"))]
+    fn memory_usage(&self, db: &dyn Database) -> Option<Vec<crate::SlotInfo>> {
+        use parking_lot::lock_api::RawMutex;
+
+        for shard in self.shards.iter() {
+            // SAFETY: We do not hold any active mutex guards.
+            unsafe { shard.raw().lock() };
+        }
+
+        let memory_usage = self
+            .entries(db)
+            // SAFETY: The memo table belongs to a value that we allocated, so it
+            // has the correct type. Additionally, we are holding the locks for all shards.
+            .map(|value| unsafe { value.memory_usage(&self.memo_table_types) })
+            .collect();
+
+        for shard in self.shards.iter() {
+            // SAFETY: We acquired the locks for all shards.
+            unsafe { shard.raw().unlock() };
+        }
+
+        Some(memory_usage)
+    }
 }
 
 impl<C> std::fmt::Debug for IngredientImpl<C>

src/lib.rs

@@ -39,6 +39,9 @@ pub use parallel::{join, par_map};
 #[cfg(feature = "macros")]
 pub use salsa_macros::{accumulator, db, input, interned, tracked, Supertype, Update};
 
+#[cfg(feature = "salsa_unstable")]
+pub use self::database::{IngredientInfo, SlotInfo};
+
 pub use self::accumulator::Accumulator;
 pub use self::active_query::Backtrace;
 pub use self::cancelled::Cancelled;

src/table.rs

@@ -253,6 +253,7 @@ impl Table {
         unsafe { page.memo_types.attach_memos_mut(memos) }
     }
 
+    #[cfg(feature = "salsa_unstable")]
     pub(crate) fn slots_of<T: Slot>(&self) -> impl Iterator<Item = &T> + '_ {
         self.pages
             .iter()
@@ -392,6 +393,7 @@ impl Page {
         PageView(self, PhantomData)
     }
 
+    #[cfg(feature = "salsa_unstable")]
     fn cast_type<T: Slot>(&self) -> Option<PageView<'_, T>> {
         if self.slot_type_id == TypeId::of::<T>() {
             Some(PageView(self, PhantomData))