From 77ce3f6cbfa08a2a695de1760575f35670418347 Mon Sep 17 00:00:00 2001
From: Jared White <git@jaredwhite.dev>
Date: Sun, 5 Apr 2026 21:36:07 -0700
Subject: [PATCH] vmm: memory_actual_size reflects hotplug state

It is desirable to be able to track the progress of memory hotplug.
Update the memory_actual_size field to query the current plugged size
from virtio-mem to enable this.

Signed-off-by: Jared White <git@jaredwhite.dev>
---
 virtio-devices/src/mem.rs |  4 ++++
 vmm/src/config.rs         | 31 +++++++++++++++++--------------
 vmm/src/lib.rs            |  4 +++-
 vmm/src/memory_manager.rs | 13 +++++++++++++
 vmm/src/vm.rs             |  8 ++++++++
 5 files changed, 45 insertions(+), 15 deletions(-)

diff --git a/virtio-devices/src/mem.rs b/virtio-devices/src/mem.rs
index aed8ed48d..067100164 100644
--- a/virtio-devices/src/mem.rs
+++ b/virtio-devices/src/mem.rs
@@ -834,6 +834,10 @@ impl Mem {
         })
     }
 
+    pub fn plugged_size(&self) -> u64 {
+        self.config.lock().unwrap().plugged_size
+    }
+
     pub fn resize(&mut self, size: u64) -> result::Result<(), Error> {
         let mut config = self.config.lock().unwrap();
         config.resize(size).map_err(|e| {
diff --git a/vmm/src/config.rs b/vmm/src/config.rs
index 93ec6c891..fc7eb1b8d 100644
--- a/vmm/src/config.rs
+++ b/vmm/src/config.rs
@@ -1039,21 +1039,24 @@ impl MemoryConfig {
     }
 
     pub fn total_size(&self) -> u64 {
-        let mut size = self.size;
-        if let Some(hotplugged_size) = self.hotplugged_size {
-            size += hotplugged_size;
-        }
+        self.size
+            + self
+                .zones
+                .iter()
+                .flatten()
+                .map(|zone| zone.size)
+                .sum::<u64>()
+            + self.hotplugged_size()
+    }
 
-        if let Some(zones) = &self.zones {
-            for zone in zones.iter() {
-                size += zone.size;
-                if let Some(hotplugged_size) = zone.hotplugged_size {
-                    size += hotplugged_size;
-                }
-            }
-        }
-
-        size
+    pub fn hotplugged_size(&self) -> u64 {
+        self.hotplugged_size.unwrap_or(0)
+            + self
+                .zones
+                .iter()
+                .flatten()
+                .filter_map(|zone| zone.hotplugged_size)
+                .sum::<u64>()
     }
 }
 
diff --git a/vmm/src/lib.rs b/vmm/src/lib.rs
index ce5664249..92bf4c6b7 100644
--- a/vmm/src/lib.rs
+++ b/vmm/src/lib.rs
@@ -1915,9 +1915,11 @@ impl RequestHandler for Vmm {
                 };
                 let config = vm_config.lock().unwrap().clone();
 
-                let mut memory_actual_size = config.memory.total_size();
+                let mut memory_actual_size =
+                    config.memory.total_size() - config.memory.hotplugged_size();
                 if let Some(vm) = &self.vm {
                     memory_actual_size = memory_actual_size.saturating_sub(vm.balloon_size());
+                    memory_actual_size += vm.virtio_mem_plugged_size();
                 }
 
                 let device_tree = self
diff --git a/vmm/src/memory_manager.rs b/vmm/src/memory_manager.rs
index 62b4522cc..b2058435b 100644
--- a/vmm/src/memory_manager.rs
+++ b/vmm/src/memory_manager.rs
@@ -2420,6 +2420,19 @@ impl MemoryManager {
         unsafe { (*stat.as_ptr()).st_nlink as usize > 0 }
     }
 
+    pub fn virtio_mem_plugged_size(&self) -> u64 {
+        self.memory_zones
+            .values()
+            .filter_map(|zone| {
+                zone.virtio_mem_zone
+                    .as_ref()?
+                    .virtio_device
+                    .as_ref()
+                    .map(|dev| dev.lock().unwrap().plugged_size())
+            })
+            .sum()
+    }
+
     pub fn memory_zones(&self) -> &MemoryZones {
         &self.memory_zones
     }
diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index c0eaba7ed..1f3bcfeeb 100644
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -2822,6 +2822,14 @@ impl Vm {
         self.device_manager.lock().unwrap().balloon_size()
     }
 
+    /// Get the actual size of the virtio_mem regions
+    pub fn virtio_mem_plugged_size(&self) -> u64 {
+        self.memory_manager
+            .lock()
+            .unwrap()
+            .virtio_mem_plugged_size()
+    }
+
     pub fn send_memory_fds(
         &mut self,
         socket: &mut UnixStream,