Make bxl profiler configuration propagate the same as others

Summary:
Thist makes bxl profiling use the same mechanism for profiler
configuration that others do (i.e. putting the configuration on the dice graph
and looking it up).

Reviewed By: JakobDegen

Differential Revision: D64281581

fbshipit-source-id: e1264b05ab1256c9ce461948f1a00e385e29ec5a

app/buck2_bxl/src/bxl/calculation.rs

@@ -15,7 +15,6 @@ use buck2_build_api::bxl::calculation::BxlComputeResult;
 use buck2_build_api::bxl::calculation::BXL_CALCULATION_IMPL;
 use buck2_core::deferred::base_deferred_key::BaseDeferredKeyBxl;
 use buck2_futures::cancellation::CancellationContext;
-use buck2_interpreter::starlark_profiler::mode::StarlarkProfileMode;
 use dice::DiceComputations;
 use dice::Key;
 use dupe::Dupe;
@@ -65,7 +64,7 @@ impl Key for internal::BxlComputeKey {
         cancellation
             .with_structured_cancellation(|observer| {
                 async move {
-                    eval(ctx, key, StarlarkProfileMode::None, observer)
+                    eval(ctx, key, observer)
                         .await
                         .map_err(buck2_error::Error::from)
                         .map(|(result, _)| BxlComputeResult(Arc::new(result)))

app/buck2_bxl/src/bxl/eval.rs

@@ -43,11 +43,8 @@ use buck2_interpreter::load_module::InterpreterCalculation;
 use buck2_interpreter::paths::module::StarlarkModulePath;
 use buck2_interpreter::print_handler::EventDispatcherPrintHandler;
 use buck2_interpreter::soft_error::Buck2StarlarkSoftErrorHandler;
-use buck2_interpreter::starlark_profiler::data::ProfileTarget;
+use buck2_interpreter::starlark_profiler::config::GetStarlarkProfilerInstrumentation;
 use buck2_interpreter::starlark_profiler::data::StarlarkProfileDataAndStats;
-use buck2_interpreter::starlark_profiler::mode::StarlarkProfileMode;
-use buck2_interpreter::starlark_profiler::profiler::StarlarkProfiler;
-use buck2_interpreter::starlark_profiler::profiler::StarlarkProfilerOpt;
 use clap::error::ErrorKind;
 use dice::DiceComputations;
 use dice::DiceTransaction;
@@ -103,7 +100,6 @@ impl LimitedExecutor {
 pub(crate) async fn eval(
     ctx: &mut DiceComputations<'_>,
     key: BxlKey,
-    profile_mode_or_instrumentation: StarlarkProfileMode,
     liveness: CancellationObserver,
 ) -> buck2_error::Result<(BxlResult, Option<StarlarkProfileDataAndStats>)> {
     // Note: because we use `block_in_place`, that will prevent the inner future from being polled
@@ -126,13 +122,7 @@ pub(crate) async fn eval(
         // to terminate.
         scope_and_collect_with_dice(ctx, |ctx, s| {
             s.spawn_cancellable(
-                limited_executor.execute(eval_bxl_inner(
-                    ctx,
-                    dispatcher,
-                    key,
-                    profile_mode_or_instrumentation,
-                    liveness,
-                )),
+                limited_executor.execute(eval_bxl_inner(ctx, dispatcher, key, liveness)),
                 || Err(buck2_error!(buck2_error::ErrorTag::Tier0, "cancelled")),
             )
         })
@@ -277,7 +267,6 @@ async fn eval_bxl_inner(
     ctx: &mut DiceComputations<'_>,
     dispatcher: EventDispatcher,
     key: BxlKey,
-    profile_mode_or_instrumentation: StarlarkProfileMode,
     liveness: CancellationObserver,
 ) -> buck2_error::Result<(BxlResult, Option<StarlarkProfileDataAndStats>)> {
     let bxl_module = ctx
@@ -294,15 +283,6 @@ async fn eval_bxl_inner(
     // futures that requires work to be done on the current thread, so using block_in_place
     // should have no noticeable different compared to spawn_blocking
 
-    let mut profiler_opt = profile_mode_or_instrumentation
-        .profile_mode()
-        .map(|profile_mode| StarlarkProfiler::new(profile_mode.dupe(), ProfileTarget::Bxl));
-
-    let mut profiler = match &mut profiler_opt {
-        None => StarlarkProfilerOpt::disabled(),
-        Some(profiler) => StarlarkProfilerOpt::for_profiler(profiler),
-    };
-
     let eval_ctx = BxlInnerEvaluator {
         data: core_data,
         module: bxl_module,
@@ -311,15 +291,17 @@ async fn eval_bxl_inner(
         dispatcher,
     };
 
+    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(
         ctx,
-        &mut profiler,
-        &key.as_starlark_eval_kind(),
+        &mut profiler.as_mut(),
+        &eval_kind,
         move |provider, ctx| eval_ctx.do_eval(provider, ctx),
     )
     .await?;
 
-    let profile_data = profiler_opt.map(|p| p.finish()).transpose()?;
+    let profile_data = profiler.finish()?;
     Ok((bxl_result, profile_data))
 }
 

app/buck2_bxl/src/profile_command.rs

@@ -19,10 +19,9 @@ use buck2_core::fs::paths::abs_path::AbsPath;
 use buck2_error::buck2_error;
 use buck2_error::internal_error;
 use buck2_error::BuckErrorContext;
-use buck2_interpreter::starlark_profiler::config::StarlarkProfilerConfiguration;
+use buck2_interpreter::starlark_profiler::config::GetStarlarkProfilerInstrumentation;
 use buck2_interpreter::starlark_profiler::mode::StarlarkProfileMode;
 use buck2_profile::get_profile_response;
-use buck2_profile::starlark_profiler_configuration_from_request;
 use buck2_server_ctx::ctx::ServerCommandContextTrait;
 use buck2_server_ctx::global_cfg_options::global_cfg_options_from_client_context;
 use buck2_server_ctx::partial_result_dispatcher::NoPartialResult;
@@ -82,13 +81,6 @@ impl ServerCommandTemplate for BxlProfileServerCommand {
 
         let output = AbsPath::new(Path::new(&self.req.destination_path))?;
 
-        let profile_mode =
-            starlark_profiler_configuration_from_request(&self.req, server_ctx.project_root())?;
-
-        let StarlarkProfilerConfiguration::ProfileBxl(profile_mode) = profile_mode else {
-            return Err(internal_error!("Incorrect profile mode"));
-        };
-
         let cwd = server_ctx.working_dir();
 
         let cell_resolver = ctx.get_cell_resolver().await?;
@@ -122,21 +114,23 @@ impl ServerCommandTemplate for BxlProfileServerCommand {
             global_cfg_options,
         );
 
+        if let StarlarkProfileMode::None = ctx
+            .get_starlark_profiler_mode(&bxl_key.as_starlark_eval_kind())
+            .await?
+        {
+            return Err(internal_error!("Incorrect profile mode"));
+        }
+
         let profile_data = server_ctx
             .cancellation_context()
             .with_structured_cancellation(|observer| {
                 async move {
                     buck2_error::Ok(
-                        eval(
-                            &mut ctx,
-                            bxl_key,
-                            StarlarkProfileMode::Profile(profile_mode),
-                            observer,
-                        )
-                        .await?
-                        .1
-                        .map(Arc::new)
-                        .expect("No bxl profile data found"),
+                        eval(&mut ctx, bxl_key, observer)
+                            .await?
+                            .1
+                            .map(Arc::new)
+                            .expect("No bxl profile data found"),
                     )
                 }
                 .boxed()

app/buck2_interpreter/src/starlark_profiler/config.rs

@@ -205,7 +205,10 @@ impl ProjectionKey for StarlarkProfileModeForKind {
                     _ => Ok(StarlarkProfileMode::None),
                 }
             }
-            StarlarkProfilerConfigurationResolved::ProfileBxl(_) => Ok(StarlarkProfileMode::None),
+            StarlarkProfilerConfigurationResolved::ProfileBxl(mode) => match &self.0 {
+                StarlarkEvalKind::Bxl(..) => Ok(StarlarkProfileMode::Profile(mode.dupe())),
+                _ => Ok(StarlarkProfileMode::None),
+            },
         }
     }