devices: support balloon free-page reporting

Added VIRTIO_BALLOON_F_PAGE_REPORTING bit.
Added reporting queue for handling messages from free-page reporting.
Balloon is expanding and now requires 8 futures. Added select8.
Added commandline flag balloon_page_reporting.
Modified name: inflate_tube -> release_memory_tube

Bug=b:235926042
Test=cd /devices/src/virtio && cargo test
Test=run_tests
Test=boot crosvm, set balloon, use up memory, free memory.

Change-Id: Iadb2f5d52cc4dd58e7b681b3983972d225e29861
Reviewed-on: https://chromium-review.googlesource.com/c/crosvm/crosvm/+/3774495
Tested-by: Elie Kheirallah <khei@google.com>
Commit-Queue: Elie Kheirallah <khei@google.com>
Reviewed-by: David Stevens <stevensd@chromium.org>

cros_async/src/lib.rs

@@ -382,6 +382,37 @@ pub async fn select7<
 ) {
     select::Select7::new(f1, f2, f3, f4, f5, f6, f7).await
 }
+
+pub async fn select8<
+    F1: Future + Unpin,
+    F2: Future + Unpin,
+    F3: Future + Unpin,
+    F4: Future + Unpin,
+    F5: Future + Unpin,
+    F6: Future + Unpin,
+    F7: Future + Unpin,
+    F8: Future + Unpin,
+>(
+    f1: F1,
+    f2: F2,
+    f3: F3,
+    f4: F4,
+    f5: F5,
+    f6: F6,
+    f7: F7,
+    f8: F8,
+) -> (
+    SelectResult<F1>,
+    SelectResult<F2>,
+    SelectResult<F3>,
+    SelectResult<F4>,
+    SelectResult<F5>,
+    SelectResult<F6>,
+    SelectResult<F7>,
+    SelectResult<F8>,
+) {
+    select::Select8::new(f1, f2, f3, f4, f5, f6, f7, f8).await
+}
 // Combination helpers to run until all futures are complete.
 
 /// Creates a combinator that runs the two given futures to completion, returning a tuple of the

cros_async/src/select.rs

@@ -89,4 +89,7 @@ generate! {
 
     /// _Future for the [`select7`] function.
     (Select7, <_Fut1, _Fut2, _Fut3, _Fut4, _Fut5, _Fut6, _Fut7>),
+
+    /// _Future for the [`select8`] function.
+    (Select8, <_Fut1, _Fut2, _Fut3, _Fut4, _Fut5, _Fut6, _Fut7, _Fut8>),
 }

devices/src/virtio/balloon.rs

@@ -20,7 +20,7 @@ use base::Event;
 use base::RawDescriptor;
 use base::Tube;
 use cros_async::block_on;
-use cros_async::select7;
+use cros_async::select8;
 use cros_async::sync::Mutex as AsyncMutex;
 use cros_async::AsyncTube;
 use cros_async::EventAsync;
@@ -83,6 +83,15 @@ const VIRTIO_BALLOON_PF_SIZE: u64 = 1 << VIRTIO_BALLOON_PFN_SHIFT;
 const VIRTIO_BALLOON_F_MUST_TELL_HOST: u32 = 0; // Tell before reclaiming pages
 const VIRTIO_BALLOON_F_STATS_VQ: u32 = 1; // Stats reporting enabled
 const VIRTIO_BALLOON_F_DEFLATE_ON_OOM: u32 = 2; // Deflate balloon on OOM
+const VIRTIO_BALLOON_F_PAGE_REPORTING: u32 = 5; // Page reporting virtqueue
+
+#[derive(Copy, Clone)]
+#[repr(u32)]
+// Balloon virtqueues
+pub enum BalloonFeatures {
+    // Page Reporting enabled
+    PageReporting = VIRTIO_BALLOON_F_PAGE_REPORTING,
+}
 
 // These feature bits are part of the proposal:
 //  https://lists.oasis-open.org/archives/virtio-comment/202201/msg00139.html
@@ -178,10 +187,10 @@ where
 }
 
 // Processes one message's list of addresses.
-// Unpin requests for each inflate range will be sent via `inflate_tube`
+// Unpin requests for each inflate range will be sent via `release_memory_tube`
 // if provided, and then `desc_handler` will be called for each inflate range.
 fn handle_address_chain<F>(
-    inflate_tube: &Option<Tube>,
+    release_memory_tube: &Option<Tube>,
     avail_desc: DescriptorChain,
     mem: &GuestMemory,
     desc_handler: &mut F,
@@ -225,7 +234,7 @@ where
         inflate_ranges.push((range_start, range_size));
     }
 
-    if let Some(tube) = inflate_tube {
+    if let Some(tube) = release_memory_tube {
         let unpin_ranges = inflate_ranges
             .iter()
             .map(|v| {
@@ -261,7 +270,7 @@ async fn handle_queue<F>(
     mem: &GuestMemory,
     mut queue: Queue,
     mut queue_event: EventAsync,
-    inflate_tube: &Option<Tube>,
+    release_memory_tube: &Option<Tube>,
     interrupt: Rc<RefCell<Interrupt>>,
     mut desc_handler: F,
 ) where
@@ -276,7 +285,9 @@ async fn handle_queue<F>(
             Ok(d) => d,
         };
         let index = avail_desc.index;
-        if let Err(e) = handle_address_chain(inflate_tube, avail_desc, mem, &mut desc_handler) {
+        if let Err(e) =
+            handle_address_chain(release_memory_tube, avail_desc, mem, &mut desc_handler)
+        {
             error!("balloon: failed to process inflate addresses: {}", e);
         }
         queue.add_used(mem, index, 0);
@@ -473,7 +484,7 @@ fn run_worker(
     queues: Vec<Queue>,
     command_tube: Tube,
     #[cfg(windows)] dynamic_mapping_tube: Tube,
-    inflate_tube: Option<Tube>,
+    release_memory_tube: Option<Tube>,
     interrupt: Interrupt,
     kill_evt: Event,
     mem: GuestMemory,
@@ -499,7 +510,7 @@ fn run_worker(
             &mem,
             queues.pop_front().unwrap(),
             queue_evts.pop_front().unwrap(),
-            &inflate_tube,
+            &release_memory_tube,
             interrupt.clone(),
             |guest_address, len| {
                 sys::free_memory(
@@ -552,6 +563,31 @@ fn run_worker(
         };
         pin_mut!(stats);
 
+        // The next queue is used for reporting messages
+        let reporting = if (acked_features & (1 << VIRTIO_BALLOON_F_PAGE_REPORTING)) != 0 {
+            handle_queue(
+                &mem,
+                queues.pop_front().unwrap(),
+                queue_evts.pop_front().unwrap(),
+                &release_memory_tube,
+                interrupt.clone(),
+                |guest_address, len| {
+                    sys::free_memory(
+                        &guest_address,
+                        len,
+                        #[cfg(windows)]
+                        &dynamic_mapping_tube,
+                        #[cfg(unix)]
+                        &mem,
+                    )
+                },
+            )
+            .left_future()
+        } else {
+            std::future::pending().right_future()
+        };
+        pin_mut!(reporting);
+
         // Future to handle command messages that resize the balloon.
         let command =
             handle_command_tube(&command_tube, interrupt.clone(), state.clone(), stats_tx);
@@ -582,8 +618,8 @@ fn run_worker(
         pin_mut!(events);
 
         if let Err(e) = ex
-            .run_until(select7(
+            .run_until(select8(
-                inflate, deflate, stats, command, resample, kill, events,
+                inflate, deflate, stats, reporting, command, resample, kill, events,
             ))
             .map(|_| ())
         {
@@ -591,7 +627,7 @@ fn run_worker(
         }
     }
 
-    inflate_tube
+    release_memory_tube
 }
 
 /// Virtio device for memory balloon inflation/deflation.
@@ -599,7 +635,7 @@ pub struct Balloon {
     command_tube: Option<Tube>,
     #[cfg(windows)]
     dynamic_mapping_tube: Option<Tube>,
-    inflate_tube: Option<Tube>,
+    release_memory_tube: Option<Tube>,
     state: Arc<AsyncMutex<BalloonState>>,
     features: u64,
     acked_features: u64,
@@ -619,21 +655,23 @@ pub enum BalloonMode {
 impl Balloon {
     /// Creates a new virtio balloon device.
     /// To let Balloon able to successfully release the memory which are pinned
-    /// by CoIOMMU to host, the inflate_tube will be used to send the inflate
+    /// by CoIOMMU to host, the release_memory_tube will be used to send the inflate
     /// ranges to CoIOMMU with UnpinRequest/UnpinResponse messages, so that The
     /// memory in the inflate range can be unpinned first.
     pub fn new(
         base_features: u64,
         command_tube: Tube,
         #[cfg(windows)] dynamic_mapping_tube: Tube,
-        inflate_tube: Option<Tube>,
+        release_memory_tube: Option<Tube>,
         init_balloon_size: u64,
         mode: BalloonMode,
+        enabled_features: u64,
     ) -> Result<Balloon> {
         let features = base_features
             | 1 << VIRTIO_BALLOON_F_MUST_TELL_HOST
             | 1 << VIRTIO_BALLOON_F_STATS_VQ
             | 1 << VIRTIO_BALLOON_F_EVENTS_VQ
+            | enabled_features
             | if mode == BalloonMode::Strict {
                 1 << VIRTIO_BALLOON_F_RESPONSIVE_DEVICE
             } else {
@@ -644,7 +682,7 @@ impl Balloon {
             command_tube: Some(command_tube),
             #[cfg(windows)]
             dynamic_mapping_tube: Some(dynamic_mapping_tube),
-            inflate_tube,
+            release_memory_tube,
             state: Arc::new(AsyncMutex::new(BalloonState {
                 num_pages: (init_balloon_size >> VIRTIO_BALLOON_PFN_SHIFT) as u32,
                 actual_pages: 0,
@@ -667,7 +705,9 @@ impl Balloon {
 
     fn num_expected_queues(acked_features: u64) -> usize {
         // mandatory inflate and deflate queues plus any optional ack'ed queues
-        let queue_bits = (1 << VIRTIO_BALLOON_F_STATS_VQ) | (1 << VIRTIO_BALLOON_F_EVENTS_VQ);
+        let queue_bits = (1 << VIRTIO_BALLOON_F_STATS_VQ)
+            | (1 << VIRTIO_BALLOON_F_EVENTS_VQ)
+            | (1 << VIRTIO_BALLOON_F_PAGE_REPORTING);
         2 + (acked_features & queue_bits as u64).count_ones() as usize
     }
 }
@@ -691,8 +731,8 @@ impl VirtioDevice for Balloon {
         if let Some(command_tube) = &self.command_tube {
             rds.push(command_tube.as_raw_descriptor());
         }
-        if let Some(inflate_tube) = &self.inflate_tube {
+        if let Some(release_memory_tube) = &self.release_memory_tube {
-            rds.push(inflate_tube.as_raw_descriptor());
+            rds.push(release_memory_tube.as_raw_descriptor());
         }
         rds
     }
@@ -765,7 +805,7 @@ impl VirtioDevice for Balloon {
         let command_tube = self.command_tube.take().unwrap();
         #[cfg(windows)]
         let mapping_tube = self.dynamic_mapping_tube.take().unwrap();
-        let inflate_tube = self.inflate_tube.take();
+        let release_memory_tube = self.release_memory_tube.take();
         let acked_features = self.acked_features;
         let worker_result = thread::Builder::new()
             .name("virtio_balloon".to_string())
@@ -776,7 +816,7 @@ impl VirtioDevice for Balloon {
                     command_tube,
                     #[cfg(windows)]
                     mapping_tube,
-                    inflate_tube,
+                    release_memory_tube,
                     interrupt,
                     kill_evt,
                     mem,
@@ -809,8 +849,8 @@ impl VirtioDevice for Balloon {
                     error!("{}: failed to get back resources", self.debug_label());
                     return false;
                 }
-                Ok(inflate_tube) => {
+                Ok(release_memory_tube) => {
-                    self.inflate_tube = inflate_tube;
+                    self.release_memory_tube = release_memory_tube;
                     return true;
                 }
             }
@@ -878,10 +918,11 @@ mod tests {
             Balloon::num_expected_queues(to_feature_bits(&[VIRTIO_BALLOON_F_STATS_VQ]))
         );
         assert_eq!(
-            4,
+            5,
             Balloon::num_expected_queues(to_feature_bits(&[
                 VIRTIO_BALLOON_F_STATS_VQ,
-                VIRTIO_BALLOON_F_EVENTS_VQ
+                VIRTIO_BALLOON_F_EVENTS_VQ,
+                VIRTIO_BALLOON_F_PAGE_REPORTING
             ]))
         );
     }

src/crosvm/cmdline.rs

@@ -468,6 +468,9 @@ pub struct RunCommand {
     #[argh(option, arg_name = "PATH")]
     /// path for balloon controller socket.
     pub balloon_control: Option<PathBuf>,
+    #[argh(switch)]
+    /// enable page reporting in balloon.
+    pub balloon_page_reporting: bool,
     #[argh(option)]
     /// comma separated key=value pairs for setting up battery
     /// device
@@ -1642,7 +1645,7 @@ impl TryFrom<RunCommand> for super::config::Config {
         cfg.usb = !cmd.no_usb;
         cfg.rng = !cmd.no_rng;
         cfg.balloon = !cmd.no_balloon;
-
+        cfg.balloon_page_reporting = cmd.balloon_page_reporting;
         #[cfg(feature = "audio")]
         {
             cfg.virtio_snds = cmd.virtio_snds;

src/crosvm/config.rs

@@ -1195,6 +1195,7 @@ pub struct Config {
     pub balloon: bool,
     pub balloon_bias: i64,
     pub balloon_control: Option<PathBuf>,
+    pub balloon_page_reporting: bool,
     pub battery_config: Option<BatteryConfig>,
     #[cfg(windows)]
     pub block_control_tube: Vec<Tube>,
@@ -1386,6 +1387,7 @@ impl Default for Config {
             balloon: true,
             balloon_bias: 0,
             balloon_control: None,
+            balloon_page_reporting: false,
             battery_config: None,
             #[cfg(windows)]
             block_control_tube: Vec::new(),
@@ -1742,6 +1744,10 @@ pub fn validate_config(cfg: &mut Config) -> std::result::Result<(), String> {
         return Err("'balloon-control' requires enabled balloon".to_string());
     }
 
+    if !cfg.balloon && cfg.balloon_page_reporting {
+        return Err("'balloon_page_reporting' requires enabled balloon".to_string());
+    }
+
     #[cfg(unix)]
     if cfg.lock_guest_memory && cfg.jail_config.is_none() {
         return Err("'lock-guest-memory' and 'disable-sandbox' are mutually exclusive".to_string());

src/crosvm/sys/unix.rs

@@ -70,6 +70,7 @@ use devices::virtio::vhost::user::VhostUserListener;
 use devices::virtio::vhost::user::VhostUserListenerTrait;
 use devices::virtio::vhost::vsock::VhostVsockConfig;
 #[cfg(feature = "balloon")]
+use devices::virtio::BalloonFeatures;
 use devices::virtio::BalloonMode;
 #[cfg(feature = "gpu")]
 use devices::virtio::EventDevice;
@@ -428,6 +429,8 @@ fn create_virtio_devices(
 
     #[cfg(feature = "balloon")]
     if let Some(balloon_device_tube) = balloon_device_tube {
+        let balloon_features =
+            (cfg.balloon_page_reporting as u64) << BalloonFeatures::PageReporting as u64;
         devs.push(create_balloon_device(
             cfg.protected_vm,
             &cfg.jail_config,
@@ -439,6 +442,7 @@ fn create_virtio_devices(
             balloon_device_tube,
             balloon_inflate_tube,
             init_balloon_size,
+            balloon_features,
         )?);
     }
 

src/crosvm/sys/unix/device_helpers.rs

@@ -719,6 +719,7 @@ pub fn create_balloon_device(
     tube: Tube,
     inflate_tube: Option<Tube>,
     init_balloon_size: u64,
+    enabled_features: u64,
 ) -> DeviceResult {
     let dev = virtio::Balloon::new(
         virtio::base_features(protected_vm),
@@ -726,6 +727,7 @@ pub fn create_balloon_device(
         inflate_tube,
         init_balloon_size,
         mode,
+        enabled_features,
     )
     .context("failed to create balloon")?;