vmm: virtio-devices: Restore every VirtioDevice upon creation

Following the new design proposal to improve the restore codepath when
migrating a VM, all virtio devices are supplied with an optional state
they can use to restore from. The restore() implementation every device
was providing has been removed in order to prevent from going through
the restoration twice.

Here is the list of devices now following the new restore design:

- Block (virtio-block)
- Net (virtio-net)
- Rng (virtio-rng)
- Fs (vhost-user-fs)
- Blk (vhost-user-block)
- Net (vhost-user-net)
- Pmem (virtio-pmem)
- Vsock (virtio-vsock)
- Mem (virtio-mem)
- Balloon (virtio-balloon)
- Watchdog (virtio-watchdog)
- Vdpa (vDPA)
- Console (virtio-console)
- Iommu (virtio-iommu)

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>

fuzz/fuzz_targets/balloon.rs

@@ -49,6 +49,7 @@ fuzz_target!(|bytes| {
         true,
         SeccompAction::Allow,
         EventFd::new(EFD_NONBLOCK).unwrap(),
+        None,
     )
     .unwrap();
 

fuzz/fuzz_targets/block.rs

@@ -60,6 +60,7 @@ fuzz_target!(|bytes| {
         SeccompAction::Allow,
         None,
         EventFd::new(EFD_NONBLOCK).unwrap(),
+        None,
     )
     .unwrap();
 

fuzz/fuzz_targets/iommu.rs

@@ -64,6 +64,7 @@ fuzz_target!(|bytes| {
         SeccompAction::Allow,
         EventFd::new(EFD_NONBLOCK).unwrap(),
         ((MEM_SIZE - IOVA_SPACE_SIZE) as u64, (MEM_SIZE - 1) as u64),
+        None,
     )
     .unwrap();
 

fuzz/fuzz_targets/mem.rs

@@ -152,6 +152,7 @@ fn create_dummy_virtio_mem(bytes: &[u8; VIRTIO_MEM_DATA_SIZE]) -> (Mem, Arc<Gues
             false,
             EventFd::new(EFD_NONBLOCK).unwrap(),
             blocks_state.clone(),
+            None,
         )
         .unwrap(),
         region,

fuzz/fuzz_targets/pmem.rs

@@ -127,6 +127,7 @@ fn create_dummy_pmem() -> Pmem {
         false,
         SeccompAction::Allow,
         EventFd::new(EFD_NONBLOCK).unwrap(),
+        None,
     )
     .unwrap()
 }

fuzz/fuzz_targets/rng.rs

@@ -63,6 +63,7 @@ fuzz_target!(|bytes| {
         false,
         SeccompAction::Allow,
         EventFd::new(EFD_NONBLOCK).unwrap(),
+        None,
     )
     .unwrap();
 

fuzz/fuzz_targets/watchdog.rs

@@ -38,6 +38,7 @@ fuzz_target!(|bytes| {
         EventFd::new(EFD_NONBLOCK).unwrap(),
         SeccompAction::Allow,
         EventFd::new(EFD_NONBLOCK).unwrap(),
+        None,
     )
     .unwrap();
 

virtio-devices/src/balloon.rs

@@ -360,15 +360,20 @@ impl Balloon {
         free_page_reporting: bool,
         seccomp_action: SeccompAction,
         exit_evt: EventFd,
+        state: Option<BalloonState>,
     ) -> io::Result<Self> {
         let mut queue_sizes = vec![QUEUE_SIZE; MIN_NUM_QUEUES];
+
+        let (avail_features, acked_features, config) = if let Some(state) = state {
+            info!("Restoring virtio-balloon {}", id);
+            (state.avail_features, state.acked_features, state.config)
+        } else {
             let mut avail_features = 1u64 << VIRTIO_F_VERSION_1;
             if deflate_on_oom {
                 avail_features |= 1u64 << VIRTIO_BALLOON_F_DEFLATE_ON_OOM;
             }
             if free_page_reporting {
                 avail_features |= 1u64 << VIRTIO_BALLOON_F_REPORTING;
-            queue_sizes.push(REPORTING_QUEUE_SIZE);
             }
 
             let config = VirtioBalloonConfig {
@@ -376,10 +381,18 @@ impl Balloon {
                 ..Default::default()
             };
 
+            (avail_features, 0, config)
+        };
+
+        if free_page_reporting {
+            queue_sizes.push(REPORTING_QUEUE_SIZE);
+        }
+
         Ok(Balloon {
             common: VirtioCommon {
                 device_type: VirtioDeviceType::Balloon as u32,
                 avail_features,
+                acked_features,
                 paused_sync: Some(Arc::new(Barrier::new(2))),
                 queue_sizes,
                 min_queues: MIN_NUM_QUEUES as u16,
@@ -418,12 +431,6 @@ impl Balloon {
         }
     }
 
-    fn set_state(&mut self, state: &BalloonState) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        self.config = state.config;
-    }
-
     #[cfg(fuzzing)]
     pub fn wait_for_epoll_threads(&mut self) {
         self.common.wait_for_epoll_threads();
@@ -573,11 +580,6 @@ impl Snapshottable for Balloon {
     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)?);
-        Ok(())
-    }
 }
 impl Transportable for Balloon {}
 impl Migratable for Balloon {}

virtio-devices/src/block.rs

@@ -416,7 +416,17 @@ impl Block {
         seccomp_action: SeccompAction,
         rate_limiter_config: Option<RateLimiterConfig>,
         exit_evt: EventFd,
+        state: Option<BlockState>,
     ) -> io::Result<Self> {
+        let (disk_nsectors, avail_features, acked_features, config) = if let Some(state) = state {
+            info!("Restoring virtio-block {}", id);
+            (
+                state.disk_nsectors,
+                state.avail_features,
+                state.acked_features,
+                state.config,
+            )
+        } else {
             let disk_size = disk_image.size().map_err(|e| {
                 io::Error::new(
                     io::ErrorKind::Other,
@@ -478,10 +488,14 @@ impl Block {
                 config.num_queues = num_queues as u16;
             }
 
+            (disk_nsectors, avail_features, 0, config)
+        };
+
         Ok(Block {
             common: VirtioCommon {
                 device_type: VirtioDeviceType::Block as u32,
                 avail_features,
+                acked_features,
                 paused_sync: Some(Arc::new(Barrier::new(num_queues + 1))),
                 queue_sizes: vec![queue_size; num_queues],
                 min_queues: 1,
@@ -510,14 +524,6 @@ impl Block {
         }
     }
 
-    fn set_state(&mut self, state: &BlockState) {
-        self.disk_path = state.disk_path.clone().into();
-        self.disk_nsectors = state.disk_nsectors;
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        self.config = state.config;
-    }
-
     fn update_writeback(&mut self) {
         // Use writeback from config if VIRTIO_BLK_F_CONFIG_WCE
         let writeback = if self.common.feature_acked(VIRTIO_BLK_F_CONFIG_WCE.into()) {
@@ -710,11 +716,6 @@ impl Snapshottable for Block {
     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)?);
-        Ok(())
-    }
 }
 impl Transportable for Block {}
 impl Migratable for Block {}

virtio-devices/src/console.rs

@@ -601,7 +601,7 @@ fn get_win_size(tty: &dyn AsRawFd) -> (u16, u16) {
 impl VersionMapped for ConsoleState {}
 
 impl Console {
-    /// Create a new virtio console device that gets random data from /dev/urandom.
+    /// Create a new virtio console device
     pub fn new(
         id: String,
         endpoint: Endpoint,
@@ -609,20 +609,37 @@ impl Console {
         iommu: bool,
         seccomp_action: SeccompAction,
         exit_evt: EventFd,
+        state: Option<ConsoleState>,
     ) -> io::Result<(Console, Arc<ConsoleResizer>)> {
+        let (avail_features, acked_features, config, in_buffer) = if let Some(state) = state {
+            info!("Restoring virtio-console {}", id);
+            (
+                state.avail_features,
+                state.acked_features,
+                state.config,
+                state.in_buffer.into(),
+            )
+        } else {
             let mut avail_features = 1u64 << VIRTIO_F_VERSION_1 | 1u64 << VIRTIO_CONSOLE_F_SIZE;
-
             if iommu {
                 avail_features |= 1u64 << VIRTIO_F_IOMMU_PLATFORM;
             }
 
+            (
+                avail_features,
+                0,
+                VirtioConsoleConfig::default(),
+                VecDeque::new(),
+            )
+        };
+
         let config_evt = EventFd::new(EFD_NONBLOCK).unwrap();
-        let console_config = Arc::new(Mutex::new(VirtioConsoleConfig::default()));
+        let console_config = Arc::new(Mutex::new(config));
         let resizer = Arc::new(ConsoleResizer {
             config_evt,
             config: console_config.clone(),
             tty: endpoint.out_file().as_ref().map(|t| t.try_clone().unwrap()),
-            acked_features: AtomicU64::new(0),
+            acked_features: AtomicU64::new(acked_features),
         });
 
         resizer.update_console_size();
@@ -633,6 +650,7 @@ impl Console {
                     device_type: VirtioDeviceType::Console as u32,
                     queue_sizes: QUEUE_SIZES.to_vec(),
                     avail_features,
+                    acked_features,
                     paused_sync: Some(Arc::new(Barrier::new(2))),
                     min_queues: NUM_QUEUES as u16,
                     ..Default::default()
@@ -643,7 +661,7 @@ impl Console {
                 resize_pipe,
                 endpoint,
                 seccomp_action,
-                in_buffer: Arc::new(Mutex::new(VecDeque::new())),
+                in_buffer: Arc::new(Mutex::new(in_buffer)),
                 exit_evt,
             },
             resizer,
@@ -658,13 +676,6 @@ impl Console {
             in_buffer: self.in_buffer.lock().unwrap().clone().into(),
         }
     }
-
-    fn set_state(&mut self, state: &ConsoleState) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        *(self.config.lock().unwrap()) = state.config;
-        *(self.in_buffer.lock().unwrap()) = state.in_buffer.clone().into();
-    }
 }
 
 impl Drop for Console {
@@ -786,11 +797,6 @@ impl Snapshottable for Console {
     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)?);
-        Ok(())
-    }
 }
 impl Transportable for Console {}
 impl Migratable for Console {}

virtio-devices/src/iommu.rs

@@ -850,7 +850,7 @@ type EndpointsState = Vec<(u32, u32)>;
 type DomainsState = Vec<(u32, (Vec<(u64, Mapping)>, bool))>;
 
 #[derive(Versionize)]
-struct IommuState {
+pub struct IommuState {
     avail_features: u64,
     acked_features: u64,
     endpoints: EndpointsState,
@@ -865,7 +865,37 @@ impl Iommu {
         seccomp_action: SeccompAction,
         exit_evt: EventFd,
         msi_iova_space: (u64, u64),
+        state: Option<IommuState>,
     ) -> io::Result<(Self, Arc<IommuMapping>)> {
+        let (avail_features, acked_features, endpoints, domains) = if let Some(state) = state {
+            info!("Restoring virtio-iommu {}", id);
+            (
+                state.avail_features,
+                state.acked_features,
+                state.endpoints.into_iter().collect(),
+                state
+                    .domains
+                    .into_iter()
+                    .map(|(k, v)| {
+                        (
+                            k,
+                            Domain {
+                                mappings: v.0.into_iter().collect(),
+                                bypass: v.1,
+                            },
+                        )
+                    })
+                    .collect(),
+            )
+        } else {
+            let avail_features = 1u64 << VIRTIO_F_VERSION_1
+                | 1u64 << VIRTIO_IOMMU_F_MAP_UNMAP
+                | 1u64 << VIRTIO_IOMMU_F_PROBE
+                | 1u64 << VIRTIO_IOMMU_F_BYPASS_CONFIG;
+
+            (avail_features, 0, BTreeMap::new(), BTreeMap::new())
+        };
+
         let config = VirtioIommuConfig {
             page_size_mask: VIRTIO_IOMMU_PAGE_SIZE_MASK,
             probe_size: PROBE_PROP_SIZE,
@@ -873,8 +903,8 @@ impl Iommu {
         };
 
         let mapping = Arc::new(IommuMapping {
-            endpoints: Arc::new(RwLock::new(BTreeMap::new())),
+            endpoints: Arc::new(RwLock::new(endpoints)),
-            domains: Arc::new(RwLock::new(BTreeMap::new())),
+            domains: Arc::new(RwLock::new(domains)),
             bypass: AtomicBool::new(true),
         });
 
@@ -884,10 +914,8 @@ impl Iommu {
                 common: VirtioCommon {
                     device_type: VirtioDeviceType::Iommu as u32,
                     queue_sizes: QUEUE_SIZES.to_vec(),
-                    avail_features: 1u64 << VIRTIO_F_VERSION_1
+                    avail_features,
-                        | 1u64 << VIRTIO_IOMMU_F_MAP_UNMAP
+                    acked_features,
-                        | 1u64 << VIRTIO_IOMMU_F_PROBE
-                        | 1u64 << VIRTIO_IOMMU_F_BYPASS_CONFIG,
                     paused_sync: Some(Arc::new(Barrier::new(2))),
                     min_queues: NUM_QUEUES as u16,
                     ..Default::default()
@@ -927,26 +955,6 @@ impl Iommu {
         }
     }
 
-    fn set_state(&mut self, state: &IommuState) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        *(self.mapping.endpoints.write().unwrap()) = state.endpoints.clone().into_iter().collect();
-        *(self.mapping.domains.write().unwrap()) = state
-            .domains
-            .clone()
-            .into_iter()
-            .map(|(k, v)| {
-                (
-                    k,
-                    Domain {
-                        mappings: v.0.into_iter().collect(),
-                        bypass: v.1,
-                    },
-                )
-            })
-            .collect();
-    }
-
     fn update_bypass(&mut self) {
         // Use bypass from config if VIRTIO_IOMMU_F_BYPASS_CONFIG has been negotiated
         if !self
@@ -1088,11 +1096,6 @@ impl Snapshottable for Iommu {
     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)?);
-        Ok(())
-    }
 }
 impl Transportable for Iommu {}
 impl Migratable for Iommu {}

virtio-devices/src/mem.rs

@@ -714,6 +714,7 @@ impl Mem {
         hugepages: bool,
         exit_evt: EventFd,
         blocks_state: Arc<Mutex<BlocksState>>,
+        state: Option<MemState>,
     ) -> io::Result<Mem> {
         let region_len = region.len();
 
@@ -727,6 +728,11 @@ impl Mem {
             ));
         }
 
+        let (avail_features, acked_features, config) = if let Some(state) = state {
+            info!("Restoring virtio-mem {}", id);
+            *(blocks_state.lock().unwrap()) = state.blocks_state.clone();
+            (state.avail_features, state.acked_features, state.config)
+        } else {
             let mut avail_features = 1u64 << VIRTIO_F_VERSION_1;
 
             let mut config = VirtioMemConfig {
@@ -765,6 +771,9 @@ impl Mem {
                 )
             })?;
 
+            (avail_features, 0, config)
+        };
+
         let host_fd = region
             .file_offset()
             .map(|f_offset| f_offset.file().as_raw_fd());
@@ -773,6 +782,7 @@ impl Mem {
             common: VirtioCommon {
                 device_type: VirtioDeviceType::Mem as u32,
                 avail_features,
+                acked_features,
                 paused_sync: Some(Arc::new(Barrier::new(2))),
                 queue_sizes: QUEUE_SIZES.to_vec(),
                 min_queues: 1,
@@ -870,13 +880,6 @@ impl Mem {
         }
     }
 
-    fn set_state(&mut self, state: &MemState) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        *(self.config.lock().unwrap()) = state.config;
-        *(self.blocks_state.lock().unwrap()) = state.blocks_state.clone();
-    }
-
     #[cfg(fuzzing)]
     pub fn wait_for_epoll_threads(&mut self) {
         self.common.wait_for_epoll_threads();
@@ -999,11 +1002,6 @@ impl Snapshottable for Mem {
     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)?);
-        Ok(())
-    }
 }
 impl Transportable for Mem {}
 impl Migratable for Mem {}

virtio-devices/src/net.rs

@@ -445,11 +445,21 @@ impl Net {
         seccomp_action: SeccompAction,
         rate_limiter_config: Option<RateLimiterConfig>,
         exit_evt: EventFd,
+        state: Option<NetState>,
     ) -> Result<Self> {
         assert!(!taps.is_empty());
 
         let mtu = taps[0].mtu().map_err(Error::TapError)? as u16;
 
+        let (avail_features, acked_features, config, queue_sizes) = if let Some(state) = state {
+            info!("Restoring virtio-net {}", id);
+            (
+                state.avail_features,
+                state.acked_features,
+                state.config,
+                state.queue_size,
+            )
+        } else {
             let mut avail_features = 1 << VIRTIO_NET_F_CSUM
                 | 1 << VIRTIO_NET_F_CTRL_GUEST_OFFLOADS
                 | 1 << VIRTIO_NET_F_GUEST_CSUM
@@ -474,16 +484,31 @@ impl Net {
 
             let mut config = VirtioNetConfig::default();
             if let Some(mac) = guest_mac {
-            build_net_config_space(&mut config, mac, num_queues, Some(mtu), &mut avail_features);
+                build_net_config_space(
+                    &mut config,
+                    mac,
+                    num_queues,
+                    Some(mtu),
+                    &mut avail_features,
+                );
             } else {
-            build_net_config_space_with_mq(&mut config, num_queues, Some(mtu), &mut avail_features);
+                build_net_config_space_with_mq(
+                    &mut config,
+                    num_queues,
+                    Some(mtu),
+                    &mut avail_features,
+                );
             }
 
+            (avail_features, 0, config, vec![queue_size; queue_num])
+        };
+
         Ok(Net {
             common: VirtioCommon {
                 device_type: VirtioDeviceType::Net as u32,
                 avail_features,
-                queue_sizes: vec![queue_size; queue_num],
+                acked_features,
+                queue_sizes,
                 paused_sync: Some(Arc::new(Barrier::new((num_queues / 2) + 1))),
                 min_queues: 2,
                 ..Default::default()
@@ -516,6 +541,7 @@ impl Net {
         seccomp_action: SeccompAction,
         rate_limiter_config: Option<RateLimiterConfig>,
         exit_evt: EventFd,
+        state: Option<NetState>,
     ) -> Result<Self> {
         let taps = open_tap(
             if_name,
@@ -538,6 +564,7 @@ impl Net {
             seccomp_action,
             rate_limiter_config,
             exit_evt,
+            state,
         )
     }
 
@@ -552,6 +579,7 @@ impl Net {
         seccomp_action: SeccompAction,
         rate_limiter_config: Option<RateLimiterConfig>,
         exit_evt: EventFd,
+        state: Option<NetState>,
     ) -> Result<Self> {
         let mut taps: Vec<Tap> = Vec::new();
         let num_queue_pairs = fds.len();
@@ -583,6 +611,7 @@ impl Net {
             seccomp_action,
             rate_limiter_config,
             exit_evt,
+            state,
         )
     }
 
@@ -594,13 +623,6 @@ impl Net {
             queue_size: self.common.queue_sizes.clone(),
         }
     }
-
-    fn set_state(&mut self, state: &NetState) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        self.config = state.config;
-        self.common.queue_sizes = state.queue_size.clone();
-    }
 }
 
 impl Drop for Net {
@@ -820,11 +842,6 @@ impl Snapshottable for Net {
     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)?);
-        Ok(())
-    }
 }
 impl Transportable for Net {}
 impl Migratable for Net {}

virtio-devices/src/pmem.rs

@@ -299,7 +299,12 @@ impl Pmem {
         iommu: bool,
         seccomp_action: SeccompAction,
         exit_evt: EventFd,
+        state: Option<PmemState>,
     ) -> io::Result<Pmem> {
+        let (avail_features, acked_features, config) = if let Some(state) = state {
+            info!("Restoring virtio-pmem {}", id);
+            (state.avail_features, state.acked_features, state.config)
+        } else {
             let config = VirtioPmemConfig {
                 start: addr.raw_value().to_le(),
                 size: (_region.size() as u64).to_le(),
@@ -310,6 +315,8 @@ impl Pmem {
             if iommu {
                 avail_features |= 1u64 << VIRTIO_F_IOMMU_PLATFORM;
             }
+            (avail_features, 0, config)
+        };
 
         Ok(Pmem {
             common: VirtioCommon {
@@ -317,6 +324,7 @@ impl Pmem {
                 queue_sizes: QUEUE_SIZES.to_vec(),
                 paused_sync: Some(Arc::new(Barrier::new(2))),
                 avail_features,
+                acked_features,
                 min_queues: 1,
                 ..Default::default()
             },
@@ -338,12 +346,6 @@ impl Pmem {
         }
     }
 
-    fn set_state(&mut self, state: &PmemState) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        self.config = state.config;
-    }
-
     #[cfg(fuzzing)]
     pub fn wait_for_epoll_threads(&mut self) {
         self.common.wait_for_epoll_threads();
@@ -461,11 +463,6 @@ impl Snapshottable for Pmem {
     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)?);
-        Ok(())
-    }
 }
 
 impl Transportable for Pmem {}

virtio-devices/src/rng.rs

@@ -174,20 +174,30 @@ impl Rng {
         iommu: bool,
         seccomp_action: SeccompAction,
         exit_evt: EventFd,
+        state: Option<RngState>,
     ) -> io::Result<Rng> {
         let random_file = File::open(path)?;
+
+        let (avail_features, acked_features) = if let Some(state) = state {
+            info!("Restoring virtio-rng {}", id);
+            (state.avail_features, state.acked_features)
+        } else {
             let mut avail_features = 1u64 << VIRTIO_F_VERSION_1;
 
             if iommu {
                 avail_features |= 1u64 << VIRTIO_F_IOMMU_PLATFORM;
             }
 
+            (avail_features, 0)
+        };
+
         Ok(Rng {
             common: VirtioCommon {
                 device_type: VirtioDeviceType::Rng as u32,
                 queue_sizes: QUEUE_SIZES.to_vec(),
                 paused_sync: Some(Arc::new(Barrier::new(2))),
                 avail_features,
+                acked_features,
                 min_queues: 1,
                 ..Default::default()
             },
@@ -205,11 +215,6 @@ impl Rng {
         }
     }
 
-    fn set_state(&mut self, state: &RngState) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-    }
-
     #[cfg(fuzzing)]
     pub fn wait_for_epoll_threads(&mut self) {
         self.common.wait_for_epoll_threads();
@@ -319,11 +324,6 @@ impl Snapshottable for Rng {
     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)?);
-        Ok(())
-    }
 }
 
 impl Transportable for Rng {}

virtio-devices/src/vdpa.rs

@@ -64,6 +64,8 @@ pub enum Error {
     ResetOwner(vhost::Error),
     #[error("Failed to set backend specific features: {0}")]
     SetBackendFeatures(vhost::Error),
+    #[error("Failed to set backend configuration: {0}")]
+    SetConfig(vhost::Error),
     #[error("Failed to set eventfd notifying about a configuration change: {0}")]
     SetConfigCall(vhost::Error),
     #[error("Failed to set virtio features: {0}")]
@@ -107,14 +109,11 @@ impl VersionMapped for VdpaState {}
 pub struct Vdpa {
     common: VirtioCommon,
     id: String,
-    mem: GuestMemoryAtomic<GuestMemoryMmap>,
-    device_path: String,
     vhost: Option<VhostKernVdpa<GuestMemoryAtomic<GuestMemoryMmap>>>,
     iova_range: VhostVdpaIovaRange,
     enabled_queues: BTreeMap<usize, bool>,
     backend_features: u64,
     migrating: bool,
-    buffered_maps: Vec<(u64, u64, u64, bool)>,
 }
 
 impl Vdpa {
@@ -123,30 +122,38 @@ impl Vdpa {
         device_path: &str,
         mem: GuestMemoryAtomic<GuestMemoryMmap>,
         num_queues: u16,
-        restoring: bool,
+        state: Option<VdpaState>,
     ) -> Result<Self> {
-        if restoring {
+        let mut vhost = VhostKernVdpa::new(device_path, mem).map_err(Error::CreateVhostVdpa)?;
-            return Ok(Vdpa {
-                common: VirtioCommon {
-                    queue_sizes: vec![1024; num_queues as usize],
-                    min_queues: num_queues,
-                    ..Default::default()
-                },
-                id,
-                mem,
-                device_path: device_path.to_string(),
-                vhost: None,
-                iova_range: VhostVdpaIovaRange { first: 0, last: 0 },
-                enabled_queues: BTreeMap::new(),
-                backend_features: 0,
-                migrating: false,
-                buffered_maps: Vec::new(),
-            });
-        }
-
-        let mut vhost =
-            VhostKernVdpa::new(device_path, mem.clone()).map_err(Error::CreateVhostVdpa)?;
         vhost.set_owner().map_err(Error::SetOwner)?;
+
+        let (
+            device_type,
+            avail_features,
+            acked_features,
+            queue_sizes,
+            iova_range,
+            backend_features,
+        ) = if let Some(state) = state {
+            info!("Restoring vDPA {}", id);
+
+            vhost.set_backend_features_acked(state.backend_features);
+            vhost
+                .set_config(0, state.config.as_slice())
+                .map_err(Error::SetConfig)?;
+
+            (
+                state.device_type,
+                state.avail_features,
+                state.acked_features,
+                state.queue_sizes,
+                VhostVdpaIovaRange {
+                    first: state.iova_range_first,
+                    last: state.iova_range_last,
+                },
+                state.backend_features,
+            )
+        } else {
             let device_type = vhost.get_device_id().map_err(Error::GetDeviceId)?;
             let queue_size = vhost.get_vring_num().map_err(Error::GetVringNum)?;
             let avail_features = vhost.get_features().map_err(Error::GetFeatures)?;
@@ -161,23 +168,31 @@ impl Vdpa {
                 return Err(Error::MissingAccessPlatformVirtioFeature);
             }
 
+            (
+                device_type,
+                avail_features,
+                0,
+                vec![queue_size; num_queues as usize],
+                iova_range,
+                backend_features,
+            )
+        };
+
         Ok(Vdpa {
             common: VirtioCommon {
                 device_type,
-                queue_sizes: vec![queue_size; num_queues as usize],
+                queue_sizes,
                 avail_features,
+                acked_features,
                 min_queues: num_queues,
                 ..Default::default()
             },
             id,
-            mem,
-            device_path: device_path.to_string(),
             vhost: Some(vhost),
             iova_range,
             enabled_queues: BTreeMap::new(),
             backend_features,
             migrating: false,
-            buffered_maps: Vec::new(),
         })
     }
 
@@ -318,12 +333,6 @@ impl Vdpa {
         host_vaddr: *const u8,
         readonly: bool,
     ) -> Result<()> {
-        if self.vhost.is_none() {
-            self.buffered_maps
-                .push((iova, size, host_vaddr as u64, readonly));
-            return Ok(());
-        }
-
         let iova_last = iova + size - 1;
         if iova < self.iova_range.first || iova_last > self.iova_range.last {
             return Err(Error::InvalidIovaRange(iova, iova_last));
@@ -373,33 +382,6 @@ impl Vdpa {
             backend_features: self.backend_features,
         })
     }
-
-    fn set_state(&mut self, state: &VdpaState) -> Result<()> {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        self.common.device_type = state.device_type;
-        self.common.queue_sizes = state.queue_sizes.clone();
-        self.iova_range = VhostVdpaIovaRange {
-            first: state.iova_range_first,
-            last: state.iova_range_last,
-        };
-        self.backend_features = state.backend_features;
-
-        let mut vhost = VhostKernVdpa::new(self.device_path.as_str(), self.mem.clone())
-            .map_err(Error::CreateVhostVdpa)?;
-        vhost.set_owner().map_err(Error::SetOwner)?;
-        vhost.set_backend_features_acked(self.backend_features);
-        self.vhost = Some(vhost);
-
-        self.write_config(0, state.config.as_slice());
-
-        let maps: Vec<(u64, u64, u64, bool)> = self.buffered_maps.drain(..).collect();
-        for (iova, size, host_vaddr, readonly) in maps {
-            self.dma_map(iova, size, host_vaddr as *const u8, readonly)?;
-        }
-
-        Ok(())
-    }
 }
 
 impl VirtioDevice for Vdpa {
@@ -514,14 +496,6 @@ impl Snapshottable for Vdpa {
 
         Ok(snapshot)
     }
-
-    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!("Error restoring vDPA device: {:?}", e))
-            })?;
-        Ok(())
-    }
 }
 
 impl Transportable for Vdpa {}

virtio-devices/src/vhost_user/blk.rs

@@ -69,43 +69,33 @@ impl Blk {
     pub fn new(
         id: String,
         vu_cfg: VhostUserConfig,
-        restoring: bool,
         seccomp_action: SeccompAction,
         exit_evt: EventFd,
         iommu: bool,
+        state: Option<State>,
     ) -> Result<Blk> {
         let num_queues = vu_cfg.num_queues;
 
-        if restoring {
-            // We need 'queue_sizes' to report a number of queues that will be
-            // enough to handle all the potential queues. VirtioPciDevice::new()
-            // will create the actual queues based on this information.
-            return Ok(Blk {
-                common: VirtioCommon {
-                    device_type: VirtioDeviceType::Block as u32,
-                    queue_sizes: vec![vu_cfg.queue_size; num_queues],
-                    paused_sync: Some(Arc::new(Barrier::new(2))),
-                    min_queues: DEFAULT_QUEUE_NUMBER as u16,
-                    ..Default::default()
-                },
-                vu_common: VhostUserCommon {
-                    socket_path: vu_cfg.socket,
-                    vu_num_queues: num_queues,
-                    ..Default::default()
-                },
-                id,
-                config: VirtioBlockConfig::default(),
-                guest_memory: None,
-                epoll_thread: None,
-                seccomp_action,
-                exit_evt,
-                iommu,
-            });
-        }
-
         let mut vu =
             VhostUserHandle::connect_vhost_user(false, &vu_cfg.socket, num_queues as u64, false)?;
 
+        let (avail_features, acked_features, acked_protocol_features, vu_num_queues, config) =
+            if let Some(state) = state {
+                info!("Restoring vhost-user-block {}", id);
+
+                vu.set_protocol_features_vhost_user(
+                    state.acked_features,
+                    state.acked_protocol_features,
+                )?;
+
+                (
+                    state.avail_features,
+                    state.acked_features,
+                    state.acked_protocol_features,
+                    state.vu_num_queues,
+                    state.config,
+                )
+            } else {
                 // Filling device and vring features VMM supports.
                 let mut avail_features = 1 << VIRTIO_BLK_F_SIZE_MAX
                     | 1 << VIRTIO_BLK_F_SEG_MAX
@@ -137,7 +127,8 @@ impl Blk {
                     if acked_protocol_features & VhostUserProtocolFeatures::MQ.bits() != 0 {
                         vu.socket_handle()
                             .get_queue_num()
-                    .map_err(Error::VhostUserGetQueueMaxNum)? as usize
+                            .map_err(Error::VhostUserGetQueueMaxNum)?
+                            as usize
                     } else {
                         DEFAULT_QUEUE_NUMBER
                     };
@@ -160,21 +151,31 @@ impl Blk {
                     )
                     .map_err(Error::VhostUserGetConfig)?;
                 let mut config = VirtioBlockConfig::default();
-        if let Some(backend_config) = VirtioBlockConfig::from_slice(config_space.as_slice()) {
+                if let Some(backend_config) = VirtioBlockConfig::from_slice(config_space.as_slice())
+                {
                     config = *backend_config;
                     config.num_queues = num_queues as u16;
                 }
 
+                (
+                    acked_features,
+                    // If part of the available features that have been acked,
+                    // the PROTOCOL_FEATURES bit must be already set through
+                    // the VIRTIO acked features as we know the guest would
+                    // never ack it, thus the feature would be lost.
+                    acked_features & VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits(),
+                    acked_protocol_features,
+                    num_queues,
+                    config,
+                )
+            };
+
         Ok(Blk {
             common: VirtioCommon {
                 device_type: VirtioDeviceType::Block as u32,
                 queue_sizes: vec![vu_cfg.queue_size; num_queues],
-                avail_features: acked_features,
+                avail_features,
-                // If part of the available features that have been acked, the
+                acked_features,
-                // PROTOCOL_FEATURES bit must be already set through the VIRTIO
-                // acked features as we know the guest would never ack it, thus
-                // the feature would be lost.
-                acked_features: acked_features & VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits(),
                 paused_sync: Some(Arc::new(Barrier::new(2))),
                 min_queues: DEFAULT_QUEUE_NUMBER as u16,
                 ..Default::default()
@@ -183,7 +184,7 @@ impl Blk {
                 vu: Some(Arc::new(Mutex::new(vu))),
                 acked_protocol_features,
                 socket_path: vu_cfg.socket,
-                vu_num_queues: num_queues,
+                vu_num_queues,
                 ..Default::default()
             },
             id,
@@ -205,24 +206,6 @@ impl Blk {
             vu_num_queues: self.vu_common.vu_num_queues,
         }
     }
-
-    fn set_state(&mut self, state: &State) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        self.config = state.config;
-        self.vu_common.acked_protocol_features = state.acked_protocol_features;
-        self.vu_common.vu_num_queues = state.vu_num_queues;
-
-        if let Err(e) = self
-            .vu_common
-            .restore_backend_connection(self.common.acked_features)
-        {
-            error!(
-                "Failed restoring connection with vhost-user backend: {:?}",
-                e
-            );
-        }
-    }
 }
 
 impl Drop for Blk {
@@ -392,11 +375,6 @@ impl Snapshottable for Blk {
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
         self.vu_common.snapshot(&self.id(), &self.state())
     }
-
-    fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        self.set_state(&snapshot.to_versioned_state(&self.id)?);
-        Ok(())
-    }
 }
 impl Transportable for Blk {}
 

virtio-devices/src/vhost_user/fs.rs

@@ -315,47 +315,42 @@ impl Fs {
         queue_size: u16,
         cache: Option<(VirtioSharedMemoryList, MmapRegion)>,
         seccomp_action: SeccompAction,
-        restoring: bool,
         exit_evt: EventFd,
         iommu: bool,
+        state: Option<State>,
     ) -> Result<Fs> {
         let mut slave_req_support = false;
 
         // Calculate the actual number of queues needed.
         let num_queues = NUM_QUEUE_OFFSET + req_num_queues;
 
-        if restoring {
-            // We need 'queue_sizes' to report a number of queues that will be
-            // enough to handle all the potential queues. VirtioPciDevice::new()
-            // will create the actual queues based on this information.
-            return Ok(Fs {
-                common: VirtioCommon {
-                    device_type: VirtioDeviceType::Fs as u32,
-                    queue_sizes: vec![queue_size; num_queues],
-                    paused_sync: Some(Arc::new(Barrier::new(2))),
-                    min_queues: 1,
-                    ..Default::default()
-                },
-                vu_common: VhostUserCommon {
-                    socket_path: path.to_string(),
-                    vu_num_queues: num_queues,
-                    ..Default::default()
-                },
-                id,
-                config: VirtioFsConfig::default(),
-                cache,
-                slave_req_support,
-                seccomp_action,
-                guest_memory: None,
-                epoll_thread: None,
-                exit_evt,
-                iommu,
-            });
-        }
-
         // Connect to the vhost-user socket.
         let mut vu = VhostUserHandle::connect_vhost_user(false, path, num_queues as u64, false)?;
 
+        let (
+            avail_features,
+            acked_features,
+            acked_protocol_features,
+            vu_num_queues,
+            config,
+            slave_req_support,
+        ) = if let Some(state) = state {
+            info!("Restoring vhost-user-fs {}", id);
+
+            vu.set_protocol_features_vhost_user(
+                state.acked_features,
+                state.acked_protocol_features,
+            )?;
+
+            (
+                state.avail_features,
+                state.acked_features,
+                state.acked_protocol_features,
+                state.vu_num_queues,
+                state.config,
+                state.slave_req_support,
+            )
+        } else {
             // Filling device and vring features VMM supports.
             let avail_features = DEFAULT_VIRTIO_FEATURES;
 
@@ -402,15 +397,25 @@ impl Fs {
             config.tag[..tag_bytes_vec.len()].copy_from_slice(tag_bytes_vec.as_slice());
             config.num_request_queues = req_num_queues as u32;
 
-        Ok(Fs {
+            (
-            common: VirtioCommon {
+                acked_features,
-                device_type: VirtioDeviceType::Fs as u32,
-                avail_features: acked_features,
                 // If part of the available features that have been acked, the
                 // PROTOCOL_FEATURES bit must be already set through the VIRTIO
                 // acked features as we know the guest would never ack it, thus
                 // the feature would be lost.
-                acked_features: acked_features & VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits(),
+                acked_features & VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits(),
+                acked_protocol_features,
+                num_queues,
+                config,
+                slave_req_support,
+            )
+        };
+
+        Ok(Fs {
+            common: VirtioCommon {
+                device_type: VirtioDeviceType::Fs as u32,
+                avail_features,
+                acked_features,
                 queue_sizes: vec![queue_size; num_queues],
                 paused_sync: Some(Arc::new(Barrier::new(2))),
                 min_queues: 1,
@@ -420,7 +425,7 @@ impl Fs {
                 vu: Some(Arc::new(Mutex::new(vu))),
                 acked_protocol_features,
                 socket_path: path.to_string(),
-                vu_num_queues: num_queues,
+                vu_num_queues,
                 ..Default::default()
             },
             id,
@@ -445,25 +450,6 @@ impl Fs {
             slave_req_support: self.slave_req_support,
         }
     }
-
-    fn set_state(&mut self, state: &State) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        self.config = state.config;
-        self.vu_common.acked_protocol_features = state.acked_protocol_features;
-        self.vu_common.vu_num_queues = state.vu_num_queues;
-        self.slave_req_support = state.slave_req_support;
-
-        if let Err(e) = self
-            .vu_common
-            .restore_backend_connection(self.common.acked_features)
-        {
-            error!(
-                "Failed restoring connection with vhost-user backend: {:?}",
-                e
-            );
-        }
-    }
 }
 
 impl Drop for Fs {
@@ -663,11 +649,6 @@ impl Snapshottable for Fs {
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
         self.vu_common.snapshot(&self.id(), &self.state())
     }
-
-    fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        self.set_state(&snapshot.to_versioned_state(&self.id)?);
-        Ok(())
-    }
 }
 impl Transportable for Fs {}
 

virtio-devices/src/vhost_user/net.rs

@@ -74,42 +74,43 @@ impl Net {
         vu_cfg: VhostUserConfig,
         server: bool,
         seccomp_action: SeccompAction,
-        restoring: bool,
         exit_evt: EventFd,
         iommu: bool,
+        state: Option<State>,
     ) -> Result<Net> {
         let mut num_queues = vu_cfg.num_queues;
 
-        if restoring {
+        let mut vu =
-            // We need 'queue_sizes' to report a number of queues that will be
+            VhostUserHandle::connect_vhost_user(server, &vu_cfg.socket, num_queues as u64, false)?;
-            // enough to handle all the potential queues. Including the control
+
-            // queue (with +1) will guarantee that. VirtioPciDevice::new() will
+        let (avail_features, acked_features, acked_protocol_features, vu_num_queues, config) =
-            // create the actual queues based on this information.
+            if let Some(state) = state {
-            return Ok(Net {
+                info!("Restoring vhost-user-net {}", id);
-                common: VirtioCommon {
+
-                    device_type: VirtioDeviceType::Net as u32,
+                // The backend acknowledged features must not contain
-                    queue_sizes: vec![vu_cfg.queue_size; num_queues + 1],
+                // VIRTIO_NET_F_MAC since we don't expect the backend
-                    paused_sync: Some(Arc::new(Barrier::new(2))),
+                // to handle it.
-                    min_queues: DEFAULT_QUEUE_NUMBER as u16,
+                let backend_acked_features = state.acked_features & !(1 << VIRTIO_NET_F_MAC);
-                    ..Default::default()
+
-                },
+                vu.set_protocol_features_vhost_user(
-                vu_common: VhostUserCommon {
+                    backend_acked_features,
-                    socket_path: vu_cfg.socket,
+                    state.acked_protocol_features,
-                    vu_num_queues: num_queues,
+                )?;
-                    server,
+
-                    ..Default::default()
+                // If the control queue feature has been negotiated, let's
-                },
+                // increase the number of queues.
-                id,
+                if state.acked_features & (1 << VIRTIO_NET_F_CTRL_VQ) != 0 {
-                config: VirtioNetConfig::default(),
+                    num_queues += 1;
-                guest_memory: None,
-                ctrl_queue_epoll_thread: None,
-                epoll_thread: None,
-                seccomp_action,
-                exit_evt,
-                iommu,
-            });
                 }
 
+                (
+                    state.avail_features,
+                    state.acked_features,
+                    state.acked_protocol_features,
+                    state.vu_num_queues,
+                    state.config,
+                )
+            } else {
                 // Filling device and vring features VMM supports.
                 let mut avail_features = 1 << VIRTIO_NET_F_CSUM
                     | 1 << VIRTIO_NET_F_GUEST_CSUM
@@ -134,9 +135,6 @@ impl Net {
                 let mut config = VirtioNetConfig::default();
                 build_net_config_space(&mut config, mac_addr, num_queues, mtu, &mut avail_features);
 
-        let mut vu =
-            VhostUserHandle::connect_vhost_user(server, &vu_cfg.socket, num_queues as u64, false)?;
-
                 let avail_protocol_features = VhostUserProtocolFeatures::MQ
                     | VhostUserProtocolFeatures::CONFIGURE_MEM_SLOTS
                     | VhostUserProtocolFeatures::REPLY_ACK
@@ -150,7 +148,8 @@ impl Net {
                     if acked_protocol_features & VhostUserProtocolFeatures::MQ.bits() != 0 {
                         vu.socket_handle()
                             .get_queue_num()
-                    .map_err(Error::VhostUserGetQueueMaxNum)? as usize
+                            .map_err(Error::VhostUserGetQueueMaxNum)?
+                            as usize
                     } else {
                         DEFAULT_QUEUE_NUMBER
                     };
@@ -172,17 +171,26 @@ impl Net {
                 // the guest, even if it hasn't been negotiated with the backend.
                 acked_features |= 1 << VIRTIO_NET_F_MAC;
 
+                (
+                    acked_features,
+                    // If part of the available features that have been acked,
+                    // the PROTOCOL_FEATURES bit must be already set through
+                    // the VIRTIO acked features as we know the guest would
+                    // never ack it, thus the feature would be lost.
+                    acked_features & VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits(),
+                    acked_protocol_features,
+                    vu_num_queues,
+                    config,
+                )
+            };
+
         Ok(Net {
             id,
             common: VirtioCommon {
                 device_type: VirtioDeviceType::Net as u32,
                 queue_sizes: vec![vu_cfg.queue_size; num_queues],
-                avail_features: acked_features,
+                avail_features,
-                // If part of the available features that have been acked, the
+                acked_features,
-                // PROTOCOL_FEATURES bit must be already set through the VIRTIO
-                // acked features as we know the guest would never ack it, thus
-                // the feature would be lost.
-                acked_features: acked_features & VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits(),
                 paused_sync: Some(Arc::new(Barrier::new(2))),
                 min_queues: DEFAULT_QUEUE_NUMBER as u16,
                 ..Default::default()
@@ -214,28 +222,6 @@ impl Net {
             vu_num_queues: self.vu_common.vu_num_queues,
         }
     }
-
-    fn set_state(&mut self, state: &State) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        self.config = state.config;
-        self.vu_common.acked_protocol_features = state.acked_protocol_features;
-        self.vu_common.vu_num_queues = state.vu_num_queues;
-
-        // The backend acknowledged features must not contain VIRTIO_NET_F_MAC
-        // since we don't expect the backend to handle it.
-        let backend_acked_features = self.common.acked_features & !(1 << VIRTIO_NET_F_MAC);
-
-        if let Err(e) = self
-            .vu_common
-            .restore_backend_connection(backend_acked_features)
-        {
-            error!(
-                "Failed restoring connection with vhost-user backend: {:?}",
-                e
-            );
-        }
-    }
 }
 
 impl Drop for Net {
@@ -425,11 +411,6 @@ impl Snapshottable for Net {
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
         self.vu_common.snapshot(&self.id(), &self.state())
     }
-
-    fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        self.set_state(&snapshot.to_versioned_state(&self.id)?);
-        Ok(())
-    }
 }
 impl Transportable for Net {}
 

virtio-devices/src/vsock/device.rs

@@ -337,6 +337,7 @@ where
 {
     /// Create a new virtio-vsock device with the given VM CID and vsock
     /// backend.
+    #[allow(clippy::too_many_arguments)]
     pub fn new(
         id: String,
         cid: u64,
@@ -345,17 +346,25 @@ where
         iommu: bool,
         seccomp_action: SeccompAction,
         exit_evt: EventFd,
+        state: Option<VsockState>,
     ) -> io::Result<Vsock<B>> {
+        let (avail_features, acked_features) = if let Some(state) = state {
+            info!("Restoring virtio-vsock {}", id);
+            (state.avail_features, state.acked_features)
+        } else {
             let mut avail_features = 1u64 << VIRTIO_F_VERSION_1 | 1u64 << VIRTIO_F_IN_ORDER;
 
             if iommu {
                 avail_features |= 1u64 << VIRTIO_F_IOMMU_PLATFORM;
             }
+            (avail_features, 0)
+        };
 
         Ok(Vsock {
             common: VirtioCommon {
                 device_type: VirtioDeviceType::Vsock as u32,
                 avail_features,
+                acked_features,
                 paused_sync: Some(Arc::new(Barrier::new(2))),
                 queue_sizes: QUEUE_SIZES.to_vec(),
                 min_queues: NUM_QUEUES as u16,
@@ -376,11 +385,6 @@ where
             acked_features: self.common.acked_features,
         }
     }
-
-    fn set_state(&mut self, state: &VsockState) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-    }
 }
 
 impl<B> Drop for Vsock<B>
@@ -515,11 +519,6 @@ where
     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)?);
-        Ok(())
-    }
 }
 impl<B> Transportable for Vsock<B> where B: VsockBackend + Sync + 'static {}
 impl<B> Migratable for Vsock<B> where B: VsockBackend + Sync + 'static {}

virtio-devices/src/vsock/mod.rs

@@ -276,6 +276,7 @@ mod tests {
                     false,
                     seccompiler::SeccompAction::Trap,
                     EventFd::new(EFD_NONBLOCK).unwrap(),
+                    None,
                 )
                 .unwrap(),
             }

virtio-devices/src/watchdog.rs

@@ -213,26 +213,44 @@ impl Watchdog {
         reset_evt: EventFd,
         seccomp_action: SeccompAction,
         exit_evt: EventFd,
+        state: Option<WatchdogState>,
     ) -> io::Result<Watchdog> {
-        let avail_features = 1u64 << VIRTIO_F_VERSION_1;
+        let mut last_ping_time = None;
+        let (avail_features, acked_features) = if let Some(state) = state {
+            info!("Restoring virtio-watchdog {}", id);
+
+            // When restoring enable the watchdog if it was previously enabled.
+            // We reset the timer to ensure that we don't unnecessarily reboot
+            // due to the offline time.
+            if state.enabled {
+                last_ping_time = Some(Instant::now());
+            }
+
+            (state.avail_features, state.acked_features)
+        } else {
+            (1u64 << VIRTIO_F_VERSION_1, 0)
+        };
+
         let timer_fd = timerfd_create().map_err(|e| {
             error!("Failed to create timer fd {}", e);
             e
         })?;
         let timer = unsafe { File::from_raw_fd(timer_fd) };
+
         Ok(Watchdog {
             common: VirtioCommon {
                 device_type: VirtioDeviceType::Watchdog as u32,
                 queue_sizes: QUEUE_SIZES.to_vec(),
                 paused_sync: Some(Arc::new(Barrier::new(2))),
                 avail_features,
+                acked_features,
                 min_queues: 1,
                 ..Default::default()
             },
             id,
             seccomp_action,
             reset_evt,
-            last_ping_time: Arc::new(Mutex::new(None)),
+            last_ping_time: Arc::new(Mutex::new(last_ping_time)),
             timer,
             exit_evt,
         })
@@ -246,16 +264,6 @@ impl Watchdog {
         }
     }
 
-    fn set_state(&mut self, state: &WatchdogState) {
-        self.common.avail_features = state.avail_features;
-        self.common.acked_features = state.acked_features;
-        // When restoring enable the watchdog if it was previously enabled. We reset the timer
-        // to ensure that we don't unnecessarily reboot due to the offline time.
-        if state.enabled {
-            self.last_ping_time.lock().unwrap().replace(Instant::now());
-        }
-    }
-
     #[cfg(fuzzing)]
     pub fn wait_for_epoll_threads(&mut self) {
         self.common.wait_for_epoll_threads();
@@ -409,11 +417,6 @@ impl Snapshottable for Watchdog {
     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)?);
-        Ok(())
-    }
 }
 
 impl Transportable for Watchdog {}

vm-migration/src/lib.rs

@@ -236,6 +236,23 @@ impl Snapshot {
     }
 }
 
+pub fn snapshot_from_id(snapshot: Option<&Snapshot>, id: &str) -> Option<Snapshot> {
+    snapshot.and_then(|s| s.snapshots.get(id).map(|s| *s.clone()))
+}
+
+pub fn versioned_state_from_id<T>(
+    snapshot: Option<&Snapshot>,
+    id: &str,
+) -> Result<Option<T>, MigratableError>
+where
+    T: Versionize + VersionMapped,
+{
+    snapshot
+        .and_then(|s| s.snapshots.get(id).map(|s| *s.clone()))
+        .map(|s| s.to_versioned_state(id))
+        .transpose()
+}
+
 /// A snapshottable component can be snapshotted.
 pub trait Snapshottable: Pausable {
     /// The snapshottable component id.

vmm/src/device_manager.rs

@@ -96,8 +96,8 @@ use vm_memory::{Address, GuestAddress, GuestUsize, MmapRegion};
 #[cfg(target_arch = "x86_64")]
 use vm_memory::{GuestAddressSpace, GuestMemory};
 use vm_migration::{
-    protocol::MemoryRangeTable, Migratable, MigratableError, Pausable, Snapshot,
+    protocol::MemoryRangeTable, snapshot_from_id, versioned_state_from_id, Migratable,
-    SnapshotDataSection, Snapshottable, Transportable,
+    MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable, Transportable,
 };
 use vm_virtio::AccessPlatform;
 use vm_virtio::VirtioDeviceType;
@@ -466,6 +466,9 @@ pub enum DeviceManagerError {
 
     /// Error activating virtio device
     VirtioActivate(ActivateError),
+
+    /// Failed retrieving device state from snapshot
+    RestoreGetState(MigratableError),
 }
 pub type DeviceManagerResult<T> = result::Result<T, DeviceManagerError>;
 
@@ -940,6 +943,8 @@ pub struct DeviceManager {
 
     // Addresses for ACPI platform devices e.g. ACPI PM timer, sleep/reset registers
     acpi_platform_addresses: AcpiPlatformAddresses,
+
+    snapshot: Option<Snapshot>,
 }
 
 impl DeviceManager {
@@ -958,6 +963,7 @@ impl DeviceManager {
         restoring: bool,
         boot_id_list: BTreeSet<String>,
         timestamp: Instant,
+        snapshot: Option<Snapshot>,
     ) -> DeviceManagerResult<Arc<Mutex<Self>>> {
         trace_scoped!("DeviceManager::new");
 
@@ -1085,6 +1091,7 @@ impl DeviceManager {
             timestamp,
             pending_activations: Arc::new(Mutex::new(Vec::default())),
             acpi_platform_addresses: AcpiPlatformAddresses::default(),
+            snapshot,
         };
 
         let device_manager = Arc::new(Mutex::new(device_manager));
@@ -1238,6 +1245,8 @@ impl DeviceManager {
                     .try_clone()
                     .map_err(DeviceManagerError::EventFd)?,
                 self.get_msi_iova_space(),
+                versioned_state_from_id(self.snapshot.as_ref(), iommu_id.as_str())
+                    .map_err(DeviceManagerError::RestoreGetState)?,
             )
             .map_err(DeviceManagerError::CreateVirtioIommu)?;
             let device = Arc::new(Mutex::new(device));
@@ -1886,6 +1895,8 @@ impl DeviceManager {
             self.exit_evt
                 .try_clone()
                 .map_err(DeviceManagerError::EventFd)?,
+            versioned_state_from_id(self.snapshot.as_ref(), id.as_str())
+                .map_err(DeviceManagerError::RestoreGetState)?,
         )
         .map_err(DeviceManagerError::CreateVirtioConsole)?;
         let virtio_console_device = Arc::new(Mutex::new(virtio_console_device));
@@ -2031,6 +2042,8 @@ impl DeviceManager {
 
         info!("Creating virtio-block device: {:?}", disk_cfg);
 
+        let snapshot = snapshot_from_id(self.snapshot.as_ref(), id.as_str());
+
         let (virtio_device, migratable_device) = if disk_cfg.vhost_user {
             let socket = disk_cfg.vhost_socket.as_ref().unwrap().clone();
             let vu_cfg = VhostUserConfig {
@@ -2042,12 +2055,15 @@ impl DeviceManager {
                 match virtio_devices::vhost_user::Blk::new(
                     id.clone(),
                     vu_cfg,
-                    self.restoring,
                     self.seccomp_action.clone(),
                     self.exit_evt
                         .try_clone()
                         .map_err(DeviceManagerError::EventFd)?,
                     self.force_iommu,
+                    snapshot
+                        .map(|s| s.to_versioned_state(&id))
+                        .transpose()
+                        .map_err(DeviceManagerError::RestoreGetState)?,
                 ) {
                     Ok(vub_device) => vub_device,
                     Err(e) => {
@@ -2143,6 +2159,10 @@ impl DeviceManager {
                     self.exit_evt
                         .try_clone()
                         .map_err(DeviceManagerError::EventFd)?,
+                    snapshot
+                        .map(|s| s.to_versioned_state(&id))
+                        .transpose()
+                        .map_err(DeviceManagerError::RestoreGetState)?,
                 )
                 .map_err(DeviceManagerError::CreateVirtioBlock)?,
             ));
@@ -2197,6 +2217,8 @@ impl DeviceManager {
         };
         info!("Creating virtio-net device: {:?}", net_cfg);
 
+        let snapshot = snapshot_from_id(self.snapshot.as_ref(), id.as_str());
+
         let (virtio_device, migratable_device) = if net_cfg.vhost_user {
             let socket = net_cfg.vhost_socket.as_ref().unwrap().clone();
             let vu_cfg = VhostUserConfig {
@@ -2216,11 +2238,14 @@ impl DeviceManager {
                     vu_cfg,
                     server,
                     self.seccomp_action.clone(),
-                    self.restoring,
                     self.exit_evt
                         .try_clone()
                         .map_err(DeviceManagerError::EventFd)?,
                     self.force_iommu,
+                    snapshot
+                        .map(|s| s.to_versioned_state(&id))
+                        .transpose()
+                        .map_err(DeviceManagerError::RestoreGetState)?,
                 ) {
                     Ok(vun_device) => vun_device,
                     Err(e) => {
@@ -2234,6 +2259,11 @@ impl DeviceManager {
                 vhost_user_net as Arc<Mutex<dyn Migratable>>,
             )
         } else {
+            let state = snapshot
+                .map(|s| s.to_versioned_state(&id))
+                .transpose()
+                .map_err(DeviceManagerError::RestoreGetState)?;
+
             let virtio_net = if let Some(ref tap_if_name) = net_cfg.tap {
                 Arc::new(Mutex::new(
                     virtio_devices::Net::new(
@@ -2252,6 +2282,7 @@ impl DeviceManager {
                         self.exit_evt
                             .try_clone()
                             .map_err(DeviceManagerError::EventFd)?,
+                        state,
                     )
                     .map_err(DeviceManagerError::CreateVirtioNet)?,
                 ))
@@ -2269,6 +2300,7 @@ impl DeviceManager {
                         self.exit_evt
                             .try_clone()
                             .map_err(DeviceManagerError::EventFd)?,
+                        state,
                     )
                     .map_err(DeviceManagerError::CreateVirtioNet)?,
                 ))
@@ -2290,6 +2322,7 @@ impl DeviceManager {
                         self.exit_evt
                             .try_clone()
                             .map_err(DeviceManagerError::EventFd)?,
+                        state,
                     )
                     .map_err(DeviceManagerError::CreateVirtioNet)?,
                 ))
@@ -2350,6 +2383,8 @@ impl DeviceManager {
                     self.exit_evt
                         .try_clone()
                         .map_err(DeviceManagerError::EventFd)?,
+                    versioned_state_from_id(self.snapshot.as_ref(), id.as_str())
+                        .map_err(DeviceManagerError::RestoreGetState)?,
                 )
                 .map_err(DeviceManagerError::CreateVirtioRng)?,
             ));
@@ -2400,11 +2435,12 @@ impl DeviceManager {
                     fs_cfg.queue_size,
                     None,
                     self.seccomp_action.clone(),
-                    self.restoring,
                     self.exit_evt
                         .try_clone()
                         .map_err(DeviceManagerError::EventFd)?,
                     self.force_iommu,
+                    versioned_state_from_id(self.snapshot.as_ref(), id.as_str())
+                        .map_err(DeviceManagerError::RestoreGetState)?,
                 )
                 .map_err(DeviceManagerError::CreateVirtioFs)?,
             ));
@@ -2587,6 +2623,8 @@ impl DeviceManager {
                 self.exit_evt
                     .try_clone()
                     .map_err(DeviceManagerError::EventFd)?,
+                versioned_state_from_id(self.snapshot.as_ref(), id.as_str())
+                    .map_err(DeviceManagerError::RestoreGetState)?,
             )
             .map_err(DeviceManagerError::CreateVirtioPmem)?,
         ));
@@ -2657,6 +2695,8 @@ impl DeviceManager {
                 self.exit_evt
                     .try_clone()
                     .map_err(DeviceManagerError::EventFd)?,
+                versioned_state_from_id(self.snapshot.as_ref(), id.as_str())
+                    .map_err(DeviceManagerError::RestoreGetState)?,
             )
             .map_err(DeviceManagerError::CreateVirtioVsock)?,
         ));
@@ -2715,6 +2755,8 @@ impl DeviceManager {
                             .try_clone()
                             .map_err(DeviceManagerError::EventFd)?,
                         virtio_mem_zone.blocks_state().clone(),
+                        versioned_state_from_id(self.snapshot.as_ref(), memory_zone_id.as_str())
+                            .map_err(DeviceManagerError::RestoreGetState)?,
                     )
                     .map_err(DeviceManagerError::CreateVirtioMem)?,
                 ));
@@ -2765,6 +2807,8 @@ impl DeviceManager {
                     self.exit_evt
                         .try_clone()
                         .map_err(DeviceManagerError::EventFd)?,
+                    versioned_state_from_id(self.snapshot.as_ref(), id.as_str())
+                        .map_err(DeviceManagerError::RestoreGetState)?,
                 )
                 .map_err(DeviceManagerError::CreateVirtioBalloon)?,
             ));
@@ -2807,6 +2851,8 @@ impl DeviceManager {
                 self.exit_evt
                     .try_clone()
                     .map_err(DeviceManagerError::EventFd)?,
+                versioned_state_from_id(self.snapshot.as_ref(), id.as_str())
+                    .map_err(DeviceManagerError::RestoreGetState)?,
             )
             .map_err(DeviceManagerError::CreateVirtioWatchdog)?,
         ));
@@ -2852,7 +2898,8 @@ impl DeviceManager {
                 device_path,
                 self.memory_manager.lock().unwrap().guest_memory(),
                 vdpa_cfg.num_queues as u16,
-                self.restoring,
+                versioned_state_from_id(self.snapshot.as_ref(), id.as_str())
+                    .map_err(DeviceManagerError::RestoreGetState)?,
             )
             .map_err(DeviceManagerError::CreateVdpa)?,
         ));

vmm/src/lib.rs

@@ -1232,6 +1232,7 @@ impl Vmm {
             activate_evt,
             true,
             timestamp,
+            Some(&snapshot),
         )
         .map_err(|e| {
             MigratableError::MigrateReceive(anyhow!("Error creating VM from snapshot: {:?}", e))

vmm/src/vm.rs

@@ -94,7 +94,7 @@ use vm_memory::{Address, ByteValued, GuestMemory, GuestMemoryRegion};
 use vm_memory::{Bytes, GuestAddress, GuestAddressSpace, GuestMemoryAtomic};
 use vm_migration::protocol::{Request, Response, Status};
 use vm_migration::{
-    protocol::MemoryRangeTable, Migratable, MigratableError, Pausable, Snapshot,
+    protocol::MemoryRangeTable, snapshot_from_id, Migratable, MigratableError, Pausable, Snapshot,
     SnapshotDataSection, Snapshottable, Transportable,
 };
 use vmm_sys_util::eventfd::EventFd;
@@ -497,6 +497,7 @@ impl Vm {
         activate_evt: EventFd,
         restoring: bool,
         timestamp: Instant,
+        snapshot: Option<&Snapshot>,
     ) -> Result<Self> {
         trace_scoped!("Vm::new_from_memory_manager");
 
@@ -544,6 +545,7 @@ impl Vm {
             restoring,
             boot_id_list,
             timestamp,
+            snapshot_from_id(snapshot, DEVICE_MANAGER_SNAPSHOT_ID),
         )
         .map_err(Error::DeviceManager)?;
 
@@ -769,6 +771,7 @@ impl Vm {
             activate_evt,
             false,
             timestamp,
+            None,
         )?;
 
         // The device manager must create the devices from here as it is part
@@ -835,6 +838,7 @@ impl Vm {
             activate_evt,
             true,
             timestamp,
+            Some(snapshot),
         )
     }