main: Add inflate/deflate interface for balloon

Change-Id: I0fc63abbed8db303c7d283ce392fd47777b60d19
Reviewed-on: https://chromium-review.googlesource.com/818207
Commit-Ready: Dylan Reid <dgreid@chromium.org>
Tested-by: Dylan Reid <dgreid@chromium.org>
Reviewed-by: Zach Reizner <zachr@chromium.org>

Cargo.lock

@@ -132,6 +132,7 @@ dependencies = [
 name = "vm_control"
 version = "0.1.0"
 dependencies = [
+ "byteorder 1.1.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "data_model 0.1.0",
  "kvm 0.1.0",
  "libc 0.2.32 (registry+https://github.com/rust-lang/crates.io-index)",

src/linux.rs

@@ -253,7 +253,7 @@ pub fn run_config(cfg: Config) -> Result<()> {
     device_manager.register_mmio(rng_box, rng_jail, &mut cmdline)
         .map_err(Error::RegisterRng)?;
 
-    let (_, balloon_device_socket) = UnixDatagram::pair().map_err(Error::Socket)?;
+    let (balloon_host_socket, balloon_device_socket) = UnixDatagram::pair().map_err(Error::Socket)?;
     let balloon_box = Box::new(devices::virtio::Balloon::new(balloon_device_socket)
                                    .map_err(Error::BalloonDeviceNew)?);
     let balloon_jail = if cfg.multiprocess {
@@ -387,7 +387,8 @@ pub fn run_config(cfg: Config) -> Result<()> {
             cfg.vcpu_count.unwrap_or(1),
             guest_mem,
             &device_manager.bus,
-            control_sockets)
+            control_sockets,
+            balloon_host_socket)
 }
 
 fn run_kvm(requests: Vec<VmRequest>,
@@ -396,7 +397,8 @@ fn run_kvm(requests: Vec<VmRequest>,
            vcpu_count: u32,
            guest_mem: GuestMemory,
            mmio_bus: &devices::Bus,
-           control_sockets: Vec<UnlinkUnixDatagram>)
+           control_sockets: Vec<UnlinkUnixDatagram>,
+           balloon_host_socket: UnixDatagram)
            -> Result<()> {
     let kvm = Kvm::new().map_err(Error::Kvm)?;
     let kernel_start_addr = GuestAddress(KERNEL_START_OFFSET);
@@ -414,7 +416,8 @@ fn run_kvm(requests: Vec<VmRequest>,
     let mut next_dev_pfn = BASE_DEV_MEMORY_PFN;
     for request in requests {
         let mut running = false;
-        if let VmResponse::Err(e) = request.execute(&mut vm, &mut next_dev_pfn, &mut running) {
+        if let VmResponse::Err(e) = request.execute(&mut vm, &mut next_dev_pfn,
+                                                    &mut running, &balloon_host_socket) {
             return Err(Error::Vm(e));
         }
         if !running {
@@ -573,7 +576,8 @@ fn run_kvm(requests: Vec<VmRequest>,
                 exit_evt,
                 sigchld_fd,
                 kill_signaled,
-                vcpu_handles)
+                vcpu_handles,
+                balloon_host_socket)
 }
 
 fn run_control(mut vm: Vm,
@@ -583,7 +587,8 @@ fn run_control(mut vm: Vm,
                exit_evt: EventFd,
                sigchld_fd: SignalFd,
                kill_signaled: Arc<AtomicBool>,
-               vcpu_handles: Vec<JoinHandle<()>>)
+               vcpu_handles: Vec<JoinHandle<()>>,
+               balloon_host_socket: UnixDatagram)
                -> Result<()> {
     const MAX_VM_FD_RECV: usize = 1;
 
@@ -665,7 +670,8 @@ fn run_control(mut vm: Vm,
                         Ok(request) => {
                             let mut running = true;
                             let response =
-                                request.execute(&mut vm, &mut next_dev_pfn, &mut running);
+                                request.execute(&mut vm, &mut next_dev_pfn,
+                                                &mut running, &balloon_host_socket);
                             if let Err(e) = response.send(&mut scm, socket.as_ref()) {
                                 error!("failed to send VmResponse: {:?}", e);
                             }

src/main.rs

@@ -382,6 +382,36 @@ fn stop_vms(args: std::env::Args) {
     }
 }
 
+fn balloon_vms(mut args: std::env::Args) {
+    let mut scm = Scm::new(1);
+    if args.len() < 2 {
+        print_help("crosvm balloon", "PAGE_ADJUST VM_SOCKET...", &[]);
+        println!("Adjust the ballon size of the crosvm instance by `PAGE_ADJUST` pages, `PAGE_ADJUST` can be negative to shrink the balloon.");
+    }
+    let num_pages: i32 = match args.nth(0).unwrap().parse::<i32>() {
+        Ok(n) => n,
+        Err(_) => {
+            error!("Failed to parse number of pages");
+            return;
+        },
+    };
+
+    for socket_path in args {
+        match UnixDatagram::unbound().and_then(|s| {
+                                                   s.connect(&socket_path)?;
+                                                   Ok(s)
+                                               }) {
+            Ok(s) => {
+                if let Err(e) = VmRequest::BalloonAdjust(num_pages).send(&mut scm, &s) {
+                    error!("failed to send balloon request to socket at '{}': {:?}",
+                           socket_path,
+                           e);
+                }
+            }
+            Err(e) => error!("failed to connect to socket at '{}': {}", socket_path, e),
+        }
+    }
+}
 
 fn print_usage() {
     print_help("crosvm", "[stop|run]", &[]);
@@ -410,6 +440,9 @@ fn main() {
         Some("run") => {
             run_vm(args);
         }
+        Some("balloon") => {
+            balloon_vms(args);
+        }
         Some(c) => {
             println!("invalid subcommand: {:?}", c);
             print_usage();

vm_control/Cargo.toml

@@ -4,6 +4,7 @@ version = "0.1.0"
 authors = ["The Chromium OS Authors"]
 
 [dependencies]
+byteorder = "*"
 data_model = { path = "../data_model" }
 kvm = { path = "../kvm" }
 libc = "*"

vm_control/src/lib.rs

@@ -10,6 +10,7 @@
 //! The wire message format is a little-endian C-struct of fixed size, along with a file descriptor
 //! if the request type expects one.
 
+extern crate byteorder;
 extern crate data_model;
 extern crate kvm;
 extern crate libc;
@@ -22,6 +23,7 @@ use std::result;
 
 use libc::{ERANGE, EINVAL};
 
+use byteorder::{LittleEndian, WriteBytesExt};
 use data_model::{DataInit, Le32, Le64, VolatileMemory};
 use sys_util::{EventFd, Error as SysError, MmapError, MemoryMapping, Scm, GuestAddress};
 use kvm::{IoeventAddress, Vm};
@@ -65,6 +67,8 @@ impl AsRawFd for MaybeOwnedFd {
 ///
 /// Unless otherwise noted, each request should expect a `VmResponse::Ok` to be received on success.
 pub enum VmRequest {
+    /// Try to grow or shrink the VM's balloon.
+    BalloonAdjust(i32),
     /// Break the VM's run loop and exit.
     Exit,
     /// Register the given ioevent address along with given datamatch to trigger the `EventFd`.
@@ -81,7 +85,8 @@ pub enum VmRequest {
 const VM_REQUEST_TYPE_EXIT: u32 = 1;
 const VM_REQUEST_TYPE_REGISTER_MEMORY: u32 = 2;
 const VM_REQUEST_TYPE_UNREGISTER_MEMORY: u32 = 3;
-const VM_REQUEST_SIZE: usize = 16;
+const VM_REQUEST_TYPE_BALLOON_ADJUST: u32 = 4;
+const VM_REQUEST_SIZE: usize = 24;
 
 #[repr(C)]
 #[derive(Clone, Copy, Default)]
@@ -89,6 +94,7 @@ struct VmRequestStruct {
     type_: Le32,
     slot: Le32,
     size: Le64,
+    num_pages: Le32,
 }
 
 // Safe because it only has data and has no implicit padding.
@@ -99,6 +105,7 @@ impl VmRequest {
     ///
     /// A `VmResponse` should be sent out over the given socket before another request is received.
     pub fn recv(scm: &mut Scm, s: &UnixDatagram) -> VmControlResult<VmRequest> {
+        assert_eq!(VM_REQUEST_SIZE, std::mem::size_of::<VmRequestStruct>());
         let mut buf = [0; VM_REQUEST_SIZE];
         let mut fds = Vec::new();
         let read = scm.recv(s, &mut [&mut buf], &mut fds)
@@ -118,6 +125,9 @@ impl VmRequest {
                                              req.size.to_native() as usize))
             }
             VM_REQUEST_TYPE_UNREGISTER_MEMORY => Ok(VmRequest::UnregisterMemory(req.slot.into())),
+            VM_REQUEST_TYPE_BALLOON_ADJUST => {
+                Ok(VmRequest::BalloonAdjust(req.num_pages.to_native() as i32))
+            },
             _ => Err(VmControlError::InvalidType),
         }
     }
@@ -127,6 +137,7 @@ impl VmRequest {
     /// After this request is a sent, a `VmResponse` should be received before sending another
     /// request.
     pub fn send(&self, scm: &mut Scm, s: &UnixDatagram) -> VmControlResult<()> {
+        assert_eq!(VM_REQUEST_SIZE, std::mem::size_of::<VmRequestStruct>());
         let mut req = VmRequestStruct::default();
         let mut fd_buf = [0; 1];
         let mut fd_len = 0;
@@ -142,6 +153,10 @@ impl VmRequest {
                 req.type_ = Le32::from(VM_REQUEST_TYPE_UNREGISTER_MEMORY);
                 req.slot = Le32::from(slot);
             }
+            &VmRequest::BalloonAdjust(pages) => {
+                req.type_ = Le32::from(VM_REQUEST_TYPE_BALLOON_ADJUST);
+                req.num_pages = Le32::from(pages as u32);
+            },
             _ => return Err(VmControlError::InvalidType),
         }
         let mut buf = [0; VM_REQUEST_SIZE];
@@ -162,7 +177,8 @@ impl VmRequest {
     /// This does not return a result, instead encapsulating the success or failure in a
     /// `VmResponse` with the intended purpose of sending the response back over the  socket that
     /// received this `VmRequest`.
-    pub fn execute(&self, vm: &mut Vm, next_mem_pfn: &mut u64, running: &mut bool) -> VmResponse {
+    pub fn execute(&self, vm: &mut Vm, next_mem_pfn: &mut u64, running: &mut bool,
+                   balloon_host_socket: &UnixDatagram) -> VmResponse {
         *running = true;
         match self {
             &VmRequest::Exit => {
@@ -209,6 +225,15 @@ impl VmRequest {
                     Err(e) => VmResponse::Err(e),
                 }
             }
+            &VmRequest::BalloonAdjust(num_pages) => {
+                let mut buf = [0u8; 4];
+                // write_i32 can't fail as the buffer is 4 bytes long.
+                (&mut buf[0..]).write_i32::<LittleEndian>(num_pages).unwrap();
+                match balloon_host_socket.send(&buf) {
+                    Ok(_) => VmResponse::Ok,
+                    Err(_) => VmResponse::Err(SysError::last()),
+                }
+            },
         }
     }
 }