vmm/devices/virtio/pmem/
persist.rs

1// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use serde::{Deserialize, Serialize};
5use vm_memory::GuestAddress;
6
7use super::device::{ConfigSpace, Pmem, PmemError};
8use crate::Vm;
9use crate::devices::virtio::device::DeviceState;
10use crate::devices::virtio::generated::virtio_ids::VIRTIO_ID_PMEM;
11use crate::devices::virtio::persist::{PersistError as VirtioStateError, VirtioDeviceState};
12use crate::devices::virtio::pmem::{PMEM_NUM_QUEUES, PMEM_QUEUE_SIZE};
13use crate::snapshot::Persist;
14use crate::vmm_config::pmem::PmemConfig;
15use crate::vstate::memory::{GuestMemoryMmap, GuestRegionMmap};
16use crate::vstate::vm::VmError;
17
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct PmemState {
20    pub virtio_state: VirtioDeviceState,
21    pub config_space: ConfigSpace,
22    pub config: PmemConfig,
23}
24
25#[derive(Debug)]
26pub struct PmemConstructorArgs<'a> {
27    pub mem: &'a GuestMemoryMmap,
28    pub vm: &'a Vm,
29}
30
31#[derive(Debug, thiserror::Error, displaydoc::Display)]
32pub enum PmemPersistError {
33    /// Error resetting VirtIO state: {0}
34    VirtioState(#[from] VirtioStateError),
35    /// Error creating Pmem devie: {0}
36    Pmem(#[from] PmemError),
37    /// Error registering memory region: {0}
38    Vm(#[from] VmError),
39}
40
41impl<'a> Persist<'a> for Pmem {
42    type State = PmemState;
43    type ConstructorArgs = PmemConstructorArgs<'a>;
44    type Error = PmemPersistError;
45
46    fn save(&self) -> Self::State {
47        PmemState {
48            virtio_state: VirtioDeviceState::from_device(self),
49            config_space: self.config_space,
50            config: self.config.clone(),
51        }
52    }
53
54    fn restore(
55        constructor_args: Self::ConstructorArgs,
56        state: &Self::State,
57    ) -> Result<Self, Self::Error> {
58        let queues = state.virtio_state.build_queues_checked(
59            constructor_args.mem,
60            VIRTIO_ID_PMEM,
61            PMEM_NUM_QUEUES,
62            PMEM_QUEUE_SIZE,
63        )?;
64
65        let mut pmem = Pmem::new_with_queues(state.config.clone(), queues)?;
66        pmem.config_space = state.config_space;
67        pmem.avail_features = state.virtio_state.avail_features;
68        pmem.acked_features = state.virtio_state.acked_features;
69
70        pmem.set_mem_region(constructor_args.vm)?;
71
72        Ok(pmem)
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use vmm_sys_util::tempfile::TempFile;
79
80    use super::*;
81    use crate::arch::Kvm;
82    use crate::devices::virtio::device::VirtioDevice;
83    use crate::devices::virtio::test_utils::default_mem;
84    use crate::snapshot::Snapshot;
85
86    #[test]
87    fn test_persistence() {
88        // We create the backing file here so that it exists for the whole lifetime of the test.
89        let dummy_file = TempFile::new().unwrap();
90        dummy_file.as_file().set_len(0x20_0000);
91        let dummy_path = dummy_file.as_path().to_str().unwrap().to_string();
92        let config = PmemConfig {
93            id: "1".into(),
94            path_on_host: dummy_path,
95            root_device: true,
96            read_only: false,
97        };
98        let pmem = Pmem::new(config).unwrap();
99        let guest_mem = default_mem();
100        let kvm = Kvm::new(vec![]).unwrap();
101        let vm = Vm::new(&kvm).unwrap();
102
103        // Save the block device.
104        let mut mem = vec![0; 4096];
105
106        Snapshot::new(pmem.save())
107            .save(&mut mem.as_mut_slice())
108            .unwrap();
109
110        // Restore the block device.
111        let restored_pmem = Pmem::restore(
112            PmemConstructorArgs {
113                mem: &guest_mem,
114                vm: &vm,
115            },
116            &Snapshot::load_without_crc_check(mem.as_slice())
117                .unwrap()
118                .data,
119        )
120        .unwrap();
121
122        // Test that virtio specific fields are the same.
123        assert_eq!(restored_pmem.device_type(), VIRTIO_ID_PMEM);
124        assert_eq!(restored_pmem.avail_features(), pmem.avail_features());
125        assert_eq!(restored_pmem.acked_features(), pmem.acked_features());
126        assert_eq!(restored_pmem.queues(), pmem.queues());
127        assert!(!pmem.is_activated());
128        assert!(!restored_pmem.is_activated());
129        assert_eq!(restored_pmem.config, pmem.config);
130    }
131}