Support memory update in dbs-cli

Cargo.toml

@@ -10,8 +10,11 @@ dragonball = { git = "https://github.com/kata-containers/kata-containers", branc
     "virtio-blk",
     "virtio-vsock",
     "virtio-net",
+    "virtio-mem",
+    "virtio-balloon",
     "hotplug",
     "dbs-upcall",
+    "atomic-guest-memory"
 ] }
 clap = { version = "4.0.27", features = ["derive"] }
 serde = "1.0.27"

src/api_client.rs

@@ -22,6 +22,16 @@ pub fn run_api_client(args: DBSArgs) -> Result<()> {
         let request = request_virtio_blk(&config);
         send_request(request, &args.api_sock_path)?;
     }
+    if let Some(size_mib) = args.update_args.hotplug_memory {
+        let request = request_hotplug_memory(size_mib);
+        send_request(request, &args.api_sock_path)?;
+    }
+
+    if let Some(size_mib) = args.update_args.balloon_memory {
+        let request = request_balloon_memory(size_mib);
+        send_request(request, &args.api_sock_path)?;
+    }
+
     Ok(())
 }
 
@@ -48,6 +58,20 @@ fn request_virtio_blk(virtio_blk_config: &str) -> Value {
     })
 }
 
+fn request_hotplug_memory(size_mib: usize) -> Value {
+    json!({
+        "action": "hotplug_memory",
+        "size_mib": size_mib,
+    })
+}
+
+fn request_balloon_memory(size_mib: usize) -> Value {
+    json!({
+        "action": "balloon_memory",
+        "size_mib": size_mib,
+    })
+}
+
 fn send_request(request: Value, api_sock_path: &str) -> Result<()> {
     let mut unix_stream = UnixStream::connect(api_sock_path).context("Could not create stream")?;
 

src/api_server.rs

@@ -9,7 +9,9 @@ use std::io::prelude::*;
 use std::os::unix::net::{UnixListener, UnixStream};
 
 use anyhow::{anyhow, Context, Result};
+use dragonball::device_manager::balloon_dev_mgr::BalloonDeviceConfigInfo;
 use dragonball::device_manager::blk_dev_mgr::BlockDeviceConfigInfo;
+use dragonball::device_manager::mem_dev_mgr::MemDeviceConfigInfo;
 use dragonball::device_manager::virtio_net_dev_mgr::VirtioNetDeviceConfigInfo;
 
 use crate::vmm_comm_trait::VMMComm;
@@ -105,6 +107,35 @@ impl ApiServer {
                         .context("Insert a virtio-blk device to the Dragonball")?;
                 }
             }
+            Some("hotplug_memory") => {
+                let mem_cfg = MemDeviceConfigInfo {
+                    mem_id: "virtio-mem0".to_string(),
+                    capacity_mib: 0,
+                    size_mib: v["size_mib"]
+                        .as_u64()
+                        .context("Invalid virtio-mem size input")?,
+                    multi_region: true,
+                    host_numa_node_id: None,
+                    guest_numa_node_id: None,
+                    use_generic_irq: None,
+                    use_shared_irq: None,
+                };
+                return self.insert_mem_device(mem_cfg);
+            }
+            Some("balloon_memory") => {
+                let balloon_cfg = BalloonDeviceConfigInfo {
+                    balloon_id: "virtio-balloon0".to_string(),
+                    size_mib: v["size_mib"]
+                        .as_u64()
+                        .context("Invalid virtio-balloon size input")?,
+                    use_generic_irq: None,
+                    use_shared_irq: None,
+                    f_deflate_on_oom: false,
+                    // TODO: enable free page reporting
+                    f_reporting: false,
+                };
+                return self.insert_balloon_device(balloon_cfg);
+            }
             _ => {
                 println!("Unknown Actions");
             }

src/parser/args.rs

@@ -297,4 +297,20 @@ The type of it is an array of BlockDeviceConfigInfo, e.g.
         display_order = 2
     )]
     pub hotplug_virblks: Option<String>,
+
+    #[clap(
+        long,
+        value_parser,
+        help = "Hotplug memory (MiB) through connection with dbs-cli api server",
+        display_order = 2
+    )]
+    pub hotplug_memory: Option<usize>,
+
+    #[clap(
+        long,
+        value_parser,
+        help = "Balloon memory (MiB) through connection with dbs-cli api server",
+        display_order = 2
+    )]
+    pub balloon_memory: Option<usize>,
 }

src/vmm_comm_trait.rs

@@ -9,6 +9,8 @@ use dragonball::{
         BlockDeviceConfigInfo, BootSourceConfig, VmmAction, VmmActionError, VmmData, VmmRequest,
         VmmResponse, VsockDeviceConfigInfo,
     },
+    device_manager::balloon_dev_mgr::BalloonDeviceConfigInfo,
+    device_manager::mem_dev_mgr::MemDeviceConfigInfo,
     device_manager::virtio_net_dev_mgr::VirtioNetDeviceConfigInfo,
     vcpu::VcpuResizeInfo,
     vm::VmConfigInfo,
@@ -145,4 +147,18 @@ pub trait VMMComm {
             .with_context(|| format!("Failed to insert virtio-blk device {:?}", blk_cfg))?;
         Ok(())
     }
+
+    fn insert_mem_device(&self, mem_cfg: MemDeviceConfigInfo) -> Result<()> {
+        self.handle_request_with_retry(Request::Sync(VmmAction::InsertMemDevice(mem_cfg.clone())))
+            .with_context(|| format!("Failed to insert mem device {:?}", mem_cfg))?;
+        Ok(())
+    }
+
+    fn insert_balloon_device(&self, balloon_cfg: BalloonDeviceConfigInfo) -> Result<()> {
+        self.handle_request_with_retry(Request::Sync(VmmAction::InsertBalloonDevice(
+            balloon_cfg.clone(),
+        )))
+        .with_context(|| format!("Failed to insert balloon device {:?}", balloon_cfg))?;
+        Ok(())
+    }
 }