From 3819f0ee6f5765c35bd3d201c49a4181f15f10b8 Mon Sep 17 00:00:00 2001
From: Songqian Li <mail@lisongqian.cn>
Date: Wed, 18 Oct 2023 14:03:18 +0800
Subject: [PATCH] dragonball: output balloon device metrics to runtime

Balloon device manager adds balloon device metrics to METRICS when a device is created and remove metrics when a device is dropped.

Fixes: #7248

Signed-off-by: Songqian Li <mail@lisongqian.cn>
---
 .../src/device_manager/balloon_dev_mgr.rs      | 18 ++++++++++++++++++
 src/dragonball/src/metric.rs                   |  5 +++++
 2 files changed, 23 insertions(+)

diff --git a/src/dragonball/src/device_manager/balloon_dev_mgr.rs b/src/dragonball/src/device_manager/balloon_dev_mgr.rs
index b0ee2bd37e..329966adf4 100644
--- a/src/dragonball/src/device_manager/balloon_dev_mgr.rs
+++ b/src/dragonball/src/device_manager/balloon_dev_mgr.rs
@@ -11,6 +11,7 @@ use crate::address_space_manager::GuestAddressSpaceImpl;
 use crate::config_manager::{ConfigItem, DeviceConfigInfo, DeviceConfigInfos};
 use crate::device_manager::DbsMmioV2Device;
 use crate::device_manager::{DeviceManager, DeviceMgrError, DeviceOpContext};
+use crate::metric::METRICS;
 
 // The flag of whether to use the shared irq.
 const USE_SHARED_IRQ: bool = true;
@@ -175,6 +176,11 @@ impl BalloonDeviceMgr {
                 )
                 .map_err(BalloonDeviceError::CreateBalloonDevice)?,
             );
+            METRICS
+                .write()
+                .unwrap()
+                .balloon
+                .insert(balloon_cfg.balloon_id.clone(), device.metrics());
 
             let mmio_dev =
                 DeviceManager::create_mmio_virtio_device_with_device_change_notification(
@@ -220,6 +226,11 @@ impl BalloonDeviceMgr {
                 },
             )
             .map_err(BalloonDeviceError::CreateBalloonDevice)?;
+            METRICS
+                .write()
+                .unwrap()
+                .balloon
+                .insert(info.config.balloon_id.clone(), device.metrics());
             let mmio_dev =
                 DeviceManager::create_mmio_virtio_device_with_device_change_notification(
                     Box::new(device),
@@ -275,6 +286,13 @@ impl Default for BalloonDeviceMgr {
     }
 }
 
+impl Drop for BalloonDeviceMgr {
+    // todo: move METIRCS oprations to remove_device. issue #8207.
+    fn drop(&mut self) {
+        METRICS.write().unwrap().balloon.clear();
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/src/dragonball/src/metric.rs b/src/dragonball/src/metric.rs
index 01e61395b1..f1fa46281f 100644
--- a/src/dragonball/src/metric.rs
+++ b/src/dragonball/src/metric.rs
@@ -6,6 +6,7 @@ use std::collections::HashMap;
 use std::sync::{Arc, RwLock};
 
 use dbs_utils::metric::SharedIncMetric;
+use dbs_virtio_devices::balloon::BalloonDeviceMetrics;
 use lazy_static::lazy_static;
 use serde::Serialize;
 
@@ -52,6 +53,8 @@ pub struct SignalMetrics {
 }
 
 /// Structure storing all metrics while enforcing serialization support on them.
+/// The type of the device metrics is HashMap<DeviceId, Arc<DeviceMetrics>> and the type of
+/// non-device metrics is XXMetrics.
 #[derive(Default, Serialize)]
 pub struct DragonballMetrics {
     /// Metrics related to a vcpu's functioning.
@@ -60,6 +63,8 @@ pub struct DragonballMetrics {
     pub seccomp: SeccompMetrics,
     /// Metrics related to signals.
     pub signals: SignalMetrics,
+    /// Metrics related to balloon device.
+    pub balloon: HashMap<String, Arc<BalloonDeviceMetrics>>,
 }
 
 #[cfg(test)]