dbs-cli: support virtio-balloon feature

We can use dbs-cli to simply reuse unused virtual machine on host. An
empty virtio-balloon device should be inserted first before we ask guest
to balloon some free memory. It should use update action, and only
parameter is size in mib.

Example:

First time:
./dbs-cli  --api-sock-path /tmp/api-sock --balloon-memory 0 update

Next:
./dbs-cli  --api-sock-path /tmp/api-sock --balloon-memory 512 update

Fixes: #23

Signed-off-by: Helin Guo <helinguo@linux.alibaba.com>

Author: HerlinCoder

Cargo.toml

@@ -11,6 +11,7 @@ dragonball = { git = "https://github.com/kata-containers/kata-containers", branc
     "virtio-vsock",
     "virtio-net",
     "virtio-mem",
+    "virtio-balloon",
     "hotplug",
     "dbs-upcall",
     "atomic-guest-memory"

src/api_client.rs

@@ -27,6 +27,11 @@ pub fn run_api_client(args: DBSArgs) -> Result<()> {
         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(())
 }
 
@@ -60,6 +65,13 @@ fn request_hotplug_memory(size_mib: usize) -> Value {
     })
 }
 
+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

@@ -10,6 +10,7 @@ use std::os::unix::net::{UnixListener, UnixStream};
 
 use anyhow::{anyhow, Context, Result};
 use dragonball::device_manager::blk_dev_mgr::BlockDeviceConfigInfo;
+use dragonball::device_manager::balloon_dev_mgr::BalloonDeviceConfigInfo;
 use dragonball::device_manager::mem_dev_mgr::MemDeviceConfigInfo;
 use dragonball::device_manager::virtio_net_dev_mgr::VirtioNetDeviceConfigInfo;
 
@@ -121,6 +122,20 @@ impl ApiServer {
                 };
                 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

@@ -305,4 +305,12 @@ The type of it is an array of BlockDeviceConfigInfo, e.g.
         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,7 @@ 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,
@@ -146,10 +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(())
+    }
 }