Skip to content

Commit 0fcb5ec

Browse files
ManciukicManciukic
committed
feat(virtio-mem): implement snapshot/restore
Implements basic snapshot/restore functionality for the dummy virtio-mem device. Signed-off-by: Riccardo Mancini <mancio@amazon.com>
1 parent 6e57c75 commit 0fcb5ec

File tree

6 files changed

+276
-6
lines changed

6 files changed

+276
-6
lines changed

‎src/vmm/src/device_manager/pci_mngr.rs‎

Lines changed: 59 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@ use crate::devices::virtio::block::device::Block;
2020
use crate::devices::virtio::block::persist::{BlockConstructorArgs, BlockState};
2121
use crate::devices::virtio::device::VirtioDevice;
2222
use crate::devices::virtio::generated::virtio_ids;
23+
use crate::devices::virtio::mem::VirtioMem;
24+
use crate::devices::virtio::mem::persist::{VirtioMemConstructorArgs, VirtioMemState};
2325
use crate::devices::virtio::net::Net;
2426
use crate::devices::virtio::net::persist::{NetConstructorArgs, NetState};
2527
use crate::devices::virtio::rng::Entropy;
@@ -240,6 +242,8 @@ pub struct PciDevicesState {
240242
pub mmds: Option<MmdsState>,
241243
/// Entropy device state.
242244
pub entropy_device: Option<VirtioDeviceState<EntropyState>>,
245+
/// Memory device state.
246+
pub memory_device: Option<VirtioDeviceState<VirtioMemState>>,
243247
}
244248

245249
pub struct PciDevicesConstructorArgs<'a> {
@@ -388,6 +392,20 @@ impl<'a> Persist<'a> for PciDevices {
388392
transport_state,
389393
})
390394
}
395+
virtio_ids::VIRTIO_ID_MEM => {
396+
let mem_dev = locked_virtio_dev
397+
.as_mut_any()
398+
.downcast_mut::<VirtioMem>()
399+
.unwrap();
400+
let device_state = mem_dev.save();
401+
402+
state.memory_device = Some(VirtioDeviceState {
403+
device_id: mem_dev.id().to_string(),
404+
pci_device_bdf,
405+
device_state,
406+
transport_state,
407+
})
408+
}
391409
_ => unreachable!(),
392410
}
393411
}
@@ -566,6 +584,29 @@ impl<'a> Persist<'a> for PciDevices {
566584
.unwrap()
567585
}
568586

587+
if let Some(memory_device) = &state.memory_device {
588+
let ctor_args = VirtioMemConstructorArgs::new(Arc::clone(constructor_args.vm));
589+
590+
let device = Arc::new(Mutex::new(
591+
VirtioMem::restore(ctor_args, &memory_device.device_state).unwrap(),
592+
));
593+
594+
constructor_args
595+
.vm_resources
596+
.update_from_restored_device(SharedDeviceType::VirtioMem(device.clone()))
597+
.unwrap();
598+
599+
pci_devices
600+
.restore_pci_device(
601+
constructor_args.vm,
602+
device,
603+
&memory_device.device_id,
604+
&memory_device.transport_state,
605+
constructor_args.event_manager,
606+
)
607+
.unwrap()
608+
}
609+
569610
Ok(pci_devices)
570611
}
571612
}
@@ -583,6 +624,7 @@ mod tests {
583624
use crate::snapshot::Snapshot;
584625
use crate::vmm_config::balloon::BalloonDeviceConfig;
585626
use crate::vmm_config::entropy::EntropyDeviceConfig;
627+
use crate::vmm_config::memory_hotplug::MemoryHotplugConfig;
586628
use crate::vmm_config::net::NetworkInterfaceConfig;
587629
use crate::vmm_config::vsock::VsockDeviceConfig;
588630

@@ -645,6 +687,18 @@ mod tests {
645687
let entropy_config = EntropyDeviceConfig::default();
646688
insert_entropy_device(&mut vmm, &mut cmdline, &mut event_manager, entropy_config);
647689

690+
let memory_hotplug_config = MemoryHotplugConfig {
691+
total_size_mib: 1024,
692+
block_size_mib: 2,
693+
slot_size_mib: 128,
694+
};
695+
insert_virtio_mem_device(
696+
&mut vmm,
697+
&mut cmdline,
698+
&mut event_manager,
699+
memory_hotplug_config,
700+
);
701+
648702
Snapshot::new(vmm.device_manager.save())
649703
.save(&mut buf.as_mut_slice())
650704
.unwrap();
@@ -674,7 +728,6 @@ mod tests {
674728
let _restored_dev_manager =
675729
PciDevices::restore(restore_args, &device_manager_state.pci_state).unwrap();
676730

677-
// TODO(virtio-mem): add memory-hotplug device when snapshot-restore is implemented
678731
let expected_vm_resources = format!(
679732
r#"{{
680733
"balloon": {{
@@ -734,7 +787,11 @@ mod tests {
734787
"entropy": {{
735788
"rate_limiter": null
736789
}},
737-
"memory-hotplug": null
790+
"memory-hotplug": {{
791+
"total_size_mib": 1024,
792+
"block_size_mib": 2,
793+
"slot_size_mib": 128
794+
}}
738795
}}"#,
739796
_block_files.last().unwrap().as_path().to_str().unwrap(),
740797
tmp_sock_file.as_path().to_str().unwrap()

‎src/vmm/src/device_manager/persist.rs‎

Lines changed: 66 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,10 @@ use crate::devices::virtio::block::device::Block;
2525
use crate::devices::virtio::block::persist::{BlockConstructorArgs, BlockState};
2626
use crate::devices::virtio::device::VirtioDevice;
2727
use crate::devices::virtio::generated::virtio_ids;
28+
use crate::devices::virtio::mem::VirtioMem;
29+
use crate::devices::virtio::mem::persist::{
30+
VirtioMemConstructorArgs, VirtioMemPersistError, VirtioMemState,
31+
};
2832
use crate::devices::virtio::net::Net;
2933
use crate::devices::virtio::net::persist::{
3034
NetConstructorArgs, NetPersistError as NetError, NetState,
@@ -72,6 +76,8 @@ pub enum DevicePersistError {
7276
MmdsConfig(#[from] MmdsConfigError),
7377
/// Entropy: {0}
7478
Entropy(#[from] EntropyError),
79+
/// virtio-mem: {0}
80+
VirtioMem(#[from] VirtioMemPersistError),
7581
/// Resource misconfiguration: {0}. Is the snapshot file corrupted?
7682
ResourcesError(#[from] ResourcesError),
7783
/// Could not activate device: {0}
@@ -125,6 +131,8 @@ pub struct DeviceStates {
125131
pub mmds: Option<MmdsState>,
126132
/// Entropy device state.
127133
pub entropy_device: Option<VirtioDeviceState<EntropyState>>,
134+
/// Memory device state.
135+
pub memory_device: Option<VirtioDeviceState<VirtioMemState>>,
128136
}
129137

130138
/// A type used to extract the concrete `Arc<Mutex<T>>` for each of the device
@@ -136,6 +144,7 @@ pub enum SharedDeviceType {
136144
Balloon(Arc<Mutex<Balloon>>),
137145
Vsock(Arc<Mutex<Vsock<VsockUnixBackend>>>),
138146
Entropy(Arc<Mutex<Entropy>>),
147+
VirtioMem(Arc<Mutex<VirtioMem>>),
139148
}
140149

141150
pub struct MMIODevManagerConstructorArgs<'a> {
@@ -335,6 +344,20 @@ impl<'a> Persist<'a> for MMIODeviceManager {
335344
device_info,
336345
});
337346
}
347+
virtio_ids::VIRTIO_ID_MEM => {
348+
let mem = locked_device
349+
.as_mut_any()
350+
.downcast_mut::<VirtioMem>()
351+
.unwrap();
352+
let device_state = mem.save();
353+
354+
states.memory_device = Some(VirtioDeviceState {
355+
device_id,
356+
device_state,
357+
transport_state,
358+
device_info,
359+
});
360+
}
338361
_ => unreachable!(),
339362
};
340363

@@ -549,6 +572,30 @@ impl<'a> Persist<'a> for MMIODeviceManager {
549572
)?;
550573
}
551574

575+
if let Some(memory_state) = &state.memory_device {
576+
let ctor_args = VirtioMemConstructorArgs::new(Arc::clone(vm));
577+
578+
let device = Arc::new(Mutex::new(VirtioMem::restore(
579+
ctor_args,
580+
&memory_state.device_state,
581+
)?));
582+
583+
constructor_args
584+
.vm_resources
585+
.update_from_restored_device(SharedDeviceType::VirtioMem(device.clone()))?;
586+
587+
restore_helper(
588+
device.clone(),
589+
memory_state.device_state.virtio_state.activated,
590+
false,
591+
device,
592+
&memory_state.device_id,
593+
&memory_state.transport_state,
594+
&memory_state.device_info,
595+
constructor_args.event_manager,
596+
)?;
597+
}
598+
552599
Ok(dev_manager)
553600
}
554601
}
@@ -565,6 +612,7 @@ mod tests {
565612
use crate::snapshot::Snapshot;
566613
use crate::vmm_config::balloon::BalloonDeviceConfig;
567614
use crate::vmm_config::entropy::EntropyDeviceConfig;
615+
use crate::vmm_config::memory_hotplug::MemoryHotplugConfig;
568616
use crate::vmm_config::net::NetworkInterfaceConfig;
569617
use crate::vmm_config::vsock::VsockDeviceConfig;
570618

@@ -582,6 +630,7 @@ mod tests {
582630
&& self.net_devices == other.net_devices
583631
&& self.vsock_device == other.vsock_device
584632
&& self.entropy_device == other.entropy_device
633+
&& self.memory_device == other.memory_device
585634
}
586635
}
587636

@@ -666,6 +715,18 @@ mod tests {
666715
let entropy_config = EntropyDeviceConfig::default();
667716
insert_entropy_device(&mut vmm, &mut cmdline, &mut event_manager, entropy_config);
668717

718+
let memory_hotplug_config = MemoryHotplugConfig {
719+
total_size_mib: 1024,
720+
block_size_mib: 2,
721+
slot_size_mib: 128,
722+
};
723+
insert_virtio_mem_device(
724+
&mut vmm,
725+
&mut cmdline,
726+
&mut event_manager,
727+
memory_hotplug_config,
728+
);
729+
669730
Snapshot::new(vmm.device_manager.save())
670731
.save(&mut buf.as_mut_slice())
671732
.unwrap();
@@ -691,7 +752,6 @@ mod tests {
691752
let _restored_dev_manager =
692753
MMIODeviceManager::restore(restore_args, &device_manager_state.mmio_state).unwrap();
693754

694-
// TODO(virtio-mem): add memory-hotplug device when snapshot-restore is implemented
695755
let expected_vm_resources = format!(
696756
r#"{{
697757
"balloon": {{
@@ -751,7 +811,11 @@ mod tests {
751811
"entropy": {{
752812
"rate_limiter": null
753813
}},
754-
"memory-hotplug": null
814+
"memory-hotplug": {{
815+
"total_size_mib": 1024,
816+
"block_size_mib": 2,
817+
"slot_size_mib": 128
818+
}}
755819
}}"#,
756820
_block_files.last().unwrap().as_path().to_str().unwrap(),
757821
tmp_sock_file.as_path().to_str().unwrap()

‎src/vmm/src/devices/virtio/mem/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
mod device;
55
mod event_handler;
6+
pub mod persist;
67

78
use vm_memory::GuestAddress;
89

0 commit comments

Comments
 (0)