vmm/devices/virtio/mem/
persist.rs

1// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Defines the structures needed for saving/restoring virtio-mem devices.
5
6use std::sync::Arc;
7
8use bitvec::vec::BitVec;
9use serde::{Deserialize, Serialize};
10use vm_memory::Address;
11
12use crate::Vm;
13use crate::devices::virtio::generated::virtio_ids::VIRTIO_ID_MEM;
14use crate::devices::virtio::generated::virtio_mem::virtio_mem_config;
15use crate::devices::virtio::mem::{MEM_NUM_QUEUES, VirtioMem, VirtioMemError};
16use crate::devices::virtio::persist::{PersistError as VirtioStateError, VirtioDeviceState};
17use crate::devices::virtio::queue::FIRECRACKER_MAX_QUEUE_SIZE;
18use crate::snapshot::Persist;
19use crate::utils::usize_to_u64;
20use crate::vstate::memory::{GuestMemoryMmap, GuestRegionMmap};
21
22#[derive(Debug, Clone, Serialize, Deserialize)]
23pub struct VirtioMemState {
24    pub virtio_state: VirtioDeviceState,
25    addr: u64,
26    region_size: u64,
27    block_size: u64,
28    usable_region_size: u64,
29    requested_size: u64,
30    slot_size: usize,
31    plugged_blocks: Vec<bool>,
32}
33
34#[derive(Debug)]
35pub struct VirtioMemConstructorArgs {
36    vm: Arc<Vm>,
37}
38
39impl VirtioMemConstructorArgs {
40    pub fn new(vm: Arc<Vm>) -> Self {
41        Self { vm }
42    }
43}
44
45#[derive(Debug, thiserror::Error, displaydoc::Display)]
46pub enum VirtioMemPersistError {
47    /// Create virtio-mem: {0}
48    CreateVirtioMem(#[from] VirtioMemError),
49    /// Virtio state: {0}
50    VirtioState(#[from] VirtioStateError),
51}
52
53impl Persist<'_> for VirtioMem {
54    type State = VirtioMemState;
55    type ConstructorArgs = VirtioMemConstructorArgs;
56    type Error = VirtioMemPersistError;
57
58    fn save(&self) -> Self::State {
59        VirtioMemState {
60            virtio_state: VirtioDeviceState::from_device(self),
61            addr: self.config.addr,
62            region_size: self.config.region_size,
63            block_size: self.config.block_size,
64            usable_region_size: self.config.usable_region_size,
65            plugged_blocks: self.plugged_blocks.iter().by_vals().collect(),
66            requested_size: self.config.requested_size,
67            slot_size: self.slot_size,
68        }
69    }
70
71    fn restore(
72        constructor_args: Self::ConstructorArgs,
73        state: &Self::State,
74    ) -> Result<Self, Self::Error> {
75        let queues = state.virtio_state.build_queues_checked(
76            constructor_args.vm.guest_memory(),
77            VIRTIO_ID_MEM,
78            MEM_NUM_QUEUES,
79            FIRECRACKER_MAX_QUEUE_SIZE,
80        )?;
81
82        let plugged_blocks = BitVec::from_iter(state.plugged_blocks.iter());
83
84        let config = virtio_mem_config {
85            addr: state.addr,
86            region_size: state.region_size,
87            block_size: state.block_size,
88            usable_region_size: state.usable_region_size,
89            plugged_size: usize_to_u64(plugged_blocks.count_ones()) * state.block_size,
90            requested_size: state.requested_size,
91            ..Default::default()
92        };
93
94        let mut virtio_mem = VirtioMem::from_state(
95            constructor_args.vm,
96            queues,
97            config,
98            state.slot_size,
99            plugged_blocks,
100        )?;
101        virtio_mem.set_avail_features(state.virtio_state.avail_features);
102        virtio_mem.set_acked_features(state.virtio_state.acked_features);
103
104        Ok(virtio_mem)
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111    use crate::devices::virtio::device::VirtioDevice;
112    use crate::devices::virtio::mem::device::test_utils::default_virtio_mem;
113    use crate::vstate::vm::tests::setup_vm_with_memory;
114
115    #[test]
116    fn test_save_state() {
117        let dev = default_virtio_mem();
118        let state = dev.save();
119
120        assert_eq!(state.addr, dev.config.addr);
121        assert_eq!(state.region_size, dev.config.region_size);
122        assert_eq!(state.block_size, dev.config.block_size);
123        assert_eq!(state.usable_region_size, dev.config.usable_region_size);
124        assert_eq!(
125            state.plugged_blocks.iter().collect::<BitVec>(),
126            dev.plugged_blocks
127        );
128        assert_eq!(state.requested_size, dev.config.requested_size);
129        assert_eq!(state.slot_size, dev.slot_size);
130    }
131
132    #[test]
133    fn test_save_restore_state() {
134        let mut original_dev = default_virtio_mem();
135        original_dev.set_acked_features(original_dev.avail_features());
136        let state = original_dev.save();
137
138        // Create a "new" VM for restore
139        let (_, vm) = setup_vm_with_memory(0x1000);
140        let vm = Arc::new(vm);
141        let constructor_args = VirtioMemConstructorArgs::new(vm);
142        let restored_dev = VirtioMem::restore(constructor_args, &state).unwrap();
143
144        assert_eq!(original_dev.config, restored_dev.config);
145        assert_eq!(original_dev.slot_size, restored_dev.slot_size);
146        assert_eq!(original_dev.avail_features(), restored_dev.avail_features());
147        assert_eq!(original_dev.acked_features(), restored_dev.acked_features());
148    }
149}
vmm/devices/virtio/mem/persist.rs

vmm/devices/virtio/mem/
persist.rs
