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

@@ -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,

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<BxlResult> {
+    ) -> 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,15 +296,15 @@ 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,
         move |provider, ctx| eval_ctx.do_eval(provider, ctx),
     )
     .await?;
 
-    let profile_data = profiler.finish()?;
+    let profile_data = profiler.finish(Some(&frozen_module))?;
     Ok((bxl_result, profile_data))
 }
 

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<Target = DiceComputatio
         fn evaluation_complete(&mut self, eval: &mut Evaluator) -> 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)
-        }
     }
 
     {

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(())
-    }
 }

app/buck2_interpreter/src/starlark_profiler/profiler.rs

@@ -35,7 +35,6 @@ pub struct StarlarkProfiler {
     initialized_at: Option<Instant>,
     finalized_at: Option<Instant>,
     profile_data: Option<ProfileData>,
-    total_retained_bytes: Option<usize>,
 
     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<StarlarkProfileDataAndStats> {
+    pub fn finish(
+        mut self,
+        frozen_module: Option<&FrozenModule>,
+    ) -> buck2_error::Result<StarlarkProfileDataAndStats> {
+        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<Option<StarlarkProfileDataAndStats>> {
+    pub fn finish(
+        self,
+        frozen_module: Option<&FrozenModule>,
+    ) -> buck2_error::Result<Option<StarlarkProfileDataAndStats>> {
         match self {
             StarlarkProfilerOptVal::Disabled => Ok(None),
-            StarlarkProfilerOptVal::Profiler(profiler) => profiler.finish().map(Some),
+            StarlarkProfilerOptVal::Profiler(profiler) => profiler.finish(frozen_module).map(Some),
         }
     }
 }

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(),

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
                 }