vmm/vstate/
resources.rs

1// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::convert::Infallible;
5
6use serde::{Deserialize, Serialize};
7pub use vm_allocator::AllocPolicy;
8use vm_allocator::{AddressAllocator, IdAllocator};
9
10use crate::arch;
11use crate::snapshot::Persist;
12
13/// Helper function to allocate many ids from an id allocator
14fn allocate_many_ids(
15    id_allocator: &mut IdAllocator,
16    count: u32,
17) -> Result<Vec<u32>, vm_allocator::Error> {
18    let mut ids = Vec::with_capacity(count as usize);
19
20    for _ in 0..count {
21        match id_allocator.allocate_id() {
22            Ok(id) => ids.push(id),
23            Err(err) => {
24                // It is ok to unwrap here, we just allocated the GSI
25                ids.into_iter().for_each(|id| {
26                    id_allocator.free_id(id).unwrap();
27                });
28                return Err(err);
29            }
30        }
31    }
32
33    Ok(ids)
34}
35
36/// A resource manager for (de)allocating interrupt lines (GSIs) and guest memory
37///
38/// At the moment, we support:
39///
40/// * GSIs for legacy x86_64 devices
41/// * GSIs for MMIO devicecs
42/// * Memory allocations in the MMIO address space
43#[derive(Debug, Clone, Serialize, Deserialize)]
44pub struct ResourceAllocator {
45    /// Allocator for legacy device interrupt lines
46    pub gsi_legacy_allocator: IdAllocator,
47    /// Allocator for PCI device GSIs
48    pub gsi_msi_allocator: IdAllocator,
49    /// Allocator for memory in the 32-bit MMIO address space
50    pub mmio32_memory: AddressAllocator,
51    /// Allocator for memory in the 64-bit MMIO address space
52    pub mmio64_memory: AddressAllocator,
53    /// Allocator for memory after the 64-bit MMIO address space
54    pub past_mmio64_memory: AddressAllocator,
55    /// Memory allocator for system data
56    pub system_memory: AddressAllocator,
57}
58
59impl Default for ResourceAllocator {
60    fn default() -> Self {
61        ResourceAllocator::new()
62    }
63}
64
65impl ResourceAllocator {
66    /// Create a new resource allocator for Firecracker devices
67    pub fn new() -> Self {
68        // It is fine for us to unwrap the following since we know we are passing valid ranges for
69        // all allocators
70        Self {
71            gsi_legacy_allocator: IdAllocator::new(arch::GSI_LEGACY_START, arch::GSI_LEGACY_END)
72                .unwrap(),
73            gsi_msi_allocator: IdAllocator::new(arch::GSI_MSI_START, arch::GSI_MSI_END).unwrap(),
74            mmio32_memory: AddressAllocator::new(
75                arch::MEM_32BIT_DEVICES_START,
76                arch::MEM_32BIT_DEVICES_SIZE,
77            )
78            .unwrap(),
79            mmio64_memory: AddressAllocator::new(
80                arch::MEM_64BIT_DEVICES_START,
81                arch::MEM_64BIT_DEVICES_SIZE,
82            )
83            .unwrap(),
84            past_mmio64_memory: AddressAllocator::new(
85                arch::FIRST_ADDR_PAST_64BITS_MMIO,
86                arch::PAST_64BITS_MMIO_SIZE,
87            )
88            .unwrap(),
89            system_memory: AddressAllocator::new(arch::SYSTEM_MEM_START, arch::SYSTEM_MEM_SIZE)
90                .unwrap(),
91        }
92    }
93
94    /// Allocate a number of legacy GSIs
95    ///
96    /// # Arguments
97    ///
98    /// * `gsi_count` - The number of legacy GSIs to allocate
99    pub fn allocate_gsi_legacy(&mut self, gsi_count: u32) -> Result<Vec<u32>, vm_allocator::Error> {
100        allocate_many_ids(&mut self.gsi_legacy_allocator, gsi_count)
101    }
102
103    /// Allocate a number of GSIs for MSI
104    ///
105    /// # Arguments
106    ///
107    /// * `gsi_count` - The number of GSIs to allocate
108    pub fn allocate_gsi_msi(&mut self, gsi_count: u32) -> Result<Vec<u32>, vm_allocator::Error> {
109        allocate_many_ids(&mut self.gsi_msi_allocator, gsi_count)
110    }
111
112    /// Allocate a memory range in 32-bit MMIO address space
113    ///
114    /// If it succeeds, it returns the first address of the allocated range
115    ///
116    /// # Arguments
117    ///
118    /// * `size` - The size in bytes of the memory to allocate
119    /// * `alignment` - The alignment of the address of the first byte
120    /// * `policy` - A [`vm_allocator::AllocPolicy`] variant for determining the allocation policy
121    pub fn allocate_32bit_mmio_memory(
122        &mut self,
123        size: u64,
124        alignment: u64,
125        policy: AllocPolicy,
126    ) -> Result<u64, vm_allocator::Error> {
127        Ok(self
128            .mmio32_memory
129            .allocate(size, alignment, policy)?
130            .start())
131    }
132
133    /// Allocate a memory range in 64-bit MMIO address space
134    ///
135    /// If it succeeds, it returns the first address of the allocated range
136    ///
137    /// # Arguments
138    ///
139    /// * `size` - The size in bytes of the memory to allocate
140    /// * `alignment` - The alignment of the address of the first byte
141    /// * `policy` - A [`vm_allocator::AllocPolicy`] variant for determining the allocation policy
142    pub fn allocate_64bit_mmio_memory(
143        &mut self,
144        size: u64,
145        alignment: u64,
146        policy: AllocPolicy,
147    ) -> Result<u64, vm_allocator::Error> {
148        Ok(self
149            .mmio64_memory
150            .allocate(size, alignment, policy)?
151            .start())
152    }
153
154    /// Allocate a memory range for system data
155    ///
156    /// If it succeeds, it returns the first address of the allocated range
157    ///
158    /// # Arguments
159    ///
160    /// * `size` - The size in bytes of the memory to allocate
161    /// * `alignment` - The alignment of the address of the first byte
162    /// * `policy` - A [`vm_allocator::AllocPolicy`] variant for determining the allocation policy
163    pub fn allocate_system_memory(
164        &mut self,
165        size: u64,
166        alignment: u64,
167        policy: AllocPolicy,
168    ) -> Result<u64, vm_allocator::Error> {
169        Ok(self
170            .system_memory
171            .allocate(size, alignment, policy)?
172            .start())
173    }
174}
175
176impl<'a> Persist<'a> for ResourceAllocator {
177    type State = ResourceAllocator;
178    type ConstructorArgs = ();
179    type Error = Infallible;
180
181    fn save(&self) -> Self::State {
182        self.clone()
183    }
184
185    fn restore(
186        _constructor_args: Self::ConstructorArgs,
187        state: &Self::State,
188    ) -> Result<Self, Self::Error> {
189        Ok(state.clone())
190    }
191}
192
193#[cfg(test)]
194mod tests {
195    use vm_allocator::AllocPolicy;
196
197    use super::ResourceAllocator;
198    use crate::arch::{self, GSI_LEGACY_NUM, GSI_LEGACY_START, GSI_MSI_NUM, GSI_MSI_START};
199    use crate::snapshot::{Persist, Snapshot};
200
201    #[test]
202    fn test_allocate_irq() {
203        let mut allocator = ResourceAllocator::new();
204        // asking for 0 IRQs should return us an empty vector
205        assert_eq!(allocator.allocate_gsi_legacy(0), Ok(vec![]));
206        // We cannot allocate more GSIs than available
207        assert_eq!(
208            allocator.allocate_gsi_legacy(GSI_LEGACY_NUM + 1),
209            Err(vm_allocator::Error::ResourceNotAvailable)
210        );
211        // But allocating all of them at once should work
212        assert_eq!(
213            allocator.allocate_gsi_legacy(GSI_LEGACY_NUM),
214            Ok((arch::GSI_LEGACY_START..=arch::GSI_LEGACY_END).collect::<Vec<_>>())
215        );
216        // And now we ran out of GSIs
217        assert_eq!(
218            allocator.allocate_gsi_legacy(1),
219            Err(vm_allocator::Error::ResourceNotAvailable)
220        );
221        // But we should be able to ask for 0 GSIs
222        assert_eq!(allocator.allocate_gsi_legacy(0), Ok(vec![]));
223
224        let mut allocator = ResourceAllocator::new();
225        // We should be able to allocate 1 GSI
226        assert_eq!(
227            allocator.allocate_gsi_legacy(1),
228            Ok(vec![arch::GSI_LEGACY_START])
229        );
230        // We can't allocate MAX_IRQS any more
231        assert_eq!(
232            allocator.allocate_gsi_legacy(GSI_LEGACY_NUM),
233            Err(vm_allocator::Error::ResourceNotAvailable)
234        );
235        // We can allocate another one and it should be the second available
236        assert_eq!(
237            allocator.allocate_gsi_legacy(1),
238            Ok(vec![arch::GSI_LEGACY_START + 1])
239        );
240        // Let's allocate the rest in a loop
241        for i in arch::GSI_LEGACY_START + 2..=arch::GSI_LEGACY_END {
242            assert_eq!(allocator.allocate_gsi_legacy(1), Ok(vec![i]));
243        }
244    }
245
246    #[test]
247    fn test_allocate_gsi() {
248        let mut allocator = ResourceAllocator::new();
249        // asking for 0 IRQs should return us an empty vector
250        assert_eq!(allocator.allocate_gsi_msi(0), Ok(vec![]));
251        // We cannot allocate more GSIs than available
252        assert_eq!(
253            allocator.allocate_gsi_msi(GSI_MSI_NUM + 1),
254            Err(vm_allocator::Error::ResourceNotAvailable)
255        );
256        // But allocating all of them at once should work
257        assert_eq!(
258            allocator.allocate_gsi_msi(GSI_MSI_NUM),
259            Ok((arch::GSI_MSI_START..=arch::GSI_MSI_END).collect::<Vec<_>>())
260        );
261        // And now we ran out of GSIs
262        assert_eq!(
263            allocator.allocate_gsi_msi(1),
264            Err(vm_allocator::Error::ResourceNotAvailable)
265        );
266        // But we should be able to ask for 0 GSIs
267        assert_eq!(allocator.allocate_gsi_msi(0), Ok(vec![]));
268
269        let mut allocator = ResourceAllocator::new();
270        // We should be able to allocate 1 GSI
271        assert_eq!(allocator.allocate_gsi_msi(1), Ok(vec![arch::GSI_MSI_START]));
272        // We can't allocate MAX_IRQS any more
273        assert_eq!(
274            allocator.allocate_gsi_msi(GSI_MSI_NUM),
275            Err(vm_allocator::Error::ResourceNotAvailable)
276        );
277        // We can allocate another one and it should be the second available
278        assert_eq!(
279            allocator.allocate_gsi_msi(1),
280            Ok(vec![arch::GSI_MSI_START + 1])
281        );
282        // Let's allocate the rest in a loop
283        for i in arch::GSI_MSI_START + 2..=arch::GSI_MSI_END {
284            assert_eq!(allocator.allocate_gsi_msi(1), Ok(vec![i]));
285        }
286    }
287
288    fn clone_allocator(allocator: &ResourceAllocator) -> ResourceAllocator {
289        let mut buf = vec![0u8; 1024];
290        Snapshot::new(allocator.save())
291            .save(&mut buf.as_mut_slice())
292            .unwrap();
293        let restored_state: ResourceAllocator = Snapshot::load_without_crc_check(buf.as_slice())
294            .unwrap()
295            .data;
296        ResourceAllocator::restore((), &restored_state).unwrap()
297    }
298
299    #[test]
300    fn test_save_restore() {
301        let mut allocator0 = ResourceAllocator::new();
302        let irq_0 = allocator0.allocate_gsi_legacy(1).unwrap()[0];
303        assert_eq!(irq_0, GSI_LEGACY_START);
304        let gsi_0 = allocator0.allocate_gsi_msi(1).unwrap()[0];
305        assert_eq!(gsi_0, GSI_MSI_START);
306
307        let mut allocator1 = clone_allocator(&allocator0);
308        let irq_1 = allocator1.allocate_gsi_legacy(1).unwrap()[0];
309        assert_eq!(irq_1, GSI_LEGACY_START + 1);
310        let gsi_1 = allocator1.allocate_gsi_msi(1).unwrap()[0];
311        assert_eq!(gsi_1, GSI_MSI_START + 1);
312        let mmio32_mem = allocator1
313            .allocate_32bit_mmio_memory(0x42, 1, AllocPolicy::FirstMatch)
314            .unwrap();
315        assert_eq!(mmio32_mem, arch::MEM_32BIT_DEVICES_START);
316        let mmio64_mem = allocator1
317            .allocate_64bit_mmio_memory(0x42, 1, AllocPolicy::FirstMatch)
318            .unwrap();
319        assert_eq!(mmio64_mem, arch::MEM_64BIT_DEVICES_START);
320        let system_mem = allocator1
321            .allocate_system_memory(0x42, 1, AllocPolicy::FirstMatch)
322            .unwrap();
323        assert_eq!(system_mem, arch::SYSTEM_MEM_START);
324
325        let mut allocator2 = clone_allocator(&allocator1);
326        allocator2
327            .allocate_32bit_mmio_memory(0x42, 1, AllocPolicy::ExactMatch(mmio32_mem))
328            .unwrap_err();
329        allocator2
330            .allocate_64bit_mmio_memory(0x42, 1, AllocPolicy::ExactMatch(mmio64_mem))
331            .unwrap_err();
332        allocator2
333            .allocate_system_memory(0x42, 1, AllocPolicy::ExactMatch(system_mem))
334            .unwrap_err();
335
336        let irq_2 = allocator2.allocate_gsi_legacy(1).unwrap()[0];
337        assert_eq!(irq_2, GSI_LEGACY_START + 2);
338        let gsi_2 = allocator2.allocate_gsi_msi(1).unwrap()[0];
339        assert_eq!(gsi_2, GSI_MSI_START + 2);
340        let mmio32_mem = allocator1
341            .allocate_32bit_mmio_memory(0x42, 1, AllocPolicy::FirstMatch)
342            .unwrap();
343        assert_eq!(mmio32_mem, arch::MEM_32BIT_DEVICES_START + 0x42);
344        let mmio64_mem = allocator1
345            .allocate_64bit_mmio_memory(0x42, 1, AllocPolicy::FirstMatch)
346            .unwrap();
347        assert_eq!(mmio64_mem, arch::MEM_64BIT_DEVICES_START + 0x42);
348        let system_mem = allocator1
349            .allocate_system_memory(0x42, 1, AllocPolicy::FirstMatch)
350            .unwrap();
351        assert_eq!(system_mem, arch::SYSTEM_MEM_START + 0x42);
352    }
353}