diff --git a/app/buck2_bxl/src/bxl/calculation.rs b/app/buck2_bxl/src/bxl/calculation.rs index 7d414bd52c67..1a39c8a4a307 100644 --- a/app/buck2_bxl/src/bxl/calculation.rs +++ b/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))) diff --git a/app/buck2_bxl/src/bxl/eval.rs b/app/buck2_bxl/src/bxl/eval.rs index 69d1018f8793..dfe5e6c75955 100644 --- a/app/buck2_bxl/src/bxl/eval.rs +++ b/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)> { // 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)> { 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)) } diff --git a/app/buck2_bxl/src/profile_command.rs b/app/buck2_bxl/src/profile_command.rs index 18042825cc77..4d272f924ff3 100644 --- a/app/buck2_bxl/src/profile_command.rs +++ b/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() diff --git a/app/buck2_interpreter/src/starlark_profiler/config.rs b/app/buck2_interpreter/src/starlark_profiler/config.rs index 6600cbf52f0b..3af80fbf5875 100644 --- a/app/buck2_interpreter/src/starlark_profiler/config.rs +++ b/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), + }, } }