Check system_memory from recorded event not from host machine

Summary:
Currently, `buck2 replay` wouldn't display memory warnings accurately because
-  we were getting system memory from the host machine but not from the actual machine we built on.
- we didn't pass the buckconfig to replay to pick threshold from.

With D58812046,  we store the actual system memory as an instant event. This diff ensures that we use that memory (instead of host's memory) to display warnings via using the host's buckconfig. This is not perfect as we don't store the threshold from the user's machine. This could be improved later but I don't think buckconfig should change that much tbh so not prioritizing it now.

Reviewed By: stepancheg

Differential Revision: D58812023

fbshipit-source-id: d53ebe7d8cc5e695a4a0409a460e3847b71862d4

app/buck2_client/src/commands/log/replay.rs

@@ -17,7 +17,6 @@ use buck2_client_ctx::replayer::Replayer;
 use buck2_client_ctx::signal_handler::with_simple_sigint_handler;
 use buck2_client_ctx::subscribers::get::get_console_with_root;
 use buck2_client_ctx::subscribers::subscribers::EventSubscribers;
-use buck2_common::init::SystemWarningConfig;
 
 use crate::commands::log::options::EventLogOptions;
 
@@ -63,7 +62,7 @@ impl ReplayCommand {
                 let (replayer, invocation) =
                     Replayer::new(event_log.get(&ctx).await?, speed, preload).await?;
 
-                let system_warning_config = SystemWarningConfig::default();
+                let system_warning_config = ctx.system_warning_config()?;
                 let console = get_console_with_root(
                     invocation.trace_id,
                     console_opts.console_type,

app/buck2_client_ctx/src/client_ctx.rs

@@ -15,6 +15,7 @@ use buck2_cli_proto::client_context::HostPlatformOverride as GrpcHostPlatformOve
 use buck2_cli_proto::client_context::PreemptibleWhen as GrpcPreemptibleWhen;
 use buck2_cli_proto::ClientContext;
 use buck2_common::argv::Argv;
+use buck2_common::init::SystemWarningConfig;
 use buck2_common::invocation_paths::InvocationPaths;
 use buck2_core::error::buck2_hard_error_env;
 use buck2_core::fs::working_dir::WorkingDir;
@@ -214,4 +215,12 @@ impl<'a> ClientCommandContext<'a> {
     pub fn allow_vpnless(&self) -> anyhow::Result<bool> {
         Ok(self.immediate_config.daemon_startup_config()?.allow_vpnless)
     }
+
+    pub fn system_warning_config(&self) -> anyhow::Result<SystemWarningConfig> {
+        Ok(self
+            .immediate_config
+            .daemon_startup_config()?
+            .system_warning_config
+            .clone())
+    }
 }

app/buck2_client_ctx/src/subscribers/simpleconsole.rs

@@ -540,6 +540,7 @@ where
 
                     if let Some(memory_pressure) = check_memory_pressure(
                         &self.observer().two_snapshots().last,
+                        &self.observer().system_info(),
                         self.system_warning_config.memory_pressure_threshold_percent,
                     ) {
                         echo!("{}", system_memory_exceeded_msg(&memory_pressure))?;

app/buck2_client_ctx/src/subscribers/superconsole.rs

@@ -175,10 +175,12 @@ impl<'s> Component for BuckRootComponent<'s> {
         let mut draw = DrawVertical::new(dimensions);
 
         let last_snapshot_tuple = &self.state.simple_console.observer.two_snapshots().last;
+        let system_info = &self.state.simple_console.observer.system_info();
         {
             draw.draw(
                 &SystemWarningComponent {
                     last_snapshot_tuple,
+                    system_info,
                     system_warning_config: &self.state.config.system_warning_config,
                 },
                 mode,

app/buck2_client_ctx/src/subscribers/superconsole/system_warning.rs

@@ -22,6 +22,7 @@ use crate::subscribers::system_warning::system_memory_exceeded_msg;
 /// This component is used to display system warnings for a command e.g. memory pressure, low disk space etc.
 pub(crate) struct SystemWarningComponent<'a, T> {
     pub(crate) last_snapshot_tuple: &'a Option<(T, buck2_data::Snapshot)>,
+    pub(crate) system_info: &'a buck2_data::SystemInfo,
     pub(crate) system_warning_config: &'a SystemWarningConfig,
 }
 
@@ -37,6 +38,7 @@ impl<'a, T> Component for SystemWarningComponent<'a, T> {
 
         if let Some(memory_pressure) = check_memory_pressure(
             self.last_snapshot_tuple,
+            self.system_info,
             self.system_warning_config.memory_pressure_threshold_percent,
         ) {
             lines.push(warning_styled(&system_memory_exceeded_msg(

app/buck2_client_ctx/src/subscribers/system_warning.rs

@@ -9,7 +9,6 @@
 
 use buck2_core::is_open_source;
 use buck2_event_observer::humanized::HumanizedBytes;
-use buck2_util::system_stats::system_memory_stats;
 
 use crate::subscribers::recorder::process_memory;
 
@@ -34,13 +33,14 @@ pub(crate) fn system_memory_exceeded_msg(memory_pressure: &MemoryPressureHigh) -
 
 pub(crate) fn check_memory_pressure<T>(
     snapshot_tuple: &Option<(T, buck2_data::Snapshot)>,
+    system_info: &buck2_data::SystemInfo,
     memory_pressure_threshold_percent: Option<u64>,
 ) -> Option<MemoryPressureHigh> {
+    // TODO (ezgi): use the recorded threshold, not the one from host's buckconfig.
     memory_pressure_threshold_percent.and_then(|memory_pressure_threshold_percent| {
         snapshot_tuple.as_ref().and_then(|(_, snapshot)| {
             process_memory(snapshot).and_then(|process_memory| {
-                // TODO(ezgi): We should check the recorded system memory, not the real one so that replay would show the warnings properly.
-                let system_total_memory = system_memory_stats();
+                let system_total_memory = system_info.system_total_memory_bytes;
                 if (process_memory * 100 / system_total_memory) >= memory_pressure_threshold_percent
                 {
                     Some(MemoryPressureHigh {

app/buck2_event_observer/src/event_observer.rs

@@ -30,6 +30,7 @@ pub struct EventObserver<E> {
     pub action_stats: ActionStats,
     re_state: ReState,
     two_snapshots: TwoSnapshots, // NOTE: We got many more copies of this than we should.
+    system_info: buck2_data::SystemInfo,
     session_info: SessionInfo,
     test_state: TestState,
     starlark_debugger_state: StarlarkDebuggerState,
@@ -48,6 +49,7 @@ where
             action_stats: ActionStats::default(),
             re_state: ReState::new(),
             two_snapshots: TwoSnapshots::default(),
+            system_info: buck2_data::SystemInfo::default(),
             session_info: SessionInfo {
                 trace_id,
                 test_session: None,
@@ -119,6 +121,9 @@ where
                                 self.session_info.modern_dice = true;
                             }
                         }
+                        SystemInfo(system_info) => {
+                            self.system_info = system_info.clone();
+                        }
                         _ => {}
                     }
                 }
@@ -147,6 +152,10 @@ where
         &self.two_snapshots
     }
 
+    pub fn system_info(&self) -> &buck2_data::SystemInfo {
+        &self.system_info
+    }
+
     pub fn session_info(&self) -> &SessionInfo {
         &self.session_info
     }