From e9637d37337e385378a2796a19097dfbf92ea15a Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Wed, 4 Aug 2021 16:38:01 +0200
Subject: [PATCH] vmm: device_manager: Fully implement Migratable trait

This patch connects the dots between the vm.rs code and each Migratable
device, in order to make sure Migratable methods are correctly invoked
when migration happens.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
---
 vmm/src/device_manager.rs | 35 ++++++++++++++++++++++++++++++++---
 1 file changed, 32 insertions(+), 3 deletions(-)

diff --git a/vmm/src/device_manager.rs b/vmm/src/device_manager.rs
index 089daced1..ec4868c80 100644
--- a/vmm/src/device_manager.rs
+++ b/vmm/src/device_manager.rs
@@ -100,8 +100,8 @@ use vm_memory::{Address, GuestAddress, GuestUsize, MmapRegion};
 #[cfg(all(target_arch = "x86_64", feature = "cmos"))]
 use vm_memory::{GuestAddressSpace, GuestMemory};
 use vm_migration::{
-    Migratable, MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable,
-    Transportable,
+    protocol::MemoryRangeTable, Migratable, MigratableError, Pausable, Snapshot,
+    SnapshotDataSection, Snapshottable, Transportable,
 };
 use vm_virtio::{VirtioDeviceType, VirtioIommuRemapping};
 use vmm_sys_util::eventfd::EventFd;
@@ -4008,7 +4008,36 @@ impl Snapshottable for DeviceManager {
 }
 
 impl Transportable for DeviceManager {}
-impl Migratable for DeviceManager {}
+
+impl Migratable for DeviceManager {
+    fn start_dirty_log(&mut self) -> std::result::Result<(), MigratableError> {
+        for (_, device_node) in self.device_tree.lock().unwrap().iter() {
+            if let Some(migratable) = &device_node.migratable {
+                migratable.lock().unwrap().start_dirty_log()?;
+            }
+        }
+        Ok(())
+    }
+
+    fn stop_dirty_log(&mut self) -> std::result::Result<(), MigratableError> {
+        for (_, device_node) in self.device_tree.lock().unwrap().iter() {
+            if let Some(migratable) = &device_node.migratable {
+                migratable.lock().unwrap().stop_dirty_log()?;
+            }
+        }
+        Ok(())
+    }
+
+    fn dirty_log(&mut self) -> std::result::Result<MemoryRangeTable, MigratableError> {
+        let mut tables = Vec::new();
+        for (_, device_node) in self.device_tree.lock().unwrap().iter() {
+            if let Some(migratable) = &device_node.migratable {
+                tables.push(migratable.lock().unwrap().dirty_log()?);
+            }
+        }
+        Ok(MemoryRangeTable::new_from_tables(tables))
+    }
+}
 
 const PCIU_FIELD_OFFSET: u64 = 0;
 const PCID_FIELD_OFFSET: u64 = 4;