[ty] Add environment variable to dump Salsa memory usage stats (#18928)

## Summary

Setting `TY_MEMORY_REPORT=full` will generate and print a memory usage
report to the CLI after a `ty check` run:

```
=======SALSA STRUCTS=======
`Definition`                                       metadata=7.24MB   fields=17.38MB  count=181062
`Expression`                                       metadata=4.45MB   fields=5.94MB   count=92804
`member_lookup_with_policy_::interned_arguments`   metadata=1.97MB   fields=2.25MB   count=35176
...
=======SALSA QUERIES=======
`File -> ty_python_semantic::semantic_index::SemanticIndex`
    metadata=11.46MB  fields=88.86MB  count=1638
`Definition -> ty_python_semantic::types::infer::TypeInference`
    metadata=24.52MB  fields=86.68MB  count=146018
`File -> ruff_db::parsed::ParsedModule`
    metadata=0.12MB   fields=69.06MB  count=1642
...
=======SALSA SUMMARY=======
TOTAL MEMORY USAGE: 577.61MB
    struct metadata = 29.00MB
    struct fields = 35.68MB
    memo metadata = 103.87MB
    memo fields = 409.06MB
```

Eventually, we should integrate these numbers into CI in some form. The
one limitation currently is that heap allocations in salsa structs (e.g.
interned values) are not tracked, but memoized values should have full
coverage. We may also want a peak memory usage counter (that accounts
for non-salsa memory), but that is relatively simple to profile manually
(e.g. `time -v ty check`) and would require a compile-time option to
avoid runtime overhead.
This commit is contained in:
Ibraheem Ahmed 2025-06-26 17:27:51 -04:00 committed by GitHub
parent a1579d82d0
commit 6f7b1c9bb3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
79 changed files with 905 additions and 207 deletions

View file

@ -8,7 +8,7 @@ use super::path::SearchPath;
use crate::module_name::ModuleName;
/// Representation of a Python module.
#[derive(Clone, PartialEq, Eq, Hash)]
#[derive(Clone, PartialEq, Eq, Hash, get_size2::GetSize)]
pub struct Module {
inner: Arc<ModuleInner>,
}
@ -99,7 +99,7 @@ impl std::fmt::Debug for Module {
}
}
#[derive(PartialEq, Eq, Hash)]
#[derive(PartialEq, Eq, Hash, get_size2::GetSize)]
enum ModuleInner {
/// A module that resolves to a file (`lib.py` or `package/__init__.py`)
FileModule {
@ -116,7 +116,7 @@ enum ModuleInner {
NamespacePackage { name: ModuleName },
}
#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, get_size2::GetSize)]
pub enum ModuleKind {
/// A single-file module (e.g. `foo.py` or `foo.pyi`)
Module,
@ -135,7 +135,7 @@ impl ModuleKind {
}
/// Enumeration of various core stdlib modules in which important types are located
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, strum_macros::EnumString)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, strum_macros::EnumString, get_size2::GetSize)]
#[cfg_attr(test, derive(strum_macros::EnumIter))]
#[strum(serialize_all = "snake_case")]
pub enum KnownModule {

View file

@ -371,7 +371,7 @@ impl From<SitePackagesDiscoveryError> for SearchPathValidationError {
type SearchPathResult<T> = Result<T, SearchPathValidationError>;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[derive(Debug, Clone, PartialEq, Eq, Hash, get_size2::GetSize)]
enum SearchPathInner {
Extra(SystemPathBuf),
FirstParty(SystemPathBuf),
@ -406,7 +406,7 @@ enum SearchPathInner {
/// or the "Editable" category. For the "First-party", "Site-packages"
/// and "Standard-library" categories, however, there will always be exactly
/// one search path from that category in any given list of search paths.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[derive(Debug, Clone, PartialEq, Eq, Hash, get_size2::GetSize)]
pub(crate) struct SearchPath(Arc<SearchPathInner>);
impl SearchPath {

View file

@ -30,7 +30,7 @@ pub fn resolve_module(db: &dyn Db, module_name: &ModuleName) -> Option<Module> {
///
/// This query should not be called directly. Instead, use [`resolve_module`]. It only exists
/// because Salsa requires the module name to be an ingredient.
#[salsa::tracked]
#[salsa::tracked(heap_size=get_size2::GetSize::get_heap_size)]
pub(crate) fn resolve_module_query<'db>(
db: &'db dyn Db,
module_name: ModuleNameIngredient<'db>,
@ -95,7 +95,7 @@ impl std::fmt::Display for SystemOrVendoredPathRef<'_> {
/// Resolves the module for the file with the given id.
///
/// Returns `None` if the file is not a module locatable via any of the known search paths.
#[salsa::tracked]
#[salsa::tracked(heap_size=get_size2::GetSize::get_heap_size)]
pub(crate) fn file_to_module(db: &dyn Db, file: File) -> Option<Module> {
let _span = tracing::trace_span!("file_to_module", ?file).entered();
@ -297,7 +297,7 @@ impl SearchPaths {
/// The editable-install search paths for the first `site-packages` directory
/// should come between the two `site-packages` directories when it comes to
/// module-resolution priority.
#[salsa::tracked(returns(deref))]
#[salsa::tracked(returns(deref), heap_size=get_size2::GetSize::get_heap_size)]
pub(crate) fn dynamic_resolution_paths(db: &dyn Db) -> Vec<SearchPath> {
tracing::debug!("Resolving dynamic module resolution paths");