[RFC] virtio-balloon: Add free page reporting hinting

src/vmm/src/resources.rs

@@ -23,9 +23,7 @@ use crate::vmm_config::boot_source::{
 use crate::vmm_config::drive::*;
 use crate::vmm_config::entropy::*;
 use crate::vmm_config::instance_info::InstanceInfo;
-use crate::vmm_config::machine_config::{
-    HugePageConfig, MachineConfig, MachineConfigError, MachineConfigUpdate,
-};
+use crate::vmm_config::machine_config::{MachineConfig, MachineConfigError, MachineConfigUpdate};
 use crate::vmm_config::metrics::{MetricsConfig, MetricsConfigError, init_metrics};
 use crate::vmm_config::mmds::{MmdsConfig, MmdsConfigError};
 use crate::vmm_config::net::*;
@@ -245,10 +243,6 @@ impl VmResources {
             }
             SharedDeviceType::Balloon(balloon) => {
                 self.balloon.set_device(balloon);
-
-                if self.machine_config.huge_pages != HugePageConfig::None {
-                    return Err(ResourcesError::BalloonDevice(BalloonConfigError::HugePages));
-                }
             }
             SharedDeviceType::Vsock(vsock) => {
                 self.vsock.set_device(vsock);
@@ -290,9 +284,6 @@ impl VmResources {
             return Err(MachineConfigError::IncompatibleBalloonSize);
         }
 
-        if self.balloon.get().is_some() && updated.huge_pages != HugePageConfig::None {
-            return Err(MachineConfigError::BalloonAndHugePages);
-        }
         self.machine_config = updated;
 
         Ok(())
@@ -349,10 +340,6 @@ impl VmResources {
             return Err(BalloonConfigError::TooManyPagesRequested);
         }
 
-        if self.machine_config.huge_pages != HugePageConfig::None {
-            return Err(BalloonConfigError::HugePages);
-        }
-
         self.balloon.set(config)
     }
 
@@ -563,7 +550,6 @@ mod tests {
     use crate::HTTP_MAX_PAYLOAD_SIZE;
     use crate::cpu_config::templates::test_utils::TEST_TEMPLATE_JSON;
     use crate::cpu_config::templates::{CpuTemplateType, StaticCpuTemplate};
-    use crate::devices::virtio::balloon::Balloon;
     use crate::devices::virtio::block::virtio::VirtioBlockError;
     use crate::devices::virtio::block::{BlockError, CacheType};
     use crate::devices::virtio::vsock::VSOCK_DEV_ID;
@@ -1543,31 +1529,6 @@ mod tests {
             .unwrap_err();
     }
 
-    #[test]
-    fn test_negative_restore_balloon_device_with_huge_pages() {
-        let mut vm_resources = default_vm_resources();
-        vm_resources.balloon = BalloonBuilder::new();
-        vm_resources
-            .update_machine_config(&MachineConfigUpdate {
-                huge_pages: Some(HugePageConfig::Hugetlbfs2M),
-                ..Default::default()
-            })
-            .unwrap();
-        let err = vm_resources
-            .update_from_restored_device(SharedDeviceType::Balloon(Arc::new(Mutex::new(
-                Balloon::new(128, false, 0, false, false).unwrap(),
-            ))))
-            .unwrap_err();
-        assert!(
-            matches!(
-                err,
-                ResourcesError::BalloonDevice(BalloonConfigError::HugePages)
-            ),
-            "{:?}",
-            err
-        );
-    }
-
     #[test]
     fn test_set_entropy_device() {
         let mut vm_resources = default_vm_resources();

src/vmm/src/vmm_config/balloon.rs

@@ -20,8 +20,6 @@ pub enum BalloonConfigError {
     TooManyPagesRequested,
     /// Error creating the balloon device: {0}
     CreateFailure(crate::devices::virtio::balloon::BalloonError),
-    /// Firecracker's huge pages support is incompatible with memory ballooning.
-    HugePages,
 }
 
 /// This struct represents the strongly typed equivalent of the json body

src/vmm/src/vmm_config/machine_config.rs

@@ -29,8 +29,6 @@ pub enum MachineConfigError {
     SmtNotSupported,
     /// Could not determine host kernel version when checking hugetlbfs compatibility
     KernelVersion,
-    /// Firecracker's huge pages support is incompatible with memory ballooning.
-    BalloonAndHugePages,
 }
 
 /// Describes the possible (huge)page configurations for a microVM's memory.

tests/conftest.py

@@ -34,7 +34,7 @@
 from framework import defs, utils
 from framework.artifacts import disks, kernel_params
 from framework.defs import DEFAULT_BINARY_DIR
-from framework.microvm import MicroVMFactory, SnapshotType
+from framework.microvm import HugePagesConfig, MicroVMFactory, SnapshotType
 from framework.properties import global_props
 from framework.utils_cpu_templates import (
     custom_cpu_templates_params,
@@ -594,6 +594,15 @@ def pci_enabled(request):
     yield request.param
 
 
+@pytest.fixture(
+    params=[HugePagesConfig.NONE, HugePagesConfig.HUGETLBFS_2MB],
+    ids=["NO_HUGE_PAGES", "2M_HUGE_PAGES"],
+)
+def huge_pages(request):
+    """Fixture that allows configuring whether a microVM will have huge pages enabled or not"""
+    yield request.param
+
+
 def uvm_booted(
     microvm_factory,
     guest_kernel,

tests/framework/microvm.py

@@ -23,10 +23,11 @@
 from collections import namedtuple
 from dataclasses import dataclass
 from enum import Enum, auto
-from functools import lru_cache
+from functools import cached_property, lru_cache
 from pathlib import Path
 from typing import Optional
 
+import psutil
 from tenacity import Retrying, retry, stop_after_attempt, wait_fixed
 
 import host_tools.cargo_build as build_tools
@@ -472,7 +473,7 @@ def state(self):
         """Get the InstanceInfo property and return the state field."""
         return self.api.describe.get().json()["state"]
 
-    @property
+    @cached_property
     def firecracker_pid(self):
         """Return Firecracker's PID
 
@@ -491,6 +492,11 @@ def firecracker_pid(self):
             with attempt:
                 return int(self.jailer.pid_file.read_text(encoding="ascii"))
 
+    @cached_property
+    def ps(self):
+        """Returns a handle to the psutil.Process for this VM"""
+        return psutil.Process(self.firecracker_pid)
+
     @property
     def dimensions(self):
         """Gets a default set of cloudwatch dimensions describing the configuration of this microvm"""

tests/framework/utils.py

@@ -14,10 +14,12 @@
 import typing
 from collections import defaultdict, namedtuple
 from contextlib import contextmanager
+from pathlib import Path
 from typing import Dict
 
 import psutil
 import semver
+from packaging import version
 from tenacity import (
     Retrying,
     retry,
@@ -259,6 +261,48 @@ def get_free_mem_ssh(ssh_connection):
     raise Exception("Available memory not found in `/proc/meminfo")
 
 
+def get_stable_rss_mem_by_pid(process, percentage_delta=1):
+    """
+    Get the RSS memory that a guest uses, given the pid of the guest.
+
+    Wait till the fluctuations in RSS drop below percentage_delta.
+    Or print a warning if this does not happen.
+    """
+
+    # All values are reported as KiB
+
+    def get_rss_from_pmap():
+        """Returns current memory utilization in KiB, including used HugeTLBFS"""
+
+        proc_status = Path("/proc", str(process.pid), "status").read_text("utf-8")
+        for line in proc_status.splitlines():
+            if line.startswith("HugetlbPages:"):  # entry is in KiB
+                hugetlbfs_usage = int(line.split()[1])
+                break
+        else:
+            assert False, f"HugetlbPages not found in {str(proc_status)}"
+        return hugetlbfs_usage + process.memory_info().rss // 1024
+
+    first_rss = 0
+    second_rss = 0
+    for _ in range(5):
+        first_rss = get_rss_from_pmap()
+        time.sleep(1)
+        second_rss = get_rss_from_pmap()
+        abs_diff = abs(first_rss - second_rss)
+        abs_delta = abs_diff / first_rss * 100
+        print(
+            f"RSS readings: old: {first_rss} new: {second_rss} abs_diff: {abs_diff} abs_delta: {abs_delta}"
+        )
+        if abs_delta < percentage_delta:
+            return second_rss
+
+        time.sleep(1)
+
+    print("WARNING: RSS readings did not stabilize")
+    return second_rss
+
+
 def _format_output_message(proc, stdout, stderr):
     output_message = f"\n[{proc.pid}] Command:\n{proc.args}"
     # Append stdout/stderr to the output message
@@ -417,6 +461,11 @@ def get_kernel_version(level=2):
     return linux_version
 
 
+def supports_hugetlbfs_discard():
+    """Returns True if the kernel supports hugetlbfs discard"""
+    return version.parse(get_kernel_version()) >= version.parse("5.18.0")
+
+
 def generate_mmds_session_token(
     ssh_connection, ipv4_address, token_ttl, imds_compat=False
 ):