From ef92e55998a529219a24eb30b13b9cea6c5c1871 Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Fri, 25 Nov 2022 16:04:37 +0100
Subject: [PATCH] devices, vmm: Move Ioapic to new restore design

Moving the Ioapic object to the new restore design, meaning the Ioapic
is created directly with the right state, and it shares the same
codepath as when it's created from scratch.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
---
 devices/src/ioapic.rs     | 70 +++++++++++++++++++++------------------
 vmm/src/device_manager.rs |  2 ++
 2 files changed, 39 insertions(+), 33 deletions(-)

diff --git a/devices/src/ioapic.rs b/devices/src/ioapic.rs
index 8b072946c..a9d5ed442 100644
--- a/devices/src/ioapic.rs
+++ b/devices/src/ioapic.rs
@@ -10,7 +10,6 @@
 // See https://pdos.csail.mit.edu/6.828/2016/readings/ia32/ioapic.pdf for a specification.
 
 use super::interrupt_controller::{Error, InterruptController};
-use anyhow::anyhow;
 use byteorder::{ByteOrder, LittleEndian};
 use std::result;
 use std::sync::{Arc, Barrier};
@@ -194,6 +193,7 @@ impl Ioapic {
         id: String,
         apic_address: GuestAddress,
         interrupt_manager: Arc<dyn InterruptManager<GroupConfig = MsiIrqGroupConfig>>,
+        state: Option<IoapicState>,
     ) -> Result<Ioapic> {
         let interrupt_source_group = interrupt_manager
             .create_group(MsiIrqGroupConfig {
@@ -202,17 +202,47 @@ impl Ioapic {
             })
             .map_err(Error::CreateInterruptSourceGroup)?;
 
+        let (id_reg, reg_sel, reg_entries, used_entries, apic_address) = if let Some(state) = &state
+        {
+            (
+                state.id_reg,
+                state.reg_sel,
+                state.reg_entries,
+                state.used_entries,
+                GuestAddress(state.apic_address),
+            )
+        } else {
+            (
+                0,
+                0,
+                [0x10000; NUM_IOAPIC_PINS],
+                [false; NUM_IOAPIC_PINS],
+                apic_address,
+            )
+        };
+
         // The IOAPIC is created with entries already masked. The guest will be
         // in charge of unmasking them if/when necessary.
-        Ok(Ioapic {
+        let ioapic = Ioapic {
             id,
-            id_reg: 0,
-            reg_sel: 0,
-            reg_entries: [0x10000; NUM_IOAPIC_PINS],
-            used_entries: [false; NUM_IOAPIC_PINS],
+            id_reg,
+            reg_sel,
+            reg_entries,
+            used_entries,
             apic_address,
             interrupt_source_group,
-        })
+        };
+
+        // When restoring the Ioapic, we must enable used entries.
+        if state.is_some() {
+            for (irq, entry) in ioapic.used_entries.iter().enumerate() {
+                if *entry {
+                    ioapic.update_entry(irq)?;
+                }
+            }
+        }
+
+        Ok(ioapic)
     }
 
     fn ioapic_write(&mut self, val: u32) {
@@ -299,21 +329,6 @@ impl Ioapic {
         }
     }
 
-    fn set_state(&mut self, state: &IoapicState) -> Result<()> {
-        self.id_reg = state.id_reg;
-        self.reg_sel = state.reg_sel;
-        self.reg_entries = state.reg_entries;
-        self.used_entries = state.used_entries;
-        self.apic_address = GuestAddress(state.apic_address);
-        for (irq, entry) in self.used_entries.iter().enumerate() {
-            if *entry {
-                self.update_entry(irq)?;
-            }
-        }
-
-        Ok(())
-    }
-
     fn update_entry(&self, irq: usize) -> Result<()> {
         let entry = self.reg_entries[irq];
 
@@ -425,17 +440,6 @@ impl Snapshottable for Ioapic {
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
         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)?)
-            .map_err(|e| {
-                MigratableError::Restore(anyhow!(
-                    "Could not restore state for {}: {:?}",
-                    self.id,
-                    e
-                ))
-            })
-    }
 }
 
 impl Pausable for Ioapic {}
diff --git a/vmm/src/device_manager.rs b/vmm/src/device_manager.rs
index dde88b6de..d6e64c0d0 100644
--- a/vmm/src/device_manager.rs
+++ b/vmm/src/device_manager.rs
@@ -1408,6 +1408,8 @@ impl DeviceManager {
                 id.clone(),
                 APIC_START,
                 Arc::clone(&self.msi_interrupt_manager),
+                versioned_state_from_id(self.snapshot.as_ref(), id.as_str())
+                    .map_err(DeviceManagerError::RestoreGetState)?,
             )
             .map_err(DeviceManagerError::CreateInterruptController)?,
         ));