diff --git a/app/buck2_server/src/daemon/server.rs b/app/buck2_server/src/daemon/server.rs index 5b8c471d2d17e..2502e6ab01ac8 100644 --- a/app/buck2_server/src/daemon/server.rs +++ b/app/buck2_server/src/daemon/server.rs @@ -830,7 +830,8 @@ impl DaemonApi for BuckdServer { data.materializer.dupe(), data.scribe_sink.dupe() as _, ) - .create_snapshot(), + .create_snapshot() + .await, ) } else { None diff --git a/app/buck2_server/src/heartbeat_guard.rs b/app/buck2_server/src/heartbeat_guard.rs index d92446647404c..be2d0237482f7 100644 --- a/app/buck2_server/src/heartbeat_guard.rs +++ b/app/buck2_server/src/heartbeat_guard.rs @@ -46,7 +46,7 @@ impl HeartbeatGuard { let mut interval = tokio::time::interval(Duration::from_secs(1)); interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); loop { - let snapshot = collector.create_snapshot(); + let snapshot = collector.create_snapshot().await; match events.lock().expect("Poisoned lock").as_ref() { Some(events) => events.instant_event(Box::new(snapshot)), None => break, @@ -70,7 +70,10 @@ impl Drop for HeartbeatGuard { // Synchronously remove access for sending new heartbeats. if let Some(events) = maybe_events.take() { // Send one last snapshot. - events.instant_event(Box::new(self.collector.create_snapshot())); + let collector = self.collector.dupe(); + tokio::spawn(async move { + events.instant_event(Box::new(collector.create_snapshot().await)); + }); } // Cancel the task as well. self.handle.abort(); diff --git a/app/buck2_server/src/snapshot.rs b/app/buck2_server/src/snapshot.rs index 1218ac09421b1..7ca7e88238b04 100644 --- a/app/buck2_server/src/snapshot.rs +++ b/app/buck2_server/src/snapshot.rs @@ -62,12 +62,12 @@ impl SnapshotCollector { } /// Create a new Snapshot. - pub fn create_snapshot(&self) -> buck2_data::Snapshot { + pub async fn create_snapshot(&self) -> buck2_data::Snapshot { let mut snapshot = Self::pre_initialization_snapshot(self.daemon_start_time); self.add_daemon_metrics(&mut snapshot); self.add_re_metrics(&mut snapshot); self.add_io_metrics(&mut snapshot); - self.add_dice_metrics(&mut snapshot); + self.add_dice_metrics(&mut snapshot).await; self.add_materializer_metrics(&mut snapshot); self.add_sink_metrics(&mut snapshot); snapshot @@ -156,8 +156,8 @@ impl SnapshotCollector { } } - fn add_dice_metrics(&self, snapshot: &mut buck2_data::Snapshot) { - let metrics = self.dice.metrics(); + async fn add_dice_metrics(&self, snapshot: &mut buck2_data::Snapshot) { + let metrics = self.dice.metrics().await; snapshot.dice_key_count = metrics.key_count as u64; snapshot.dice_currently_running_key_count = metrics.currently_running_key_count as u64; snapshot.dice_active_transaction_count = metrics.active_transaction_count; diff --git a/dice/dice/src/api/dice.rs b/dice/dice/src/api/dice.rs index 9a55acf104bd9..2de1a50234dab 100644 --- a/dice/dice/src/api/dice.rs +++ b/dice/dice/src/api/dice.rs @@ -253,8 +253,8 @@ impl Dice { self.implementation.detect_cycles() } - pub fn metrics(&self) -> Metrics { - self.implementation.metrics() + pub async fn metrics(&self) -> Metrics { + self.implementation.metrics().await } /// Wait until all active versions have exited. diff --git a/dice/dice/src/lib.rs b/dice/dice/src/lib.rs index f151578329e52..83e517b21b11e 100644 --- a/dice/dice/src/lib.rs +++ b/dice/dice/src/lib.rs @@ -316,7 +316,7 @@ impl DiceImplementation { } } - pub fn metrics(&self) -> Metrics { + pub async fn metrics(&self) -> Metrics { match self { DiceImplementation::Legacy(dice) => dice.metrics(), DiceImplementation::Modern(dice) => dice.metrics(),