diff --git a/tests/integration.rs b/tests/integration.rs index 7313c06d7..e78bf282d 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -4994,6 +4994,7 @@ mod tests { .args(&["--api-socket", &api_socket]) .args(&["--cpus", "boot=4"]) .args(&["--memory", "size=4G"]) + .args(&["--balloon", "size=0"]) .args(&["--kernel", kernel_path.to_str().unwrap()]) .args(&[ "--disk", @@ -5022,6 +5023,12 @@ mod tests { assert_eq!(guest.get_cpu_count().unwrap_or_default(), 4); // Check the guest RAM assert!(guest.get_total_memory().unwrap_or_default() > 3_840_000); + // Use balloon to remove RAM from the VM + resize_command(&api_socket, None, None, Some(1 << 30)); + thread::sleep(std::time::Duration::new(5, 0)); + let total_memory = guest.get_total_memory().unwrap_or_default(); + assert!(total_memory > 2_880_000); + assert!(total_memory < 3_840_000); // Check the guest virtio-devices, e.g. block, rng, vsock, console, and net guest.check_devices_common(Some(&socket), Some(&console_text)); @@ -5098,6 +5105,12 @@ mod tests { // Perform same checks to validate VM has been properly restored assert_eq!(guest.get_cpu_count().unwrap_or_default(), 4); + let total_memory = guest.get_total_memory().unwrap_or_default(); + assert!(total_memory > 2_880_000); + assert!(total_memory < 3_840_000); + // Deflate balloon to restore entire RAM to the VM + resize_command(&api_socket, None, None, Some(0)); + thread::sleep(std::time::Duration::new(5, 0)); assert!(guest.get_total_memory().unwrap_or_default() > 3_840_000); guest.check_devices_common(Some(&socket), Some(&console_text)); }); diff --git a/virtio-devices/src/balloon.rs b/virtio-devices/src/balloon.rs index ce2cb6d34..5ebd7b632 100644 --- a/virtio-devices/src/balloon.rs +++ b/virtio-devices/src/balloon.rs @@ -29,12 +29,15 @@ use std::result; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::mpsc; use std::sync::{Arc, Barrier, Mutex}; +use versionize::{VersionMap, Versionize, VersionizeResult}; +use versionize_derive::Versionize; use vm_memory::GuestMemory; use vm_memory::{ Address, ByteValued, Bytes, GuestAddress, GuestAddressSpace, GuestMemoryAtomic, GuestMemoryError, }; -use vm_migration::{Migratable, MigratableError, Pausable, Snapshottable, Transportable}; +use vm_migration::VersionMapped; +use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable}; use vmm_sys_util::eventfd::EventFd; const QUEUE_SIZE: u16 = 128; @@ -82,8 +85,8 @@ pub enum Error { // Got from include/uapi/linux/virtio_balloon.h #[repr(C)] -#[derive(Copy, Clone, Debug, Default)] -struct VirtioBalloonConfig { +#[derive(Copy, Clone, Debug, Default, Versionize)] +pub struct VirtioBalloonConfig { // Number of pages host wants Guest to give up. num_pages: u32, // Number of pages we've actually got in balloon. @@ -120,11 +123,11 @@ struct VirtioBalloonResize { } impl VirtioBalloonResize { - pub fn new() -> io::Result { + pub fn new(size: u64) -> io::Result { let (tx, rx) = mpsc::channel(); Ok(Self { - size: Arc::new(AtomicU64::new(0)), + size: Arc::new(AtomicU64::new(size)), tx, rx, evt: EventFd::new(EFD_NONBLOCK)?, @@ -335,6 +338,15 @@ impl EpollHelperHandler for BalloonEpollHandler { } } +#[derive(Versionize)] +pub struct BalloonState { + pub avail_features: u64, + pub acked_features: u64, + pub config: VirtioBalloonConfig, +} + +impl VersionMapped for BalloonState {} + // Virtio device for exposing entropy to the guest OS through virtio. pub struct Balloon { common: VirtioCommon, @@ -374,7 +386,7 @@ impl Balloon { ..Default::default() }, id, - resize: VirtioBalloonResize::new()?, + resize: VirtioBalloonResize::new(size)?, config: Arc::new(Mutex::new(config)), seccomp_action, exit_evt, @@ -389,6 +401,20 @@ impl Balloon { pub fn get_actual(&self) -> u64 { (self.config.lock().unwrap().actual as u64) << VIRTIO_BALLOON_PFN_SHIFT } + + fn state(&self) -> BalloonState { + BalloonState { + avail_features: self.common.avail_features, + acked_features: self.common.acked_features, + config: *(self.config.lock().unwrap()), + } + } + + fn set_state(&mut self, state: &BalloonState) { + self.common.avail_features = state.avail_features; + self.common.acked_features = state.acked_features; + *(self.config.lock().unwrap()) = state.config; + } } impl Drop for Balloon { @@ -503,6 +529,15 @@ impl Snapshottable for Balloon { fn id(&self) -> String { self.id.clone() } + + fn snapshot(&mut self) -> std::result::Result { + Snapshot::new_from_versioned_state(&self.id(), &self.state()) + } + + fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> { + self.set_state(&snapshot.to_versioned_state(&self.id)?); + Ok(()) + } } impl Transportable for Balloon {} impl Migratable for Balloon {}