refactor: VirtIO MMIO persistence logic

VirtIO MMIO restore logic activates the device the moment we restore the
device state, if the device was activated when snapshotted. Move the
activation responsibility to the logic the restores the MMIO transport.
The reason for this change is that that's how it will be done for the
PCI transport. Unifying this will allow us reusing the same types for
restoring the non-transport state of devices.

Note that we needed to change the way Net devices are saved/restored.
RxBuffer type of Net devices holds RX descriptors that we have parsed
from the Queue ahead of time. The way we restored this info was
manipulating the queue to re-parse the RX descriptors during the restore
phase. However, we need the device to be activated to do so, which now
isn't. So, instead of storing this info inside the snapshot make sure we
have flushed everything before taking the snapshot.

Also, simplify a bit the types that we use for serializing/deserializing
the state of a device.

Signed-off-by: Babis Chalios <[email protected]>

Author: bchalios

src/vmm/src/device_manager/persist.rs

@@ -87,16 +87,14 @@ pub struct BalloonState {
     stats_desc_index: Option<u16>,
     latest_stats: BalloonStatsState,
     config_space: BalloonConfigSpaceState,
-    virtio_state: VirtioDeviceState,
+    pub virtio_state: VirtioDeviceState,
 }
 
 /// Auxiliary structure for creating a device when resuming from a snapshot.
 #[derive(Debug)]
 pub struct BalloonConstructorArgs {
     /// Pointer to guest memory.
     pub mem: GuestMemoryMmap,
-    /// Interrupt used from the device.
-    pub interrupt: Arc<dyn VirtioInterrupt>,
     pub restored_from_file: bool,
 }
 
@@ -154,25 +152,18 @@ impl Persist<'_> for Balloon {
             actual_pages: state.config_space.actual_pages,
         };
 
-        if state.virtio_state.activated {
-            balloon.device_state = DeviceState::Activated(ActiveState {
-                mem: constructor_args.mem,
-                interrupt: constructor_args.interrupt,
-            });
-
-            if balloon.stats_enabled() {
-                // Restore the stats descriptor.
-                balloon.set_stats_desc_index(state.stats_desc_index);
-
-                // Restart timer if needed.
-                let timer_state = TimerState::Periodic {
-                    current: Duration::from_secs(u64::from(state.stats_polling_interval_s)),
-                    interval: Duration::from_secs(u64::from(state.stats_polling_interval_s)),
-                };
-                balloon
-                    .stats_timer
-                    .set_state(timer_state, SetTimeFlags::Default);
-            }
+        if state.virtio_state.activated && balloon.stats_enabled() {
+            // Restore the stats descriptor.
+            balloon.set_stats_desc_index(state.stats_desc_index);
+
+            // Restart timer if needed.
+            let timer_state = TimerState::Periodic {
+                current: Duration::from_secs(u64::from(state.stats_polling_interval_s)),
+                interval: Duration::from_secs(u64::from(state.stats_polling_interval_s)),
+            };
+            balloon
+                .stats_timer
+                .set_state(timer_state, SetTimeFlags::Default);
         }
 
         Ok(balloon)
@@ -202,7 +193,6 @@ mod tests {
         let restored_balloon = Balloon::restore(
             BalloonConstructorArgs {
                 mem: guest_mem,
-                interrupt: default_interrupt(),
                 restored_from_file: true,
             },
             &Snapshot::deserialize(&mut mem.as_slice()).unwrap(),

src/vmm/src/devices/virtio/balloon/persist.rs

@@ -17,9 +17,17 @@ pub enum BlockState {
     VhostUser(VhostUserBlockState),
 }
 
+impl BlockState {
+    pub fn is_activated(&self) -> bool {
+        match self {
+            BlockState::Virtio(virtio_block_state) => virtio_block_state.virtio_state.activated,
+            BlockState::VhostUser(vhost_user_block_state) => false,
+        }
+    }
+}
+
 /// Auxiliary structure for creating a device when resuming from a snapshot.
 #[derive(Debug)]
 pub struct BlockConstructorArgs {
     pub mem: GuestMemoryMmap,
-    pub interrupt: Arc<dyn VirtioInterrupt>,
 }

src/vmm/src/devices/virtio/block/persist.rs

@@ -58,7 +58,7 @@ pub struct VirtioBlockState {
     cache_type: CacheType,
     root_device: bool,
     disk_path: String,
-    virtio_state: VirtioDeviceState,
+    pub virtio_state: VirtioDeviceState,
     rate_limiter_state: RateLimiterState,
     file_engine_type: FileEngineTypeState,
 }
@@ -111,15 +111,6 @@ impl Persist<'_> for VirtioBlock {
         let avail_features = state.virtio_state.avail_features;
         let acked_features = state.virtio_state.acked_features;
 
-        let device_state = if state.virtio_state.activated {
-            DeviceState::Activated(ActiveState {
-                mem: constructor_args.mem,
-                interrupt: constructor_args.interrupt,
-            })
-        } else {
-            DeviceState::Inactive
-        };
-
         let config_space = ConfigSpace {
             capacity: disk_properties.nsectors.to_le(),
         };
@@ -132,7 +123,7 @@ impl Persist<'_> for VirtioBlock {
 
             queues,
             queue_evts,
-            device_state,
+            device_state: DeviceState::Inactive,
 
             id: state.id.clone(),
             partuuid: state.partuuid.clone(),
@@ -227,10 +218,7 @@ mod tests {
 
         // Restore the block device.
         let restored_block = VirtioBlock::restore(
-            BlockConstructorArgs {
-                mem: guest_mem,
-                interrupt: default_interrupt(),
-            },
+            BlockConstructorArgs { mem: guest_mem },
             &Snapshot::deserialize(&mut mem.as_slice()).unwrap(),
         )
         .unwrap();

src/vmm/src/devices/virtio/block/virtio/persist.rs

@@ -930,6 +930,19 @@ impl Net {
         let _ = self.resume_rx();
         let _ = self.process_tx();
     }
+
+    /// Prepare saving state
+    pub fn prepare_save(&mut self) {
+        if !self.is_activated() {
+            return;
+        }
+
+        // Give potential deferred RX frame to guest
+        self.rx_buffer.finish_frame(&mut self.queues[RX_INDEX]);
+        // Reset the parsed available descriptors, so we will re-parse them
+        self.queues[RX_INDEX].next_avail -=
+            u16::try_from(self.rx_buffer.parsed_descriptors.len()).unwrap();
+    }
 }
 
 impl VirtioDevice for Net {

src/vmm/src/devices/virtio/net/device.rs

@@ -30,27 +30,6 @@ pub struct NetConfigSpaceState {
     guest_mac: Option<MacAddr>,
 }
 
-/// Information about the parsed RX buffers
-#[derive(Debug, Default, Clone, Serialize, Deserialize)]
-pub struct RxBufferState {
-    // Number of iovecs we have parsed from the guest
-    parsed_descriptor_chains_nr: u16,
-    // Number of used descriptors
-    used_descriptors: u16,
-    // Number of used bytes
-    used_bytes: u32,
-}
-
-impl RxBufferState {
-    fn from_rx_buffers(rx_buffer: &RxBuffers) -> Self {
-        RxBufferState {
-            parsed_descriptor_chains_nr: rx_buffer.parsed_descriptors.len().try_into().unwrap(),
-            used_descriptors: rx_buffer.used_descriptors,
-            used_bytes: rx_buffer.used_bytes,
-        }
-    }
-}
-
 /// Information about the network device that are saved
 /// at snapshot.
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -62,17 +41,14 @@ pub struct NetState {
     /// The associated MMDS network stack.
     pub mmds_ns: Option<MmdsNetworkStackState>,
     config_space: NetConfigSpaceState,
-    virtio_state: VirtioDeviceState,
-    rx_buffers_state: RxBufferState,
+    pub virtio_state: VirtioDeviceState,
 }
 
 /// Auxiliary structure for creating a device when resuming from a snapshot.
 #[derive(Debug)]
 pub struct NetConstructorArgs {
     /// Pointer to guest memory.
     pub mem: GuestMemoryMmap,
-    /// Interrupt for the device.
-    pub interrupt: Arc<dyn VirtioInterrupt>,
     /// Pointer to the MMDS data store.
     pub mmds: Option<Arc<Mutex<Mmds>>>,
 }
@@ -108,7 +84,6 @@ impl Persist<'_> for Net {
                 guest_mac: self.guest_mac,
             },
             virtio_state: VirtioDeviceState::from_device(self),
-            rx_buffers_state: RxBufferState::from_rx_buffers(&self.rx_buffer),
         }
     }
 
@@ -153,25 +128,6 @@ impl Persist<'_> for Net {
         net.avail_features = state.virtio_state.avail_features;
         net.acked_features = state.virtio_state.acked_features;
 
-        if state.virtio_state.activated {
-            let supported_flags: u32 = Net::build_tap_offload_features(net.acked_features);
-            net.tap
-                .set_offload(supported_flags)
-                .map_err(NetPersistError::TapSetOffload)?;
-
-            net.device_state = DeviceState::Activated(ActiveState {
-                mem: constructor_args.mem,
-                interrupt: constructor_args.interrupt,
-            });
-
-            // Recreate `Net::rx_buffer`. We do it by re-parsing the RX queue. We're temporarily
-            // rolling back `next_avail` in the RX queue and call `parse_rx_descriptors`.
-            net.queues[RX_INDEX].next_avail -= state.rx_buffers_state.parsed_descriptor_chains_nr;
-            net.parse_rx_descriptors();
-            net.rx_buffer.used_descriptors = state.rx_buffers_state.used_descriptors;
-            net.rx_buffer.used_bytes = state.rx_buffers_state.used_bytes;
-        }
-
         Ok(net)
     }
 }
@@ -215,7 +171,6 @@ mod tests {
             match Net::restore(
                 NetConstructorArgs {
                     mem: guest_mem,
-                    interrupt: default_interrupt(),
                     mmds: mmds_ds,
                 },
                 &Snapshot::deserialize(&mut mem.as_slice()).unwrap(),

src/vmm/src/devices/virtio/net/persist.rs

@@ -19,20 +19,13 @@ use crate::vstate::memory::GuestMemoryMmap;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct EntropyState {
-    virtio_state: VirtioDeviceState,
+    pub virtio_state: VirtioDeviceState,
     rate_limiter_state: RateLimiterState,
 }
 
 #[derive(Debug)]
 pub struct EntropyConstructorArgs {
-    mem: GuestMemoryMmap,
-    interrupt: Arc<dyn VirtioInterrupt>,
-}
-
-impl EntropyConstructorArgs {
-    pub fn new(mem: GuestMemoryMmap, interrupt: Arc<dyn VirtioInterrupt>) -> Self {
-        Self { mem, interrupt }
-    }
+    pub mem: GuestMemoryMmap,
 }
 
 #[derive(Debug, thiserror::Error, displaydoc::Display)]
@@ -72,9 +65,6 @@ impl Persist<'_> for Entropy {
         let mut entropy = Entropy::new_with_queues(queues, rate_limiter)?;
         entropy.set_avail_features(state.virtio_state.avail_features);
         entropy.set_acked_features(state.virtio_state.acked_features);
-        if state.virtio_state.activated {
-            entropy.set_activated(constructor_args.mem, constructor_args.interrupt);
-        }
 
         Ok(entropy)
     }
@@ -99,7 +89,7 @@ mod tests {
 
         let guest_mem = create_virtio_mem();
         let restored = Entropy::restore(
-            EntropyConstructorArgs::new(guest_mem, default_interrupt()),
+            EntropyConstructorArgs { mem: guest_mem },
             &Snapshot::deserialize(&mut mem.as_slice()).unwrap(),
         )
         .unwrap();

src/vmm/src/devices/virtio/rng/persist.rs

@@ -31,7 +31,7 @@ pub struct VsockState {
 pub struct VsockFrontendState {
     /// Context Identifier.
     pub cid: u64,
-    virtio_state: VirtioDeviceState,
+    pub virtio_state: VirtioDeviceState,
 }
 
 /// An enum for the serializable backend state types.
@@ -53,8 +53,6 @@ pub struct VsockUdsState {
 pub struct VsockConstructorArgs<B> {
     /// Pointer to guest memory.
     pub mem: GuestMemoryMmap,
-    /// Interrupt to use for the device.
-    pub interrupt: Arc<dyn VirtioInterrupt>,
     /// The vsock Unix Backend.
     pub backend: B,
 }
@@ -123,14 +121,7 @@ where
 
         vsock.acked_features = state.virtio_state.acked_features;
         vsock.avail_features = state.virtio_state.avail_features;
-        vsock.device_state = if state.virtio_state.activated {
-            DeviceState::Activated(ActiveState {
-                mem: constructor_args.mem,
-                interrupt: constructor_args.interrupt,
-            })
-        } else {
-            DeviceState::Inactive
-        };
+        vsock.device_state = DeviceState::Inactive;
         Ok(vsock)
     }
 }
@@ -193,7 +184,6 @@ pub(crate) mod tests {
         let mut restored_device = Vsock::restore(
             VsockConstructorArgs {
                 mem: ctx.mem.clone(),
-                interrupt: default_interrupt(),
                 backend: match restored_state.backend {
                     VsockBackendState::Uds(uds_state) => {
                         assert_eq!(uds_state.path, "test".to_owned());