add memory usage information hooks

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

@@ -49,6 +49,8 @@
 //! cycle head may then iterate, which may result in a new set of iterations on the inner cycle,
 //! for each iteration of the outer cycle.
 
+use std::mem;
+
 use thin_vec::{thin_vec, ThinVec};
 
 use crate::key::DatabaseKeyIndex;
@@ -210,6 +212,10 @@ impl CycleHeads {
             }
         }
     }
+
+    pub(crate) fn allocation_size(&self) -> usize {
+        mem::size_of_val(self.0.as_slice())
+    }
 }
 
 impl IntoIterator for CycleHeads {

src/database.rs

@@ -1,6 +1,8 @@
 use std::any::Any;
 use std::borrow::Cow;
 
+use hashbrown::HashMap;
+
 use crate::zalsa::{IngredientIndex, ZalsaDatabase};
 use crate::{Durability, Revision};
 
@@ -132,4 +134,99 @@ impl dyn Database {
         let views = self.zalsa().views();
         views.downcaster_for().downcast(self)
     }
+
+    /// 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

@@ -1,9 +1,10 @@
 use std::any::Any;
 use std::fmt::{Debug, Formatter};
-use std::mem::transmute;
+use std::mem::{self, transmute};
 use std::ptr::NonNull;
 
 use crate::cycle::{empty_cycle_heads, CycleHead, CycleHeads, IterationCount, ProvisionalStatus};
+use crate::database::SlotInfo;
 use crate::function::{Configuration, IngredientImpl};
 use crate::hash::FxHashSet;
 use crate::ingredient::{Ingredient, WaitForResult};
@@ -316,6 +317,17 @@ impl<V: Send + Sync + Any> crate::table::memo::Memo for Memo<V> {
     fn origin(&self) -> QueryOriginRef<'_> {
         self.revisions.origin.as_ref()
     }
+
+    fn memory_usage(&self) -> SlotInfo {
+        let size_of = mem::size_of::<Memo<V>>() + self.revisions.allocation_size();
+
+        SlotInfo {
+            size_of_metadata: size_of - mem::size_of::<V>(),
+            debug_name: std::any::type_name::<V>(),
+            size_of_fields: mem::size_of::<V>(),
+            memos: Vec::new(),
+        }
+    }
 }
 
 pub(super) enum TryClaimHeadsResult<'me> {

src/ingredient.rs

@@ -5,6 +5,7 @@ use crate::accumulator::accumulated_map::{AccumulatedMap, InputAccumulatedValues
 use crate::cycle::{
     empty_cycle_heads, CycleHeads, CycleRecoveryStrategy, IterationCount, ProvisionalStatus,
 };
+use crate::database::SlotInfo;
 use crate::function::VerifyResult;
 use crate::plumbing::IngredientIndices;
 use crate::runtime::Running;
@@ -181,6 +182,12 @@ 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.
+    fn memory_usage(&self, _db: &dyn Database) -> Option<Vec<SlotInfo>> {
+        None
+    }
 }
 
 impl dyn Ingredient {

src/input.rs

@@ -1,6 +1,6 @@
 use std::any::{Any, TypeId};
-use std::fmt;
 use std::ops::IndexMut;
+use std::{fmt, mem};
 
 pub mod input_field;
 pub mod setter;
@@ -9,6 +9,7 @@ pub mod singleton;
 use input_field::FieldIngredientImpl;
 
 use crate::cycle::CycleHeads;
+use crate::database::SlotInfo;
 use crate::function::VerifyResult;
 use crate::id::{AsId, FromId, FromIdWithDb};
 use crate::ingredient::Ingredient;
@@ -241,6 +242,17 @@ 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.
+    fn memory_usage(&self, db: &dyn Database) -> Option<Vec<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,23 @@ 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.
+    unsafe fn memory_usage(&self, memo_table_types: &MemoTableTypes) -> SlotInfo {
+        // SAFETY: The caller guarantees this is the correct types table.
+        let memos = unsafe { memo_table_types.attach_memos(&self.memos) };
+
+        SlotInfo {
+            debug_name: C::DEBUG_NAME,
+            size_of_metadata: mem::size_of::<Self>() - mem::size_of::<C::Fields>(),
+            size_of_fields: mem::size_of::<C::Fields>(),
+            memos: memos.memory_usage(),
+        }
+    }
 }
 
 pub trait HasBuilder {

src/interned.rs

@@ -1,10 +1,10 @@
 use std::any::TypeId;
 use std::cell::{Cell, UnsafeCell};
-use std::fmt;
 use std::hash::{BuildHasher, Hash, Hasher};
 use std::marker::PhantomData;
 use std::num::NonZeroUsize;
 use std::path::{Path, PathBuf};
+use std::{fmt, mem};
 
 use crossbeam_utils::CachePadded;
 use intrusive_collections::{intrusive_adapter, LinkedList, LinkedListLink, UnsafeRef};
@@ -21,7 +21,7 @@ use crate::sync::{Arc, Mutex, OnceLock};
 use crate::table::memo::{MemoTable, MemoTableTypes, MemoTableWithTypesMut};
 use crate::table::Slot;
 use crate::zalsa::{IngredientIndex, Zalsa};
-use crate::{Database, DatabaseKeyIndex, Event, EventKind, Id, Revision};
+use crate::{Database, DatabaseKeyIndex, Event, EventKind, Id, Revision, SlotInfo};
 
 /// Trait that defines the key properties of an interned struct.
 ///
@@ -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(not(feature = "shuttle"))]
+    unsafe fn memory_usage(&self, memo_table_types: &MemoTableTypes) -> 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) };
+
+        SlotInfo {
+            debug_name: C::DEBUG_NAME,
+            size_of_metadata: mem::size_of::<Self>() - mem::size_of::<C::Fields<'_>>(),
+            size_of_fields: 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(not(feature = "shuttle"))]
+    fn memory_usage(&self, db: &dyn Database) -> Option<Vec<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

@@ -43,7 +43,7 @@ pub use self::accumulator::Accumulator;
 pub use self::active_query::Backtrace;
 pub use self::cancelled::Cancelled;
 pub use self::cycle::CycleRecoveryAction;
-pub use self::database::{AsDynDatabase, Database};
+pub use self::database::{AsDynDatabase, Database, SlotInfo};
 pub use self::database_impl::DatabaseImpl;
 pub use self::durability::Durability;
 pub use self::event::{Event, EventKind};