Skip to content

Commit 007f92f

Browse files
committed
move salsa memory dump to ProjectDatabase
1 parent ea3bb1a commit 007f92f

File tree

4 files changed

+174
-168
lines changed

4 files changed

+174
-168
lines changed

crates/ruff_db/src/source.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,9 @@ impl get_size2::GetSize for SourceTextKind {
142142
match self {
143143
SourceTextKind::Text(text) => text.get_heap_size(),
144144
// TODO: The `get-size` derive does not support ignoring enum variants.
145+
//
146+
// Jupyter notebooks are not very relevant for memory profiling, and contain
147+
// arbitrary JSON values that do not implement the `GetSize` trait.
145148
SourceTextKind::Notebook(_) => 0,
146149
}
147150
}

crates/ty/src/lib.rs

Lines changed: 4 additions & 167 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,6 @@ mod version;
55

66
pub use args::Cli;
77

8-
use std::cmp;
9-
use std::fmt::Display;
108
use std::io::{self, BufWriter, Write, stdout};
119
use std::process::{ExitCode, Termination};
1210

@@ -147,11 +145,13 @@ fn run_check(args: CheckCommand) -> anyhow::Result<ExitStatus> {
147145

148146
let mut stdout = stdout().lock();
149147
match std::env::var("TY_MEMORY_REPORT").as_deref() {
150-
Ok("short") => write!(stdout, "{}", salsa_memory_dump(&db).display_short())?,
151-
Ok("full") => write!(stdout, "{}", salsa_memory_dump(&db).display_full())?,
148+
Ok("short") => write!(stdout, "{}", db.salsa_memory_dump().display_short())?,
149+
Ok("full") => write!(stdout, "{}", db.salsa_memory_dump().display_full())?,
152150
_ => {}
153151
}
154152

153+
tracing::trace!("Counts for entire CLI run:\n{}", countme::get_all());
154+
155155
std::mem::forget(db);
156156

157157
if exit_zero {
@@ -460,166 +460,3 @@ fn setup_rayon() {
460460
.build_global()
461461
.unwrap();
462462
}
463-
464-
/// Dumps memory usage information to the CLI.
465-
#[allow(clippy::print_stdout)]
466-
fn salsa_memory_dump(db: &ProjectDatabase) -> SalsaMemoryDump {
467-
println!("Counts for entire CLI run:\n{}", countme::get_all());
468-
469-
let salsa_db = db as &dyn salsa::Database;
470-
471-
let mut ingredients = salsa_db.structs_info();
472-
let mut memos = salsa_db.queries_info().into_iter().collect::<Vec<_>>();
473-
474-
ingredients.sort_by_key(|ingredient| cmp::Reverse(ingredient.size_of_fields()));
475-
memos.sort_by_key(|(_, memo)| cmp::Reverse(memo.size_of_fields()));
476-
477-
SalsaMemoryDump { ingredients, memos }
478-
}
479-
480-
struct SalsaMemoryDump {
481-
ingredients: Vec<salsa::IngredientInfo>,
482-
memos: Vec<(&'static str, salsa::IngredientInfo)>,
483-
}
484-
485-
#[allow(clippy::cast_precision_loss)]
486-
impl SalsaMemoryDump {
487-
fn display_short(&self) -> impl Display + '_ {
488-
struct DisplayShort<'a>(&'a SalsaMemoryDump);
489-
490-
impl std::fmt::Display for DisplayShort<'_> {
491-
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
492-
let mut total_fields = 0;
493-
let mut total_metadata = 0;
494-
for ingredient in &self.0.ingredients {
495-
total_metadata += ingredient.size_of_metadata();
496-
total_fields += ingredient.size_of_fields();
497-
}
498-
499-
let mut total_memo_fields = 0;
500-
let mut total_memo_metadata = 0;
501-
for (_, memo) in &self.0.memos {
502-
total_memo_fields += memo.size_of_fields();
503-
total_memo_metadata += memo.size_of_metadata();
504-
}
505-
506-
writeln!(f, "=======SALSA SUMMARY=======")?;
507-
508-
writeln!(
509-
f,
510-
"TOTAL MEMORY USAGE: {:.2}MB",
511-
(total_metadata + total_fields + total_memo_fields + total_memo_metadata)
512-
as f64
513-
/ 1_000_000.,
514-
)?;
515-
516-
writeln!(
517-
f,
518-
" struct metadata = {:.2}MB",
519-
total_metadata as f64 / 1_000_000.,
520-
)?;
521-
writeln!(
522-
f,
523-
" struct fields = {:.2}MB",
524-
total_fields as f64 / 1_000_000.,
525-
)?;
526-
writeln!(
527-
f,
528-
" memo metadata = {:.2}MB",
529-
total_memo_metadata as f64 / 1_000_000.,
530-
)?;
531-
writeln!(
532-
f,
533-
" memo fields = {:.2}MB",
534-
total_memo_fields as f64 / 1_000_000.
535-
)?;
536-
537-
writeln!(f, "QUERY COUNT: {}", self.0.memos.len())?;
538-
writeln!(f, "STRUCT COUNT: {}", self.0.ingredients.len())?;
539-
540-
Ok(())
541-
}
542-
}
543-
544-
DisplayShort(self)
545-
}
546-
547-
fn display_full(&self) -> impl Display + '_ {
548-
struct DisplayFull<'a>(&'a SalsaMemoryDump);
549-
550-
impl std::fmt::Display for DisplayFull<'_> {
551-
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
552-
writeln!(f, "=======SALSA STRUCTS=======")?;
553-
554-
let mut total_fields = 0;
555-
let mut total_metadata = 0;
556-
for ingredient in &self.0.ingredients {
557-
total_metadata += ingredient.size_of_metadata();
558-
total_fields += ingredient.size_of_fields();
559-
560-
writeln!(
561-
f,
562-
"{:<50} metadata={:<8} fields={:<8} count={}",
563-
format!("`{}`", ingredient.debug_name()),
564-
format!("{:.2}MB", ingredient.size_of_metadata() as f64 / 1_000_000.),
565-
format!("{:.2}MB", ingredient.size_of_fields() as f64 / 1_000_000.),
566-
ingredient.count()
567-
)?;
568-
}
569-
570-
writeln!(f, "=======SALSA QUERIES=======")?;
571-
572-
let mut total_memo_fields = 0;
573-
let mut total_memo_metadata = 0;
574-
for (query_fn, memo) in &self.0.memos {
575-
total_memo_fields += memo.size_of_fields();
576-
total_memo_metadata += memo.size_of_metadata();
577-
578-
writeln!(f, "`{query_fn} -> {}`", memo.debug_name())?;
579-
580-
writeln!(
581-
f,
582-
" metadata={:<8} fields={:<8} count={}",
583-
format!("{:.2}MB", memo.size_of_metadata() as f64 / 1_000_000.),
584-
format!("{:.2}MB", memo.size_of_fields() as f64 / 1_000_000.),
585-
memo.count()
586-
)?;
587-
}
588-
589-
writeln!(f, "=======SALSA SUMMARY=======")?;
590-
writeln!(
591-
f,
592-
"TOTAL MEMORY USAGE: {:.2}MB",
593-
(total_metadata + total_fields + total_memo_fields + total_memo_metadata)
594-
as f64
595-
/ 1_000_000.,
596-
)?;
597-
598-
writeln!(
599-
f,
600-
" struct metadata = {:.2}MB",
601-
total_metadata as f64 / 1_000_000.,
602-
)?;
603-
writeln!(
604-
f,
605-
" struct fields = {:.2}MB",
606-
total_fields as f64 / 1_000_000.,
607-
)?;
608-
writeln!(
609-
f,
610-
" memo metadata = {:.2}MB",
611-
total_memo_metadata as f64 / 1_000_000.,
612-
)?;
613-
writeln!(
614-
f,
615-
" memo fields = {:.2}MB",
616-
total_memo_fields as f64 / 1_000_000.
617-
)?;
618-
619-
Ok(())
620-
}
621-
}
622-
623-
DisplayFull(self)
624-
}
625-
}

crates/ty_project/src/db.rs

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
use std::panic::{AssertUnwindSafe, RefUnwindSafe};
22
use std::sync::Arc;
3+
use std::{cmp, fmt};
34

45
use crate::metadata::settings::file_settings;
56
use crate::{DEFAULT_LINT_REGISTRY, DummyReporter};
@@ -108,6 +109,171 @@ impl ProjectDatabase {
108109
Arc::get_mut(&mut self.system)
109110
.expect("ref count should be 1 because `zalsa_mut` drops all other DB references.")
110111
}
112+
113+
/// Returns a [`SalsaMemoryDump`] that can be use to dump Salsa memory usage information
114+
/// to the CLI after a typechecker run.
115+
pub fn salsa_memory_dump(&self) -> SalsaMemoryDump {
116+
let salsa_db = self as &dyn salsa::Database;
117+
118+
let mut ingredients = salsa_db.structs_info();
119+
let mut memos = salsa_db.queries_info().into_iter().collect::<Vec<_>>();
120+
121+
ingredients.sort_by_key(|ingredient| cmp::Reverse(ingredient.size_of_fields()));
122+
memos.sort_by_key(|(_, memo)| cmp::Reverse(memo.size_of_fields()));
123+
124+
SalsaMemoryDump { ingredients, memos }
125+
}
126+
}
127+
128+
/// Stores memory usage information.
129+
pub struct SalsaMemoryDump {
130+
ingredients: Vec<salsa::IngredientInfo>,
131+
memos: Vec<(&'static str, salsa::IngredientInfo)>,
132+
}
133+
134+
#[allow(clippy::cast_precision_loss)]
135+
impl SalsaMemoryDump {
136+
/// Returns a short report that provides total memory usage information.
137+
pub fn display_short(&self) -> impl fmt::Display + '_ {
138+
struct DisplayShort<'a>(&'a SalsaMemoryDump);
139+
140+
impl fmt::Display for DisplayShort<'_> {
141+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
142+
let mut total_fields = 0;
143+
let mut total_metadata = 0;
144+
for ingredient in &self.0.ingredients {
145+
total_metadata += ingredient.size_of_metadata();
146+
total_fields += ingredient.size_of_fields();
147+
}
148+
149+
let mut total_memo_fields = 0;
150+
let mut total_memo_metadata = 0;
151+
for (_, memo) in &self.0.memos {
152+
total_memo_fields += memo.size_of_fields();
153+
total_memo_metadata += memo.size_of_metadata();
154+
}
155+
156+
writeln!(f, "=======SALSA SUMMARY=======")?;
157+
158+
writeln!(
159+
f,
160+
"TOTAL MEMORY USAGE: {:.2}MB",
161+
(total_metadata + total_fields + total_memo_fields + total_memo_metadata)
162+
as f64
163+
/ 1_000_000.,
164+
)?;
165+
166+
writeln!(
167+
f,
168+
" struct metadata = {:.2}MB",
169+
total_metadata as f64 / 1_000_000.,
170+
)?;
171+
writeln!(
172+
f,
173+
" struct fields = {:.2}MB",
174+
total_fields as f64 / 1_000_000.,
175+
)?;
176+
writeln!(
177+
f,
178+
" memo metadata = {:.2}MB",
179+
total_memo_metadata as f64 / 1_000_000.,
180+
)?;
181+
writeln!(
182+
f,
183+
" memo fields = {:.2}MB",
184+
total_memo_fields as f64 / 1_000_000.
185+
)?;
186+
187+
writeln!(f, "QUERY COUNT: {}", self.0.memos.len())?;
188+
writeln!(f, "STRUCT COUNT: {}", self.0.ingredients.len())?;
189+
190+
Ok(())
191+
}
192+
}
193+
194+
DisplayShort(self)
195+
}
196+
197+
/// Returns a short report that provides fine-grained memory usage information per
198+
/// Salsa ingredient.
199+
pub fn display_full(&self) -> impl fmt::Display + '_ {
200+
struct DisplayFull<'a>(&'a SalsaMemoryDump);
201+
202+
impl fmt::Display for DisplayFull<'_> {
203+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
204+
writeln!(f, "=======SALSA STRUCTS=======")?;
205+
206+
let mut total_fields = 0;
207+
let mut total_metadata = 0;
208+
for ingredient in &self.0.ingredients {
209+
total_metadata += ingredient.size_of_metadata();
210+
total_fields += ingredient.size_of_fields();
211+
212+
writeln!(
213+
f,
214+
"{:<50} metadata={:<8} fields={:<8} count={}",
215+
format!("`{}`", ingredient.debug_name()),
216+
format!("{:.2}MB", ingredient.size_of_metadata() as f64 / 1_000_000.),
217+
format!("{:.2}MB", ingredient.size_of_fields() as f64 / 1_000_000.),
218+
ingredient.count()
219+
)?;
220+
}
221+
222+
writeln!(f, "=======SALSA QUERIES=======")?;
223+
224+
let mut total_memo_fields = 0;
225+
let mut total_memo_metadata = 0;
226+
for (query_fn, memo) in &self.0.memos {
227+
total_memo_fields += memo.size_of_fields();
228+
total_memo_metadata += memo.size_of_metadata();
229+
230+
writeln!(f, "`{query_fn} -> {}`", memo.debug_name())?;
231+
232+
writeln!(
233+
f,
234+
" metadata={:<8} fields={:<8} count={}",
235+
format!("{:.2}MB", memo.size_of_metadata() as f64 / 1_000_000.),
236+
format!("{:.2}MB", memo.size_of_fields() as f64 / 1_000_000.),
237+
memo.count()
238+
)?;
239+
}
240+
241+
writeln!(f, "=======SALSA SUMMARY=======")?;
242+
writeln!(
243+
f,
244+
"TOTAL MEMORY USAGE: {:.2}MB",
245+
(total_metadata + total_fields + total_memo_fields + total_memo_metadata)
246+
as f64
247+
/ 1_000_000.,
248+
)?;
249+
250+
writeln!(
251+
f,
252+
" struct metadata = {:.2}MB",
253+
total_metadata as f64 / 1_000_000.,
254+
)?;
255+
writeln!(
256+
f,
257+
" struct fields = {:.2}MB",
258+
total_fields as f64 / 1_000_000.,
259+
)?;
260+
writeln!(
261+
f,
262+
" memo metadata = {:.2}MB",
263+
total_memo_metadata as f64 / 1_000_000.,
264+
)?;
265+
writeln!(
266+
f,
267+
" memo fields = {:.2}MB",
268+
total_memo_fields as f64 / 1_000_000.
269+
)?;
270+
271+
Ok(())
272+
}
273+
}
274+
275+
DisplayFull(self)
276+
}
111277
}
112278

113279
impl Upcast<dyn SemanticDb> for ProjectDatabase {

crates/ty_python_semantic/src/semantic_index/use_def.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -681,7 +681,7 @@ impl<'db> Iterator for DeclarationsIterator<'_, 'db> {
681681

682682
impl std::iter::FusedIterator for DeclarationsIterator<'_, '_> {}
683683

684-
#[derive(Debug, PartialEq, Eq, salsa::Update)]
684+
#[derive(Debug, PartialEq, Eq, salsa::Update, get_size2::GetSize)]
685685
struct ReachableDefinitions {
686686
bindings: Bindings,
687687
declarations: Declarations,

0 commit comments

Comments
 (0)