From 5bc86d21bfdbd71fde20fc7bf2ba2578c086dfeb Mon Sep 17 00:00:00 2001 From: Chris Hopman Date: Tue, 21 Jan 2025 15:55:02 -0800 Subject: [PATCH] Refactor handling of frozen module Summary: Do everything in profiler.finish() instead of split across two functions. Reviewed By: JakobDegen Differential Revision: D64281587 fbshipit-source-id: 93802f846f8b68e245863c8a7523cdbd41a289f6 --- app/buck2_analysis/src/analysis/env.rs | 7 +- app/buck2_bxl/src/bxl/eval.rs | 13 ++-- .../src/dice/starlark_provider.rs | 8 --- app/buck2_interpreter/src/factory.rs | 7 -- .../src/starlark_profiler/profiler.rs | 71 ++++++++----------- .../interpreter/dice_calculation_delegate.rs | 2 +- .../src/interpreter/interpreter_for_cell.rs | 3 - 7 files changed, 38 insertions(+), 73 deletions(-) diff --git a/app/buck2_analysis/src/analysis/env.rs b/app/buck2_analysis/src/analysis/env.rs index a84ae67632bb..d5ccc4794485 100644 --- a/app/buck2_analysis/src/analysis/env.rs +++ b/app/buck2_analysis/src/analysis/env.rs @@ -332,12 +332,7 @@ async fn run_analysis_with_env_underlying( let frozen_env = env.freeze().map_err(from_freeze_error)?; let recorded_values = registry_finalizer(&frozen_env)?; - profiler - .as_mut() - .visit_frozen_module(Some(&frozen_env)) - .buck_error_context("Profiler heap visitation failed")?; - - let profile_data = profiler.finish()?.map(Arc::new); + let profile_data = profiler.finish(Some(&frozen_env))?.map(Arc::new); let validations = transitive_validations( validations_from_deps, diff --git a/app/buck2_bxl/src/bxl/eval.rs b/app/buck2_bxl/src/bxl/eval.rs index 88f444b0829e..40c979206a1e 100644 --- a/app/buck2_bxl/src/bxl/eval.rs +++ b/app/buck2_bxl/src/bxl/eval.rs @@ -52,6 +52,7 @@ use dice::DiceTransaction; use dupe::Dupe; use itertools::Itertools; use once_cell::sync::Lazy; +use starlark::environment::FrozenModule; use starlark::environment::Module; use starlark::eval::Evaluator; use starlark::values::structs::AllocStruct; @@ -149,7 +150,7 @@ impl BxlInnerEvaluator { self, provider: &mut dyn StarlarkEvaluatorProvider, dice: &'a mut DiceComputations, - ) -> buck2_error::Result { + ) -> buck2_error::Result<(FrozenModule, BxlResult)> { let BxlInnerEvaluator { data, module, @@ -261,11 +262,7 @@ impl BxlInnerEvaluator { recorded_values, ); - provider - .visit_frozen_module(Some(&frozen_module)) - .buck_error_context("Profiler heap visitation failed")?; - - Ok(bxl_result) + Ok((frozen_module, bxl_result)) } } @@ -299,7 +296,7 @@ async fn eval_bxl_inner( let eval_kind = key.as_starlark_eval_kind(); let mut profiler = ctx.get_starlark_profiler(&eval_kind).await?; - let bxl_result = with_starlark_eval_provider( + let (frozen_module, bxl_result) = with_starlark_eval_provider( ctx, &mut profiler.as_mut(), &eval_kind, @@ -307,7 +304,7 @@ async fn eval_bxl_inner( ) .await?; - let profile_data = profiler.finish()?; + let profile_data = profiler.finish(Some(&frozen_module))?; Ok((bxl_result, profile_data)) } diff --git a/app/buck2_interpreter/src/dice/starlark_provider.rs b/app/buck2_interpreter/src/dice/starlark_provider.rs index 82f8d094b671..f2e0efb64696 100644 --- a/app/buck2_interpreter/src/dice/starlark_provider.rs +++ b/app/buck2_interpreter/src/dice/starlark_provider.rs @@ -26,7 +26,6 @@ use buck2_error::internal_error; use buck2_util::arc_str::ThinArcStr; use dice::DiceComputations; use dupe::Dupe; -use starlark::environment::FrozenModule; use starlark::environment::Module; use starlark::eval::Evaluator; @@ -181,13 +180,6 @@ pub async fn with_starlark_eval_provider<'a, D: DerefMut buck2_error::Result<()> { self.profiler.evaluation_complete(eval) } - - fn visit_frozen_module( - &mut self, - module: Option<&FrozenModule>, - ) -> buck2_error::Result<()> { - self.profiler.visit_frozen_module(module) - } } { diff --git a/app/buck2_interpreter/src/factory.rs b/app/buck2_interpreter/src/factory.rs index 8648d6380b6d..a8d643e2a8d3 100644 --- a/app/buck2_interpreter/src/factory.rs +++ b/app/buck2_interpreter/src/factory.rs @@ -7,7 +7,6 @@ * of this source tree. */ -use starlark::environment::FrozenModule; use starlark::environment::Module; use starlark::eval::Evaluator; @@ -21,8 +20,6 @@ pub trait StarlarkEvaluatorProvider { ) -> buck2_error::Result<(Evaluator<'v, 'a, 'e>, bool)>; fn evaluation_complete(&mut self, eval: &mut Evaluator) -> buck2_error::Result<()>; - - fn visit_frozen_module(&mut self, module: Option<&FrozenModule>) -> buck2_error::Result<()>; } /// Trivial provider that just constructs an Evaluator. Useful for tests (but not necessarily limited to them). @@ -39,8 +36,4 @@ impl StarlarkEvaluatorProvider for StarlarkPassthroughProvider { fn evaluation_complete(&mut self, _eval: &mut Evaluator) -> buck2_error::Result<()> { Ok(()) } - - fn visit_frozen_module(&mut self, _module: Option<&FrozenModule>) -> buck2_error::Result<()> { - Ok(()) - } } diff --git a/app/buck2_interpreter/src/starlark_profiler/profiler.rs b/app/buck2_interpreter/src/starlark_profiler/profiler.rs index a7683c339fb7..35c37e07c3a4 100644 --- a/app/buck2_interpreter/src/starlark_profiler/profiler.rs +++ b/app/buck2_interpreter/src/starlark_profiler/profiler.rs @@ -35,7 +35,6 @@ pub struct StarlarkProfiler { initialized_at: Option, finalized_at: Option, profile_data: Option, - total_retained_bytes: Option, target: ProfileTarget, } @@ -47,19 +46,40 @@ impl StarlarkProfiler { initialized_at: None, finalized_at: None, profile_data: None, - total_retained_bytes: None, target, } } /// Collect all profiling data. - pub fn finish(self) -> buck2_error::Result { + pub fn finish( + mut self, + frozen_module: Option<&FrozenModule>, + ) -> buck2_error::Result { + let total_retained_bytes = match (frozen_module, self.profile_mode.requires_frozen_module()) + { + (None, true) => { + return Err(StarlarkProfilerError::RetainedMemoryNotFrozen.into()); + } + (Some(module), requires_frozen) => { + if requires_frozen { + let profile = module + .heap_profile() + .map_err(|e| from_any_with_tag(e, buck2_error::ErrorTag::Tier0))?; + self.profile_data = Some(profile); + } + + module + .frozen_heap() + .allocated_summary() + .total_allocated_bytes() + } + _ => 0, + }; + Ok(StarlarkProfileDataAndStats { initialized_at: self.initialized_at.internal_error("did not initialize")?, finalized_at: self.finalized_at.internal_error("did not finalize")?, - total_retained_bytes: self - .total_retained_bytes - .internal_error("did not visit heap")?, + total_retained_bytes, profile_data: self .profile_data .internal_error("profile_data not initialized")?, @@ -83,27 +103,6 @@ impl StarlarkProfiler { } Ok(()) } - - fn visit_frozen_module(&mut self, module: Option<&FrozenModule>) -> buck2_error::Result<()> { - if self.profile_mode.requires_frozen_module() { - let module = module.ok_or(StarlarkProfilerError::RetainedMemoryNotFrozen)?; - let profile = module - .heap_profile() - .map_err(|e| from_any_with_tag(e, buck2_error::ErrorTag::Tier0))?; - self.profile_data = Some(profile); - } - - let total_retained_bytes = module.map_or(0, |module| { - module - .frozen_heap() - .allocated_summary() - .total_allocated_bytes() - }); - - self.total_retained_bytes = Some(total_retained_bytes); - - Ok(()) - } } enum StarlarkProfilerOptImpl<'p> { @@ -132,17 +131,6 @@ impl<'p> StarlarkProfilerOpt<'p> { } } - pub fn visit_frozen_module( - &mut self, - module: Option<&FrozenModule>, - ) -> buck2_error::Result<()> { - if let StarlarkProfilerOptImpl::Profiler(profiler) = &mut self.0 { - profiler.visit_frozen_module(module) - } else { - Ok(()) - } - } - pub fn evaluation_complete(&mut self, eval: &mut Evaluator) -> buck2_error::Result<()> { if let StarlarkProfilerOptImpl::Profiler(profiler) = &mut self.0 { profiler.evaluation_complete(eval) @@ -167,10 +155,13 @@ impl StarlarkProfilerOptVal { } } - pub fn finish(self) -> buck2_error::Result> { + pub fn finish( + self, + frozen_module: Option<&FrozenModule>, + ) -> buck2_error::Result> { match self { StarlarkProfilerOptVal::Disabled => Ok(None), - StarlarkProfilerOptVal::Profiler(profiler) => profiler.finish().map(Some), + StarlarkProfilerOptVal::Profiler(profiler) => profiler.finish(frozen_module).map(Some), } } } diff --git a/app/buck2_interpreter_for_build/src/interpreter/dice_calculation_delegate.rs b/app/buck2_interpreter_for_build/src/interpreter/dice_calculation_delegate.rs index 86c0b180803b..df0ae61f6dd2 100644 --- a/app/buck2_interpreter_for_build/src/interpreter/dice_calculation_delegate.rs +++ b/app/buck2_interpreter_for_build/src/interpreter/dice_calculation_delegate.rs @@ -578,7 +578,7 @@ impl<'c, 'd: 'c> DiceCalculationDelegate<'c, 'd> { }, ) .await?; - let profile_data = profiler.finish()?; + let profile_data = profiler.finish(None)?; if eval_result.starlark_profile.is_some() { return ( now.unwrap().elapsed(), diff --git a/app/buck2_interpreter_for_build/src/interpreter/interpreter_for_cell.rs b/app/buck2_interpreter_for_build/src/interpreter/interpreter_for_cell.rs index 3046739e7424..73d7d40f65cb 100644 --- a/app/buck2_interpreter_for_build/src/interpreter/interpreter_for_cell.rs +++ b/app/buck2_interpreter_for_build/src/interpreter/interpreter_for_cell.rs @@ -528,9 +528,6 @@ impl InterpreterForCell { eval_provider .evaluation_complete(&mut eval) .buck_error_context("Profiler finalization failed")?; - eval_provider - .visit_frozen_module(None) - .buck_error_context("Profiler heap visitation failed")?; cpu_instruction_count }