underhill_mem/
registrar.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Code to register lower VTL memory with the kernel as needed.
5//!
6//! For many kernel operations that operate on memory, such as passing a buffer
7//! to a device for DMA, the kernel requires that has allocated a `struct page`
8//! object for each page being accessed. Thanks to some optimizations for large
9//! memory allocations, the space overhead of this for guest memory is not too
10//! large, but the initialization time overhead can be significant for large
11//! VMs.
12//!
13//! To avoid this overhead, we only register memory with the kernel as needed,
14//! when a VA might leak out of a `GuestMemory` object and possibly be passed to
15//! a kernel routine.
16//!
17//! Memory is registered in caller-selected chunks. We track whether a given
18//! chunk has been registered via a small bitmap.
19
20use cvm_tracing::CVM_ALLOWED;
21use inspect::Inspect;
22use memory_range::MemoryRange;
23use memory_range::overlapping_ranges;
24use parking_lot::Mutex;
25use std::ops::Range;
26use std::sync::atomic::AtomicU64;
27use std::sync::atomic::Ordering::Acquire;
28use std::sync::atomic::Ordering::Release;
29use thiserror::Error;
30use vm_topology::memory::MemoryLayout;
31use vm_topology::memory::MemoryRangeWithNode;
32
33const PAGE_SIZE: u64 = guestmem::PAGE_SIZE as u64;
34
35#[derive(Debug)]
36pub struct MemoryRegistrar<T> {
37    registered: Bitmap,
38    chunk_count: u64,
39    state: Mutex<RegistrarState>,
40    register: T,
41    ram: Vec<MemoryRange>,
42    ram_with_node: Vec<MemoryRangeWithNode>,
43    registration_offset: u64,
44    granularity: u64,
45}
46
47impl<T> Inspect for MemoryRegistrar<T> {
48    fn inspect(&self, req: inspect::Request<'_>) {
49        req.respond()
50            .field_with("chunks_registered", || {
51                (0..self.chunk_count)
52                    .filter(|&chunk| self.registered.get(chunk))
53                    .count()
54            })
55            .field("chunk_count", self.chunk_count)
56            .field("granularity", self.granularity)
57            .hex("registration_offset", self.registration_offset);
58    }
59}
60
61#[derive(Debug)]
62struct RegistrarState {
63    failed: Bitmap,
64}
65
66#[derive(Debug, Error, PartialEq, Eq)]
67pub(crate) enum RegisterAllError {
68    #[error("failed to register memory starting at {address:#x}")]
69    RegistrationFailed { address: u64 },
70    #[error(
71        "VTL0 RAM span {span} has unregistrable edge {edge} in virtual NUMA node {vnode} range \
72         {node_range}; kernel registration requires {alignment:#x}-aligned span boundaries"
73    )]
74    UnalignedMemory {
75        span: MemoryRange,
76        edge: MemoryRange,
77        vnode: u32,
78        node_range: MemoryRange,
79        alignment: u64,
80    },
81}
82
83#[derive(Debug)]
84struct Bitmap(Vec<AtomicU64>);
85
86impl Bitmap {
87    fn new(address_space_size: u64, granularity: u64) -> Self {
88        let chunks = address_space_size.div_ceil(granularity);
89        let words = chunks.div_ceil(64);
90        let mut v = Vec::new();
91        v.resize_with(words as usize, AtomicU64::default);
92        Self(v)
93    }
94
95    fn get(&self, chunk: u64) -> bool {
96        self.0[chunk as usize / 64].load(Acquire) & (1 << (chunk % 64)) != 0
97    }
98
99    fn get_mut(&mut self, chunk: u64) -> bool {
100        *self.0[chunk as usize / 64].get_mut() & (1 << (chunk % 64)) != 0
101    }
102
103    fn set(&self, chunk: u64, value: bool) {
104        if value {
105            self.0[chunk as usize / 64].fetch_or(1 << (chunk % 64), Release);
106        } else {
107            self.0[chunk as usize / 64].fetch_and(!(1 << (chunk % 64)), Release);
108        }
109    }
110
111    fn set_mut(&mut self, chunk: u64, value: bool) {
112        if value {
113            *self.0[chunk as usize / 64].get_mut() |= 1 << (chunk % 64);
114        } else {
115            *self.0[chunk as usize / 64].get_mut() &= !(1 << (chunk % 64));
116        }
117    }
118}
119
120pub trait RegisterMemory {
121    fn register_range(&self, range: MemoryRange) -> Result<(), impl 'static + std::error::Error>;
122}
123
124impl<T: Fn(MemoryRange) -> Result<(), E>, E: 'static + std::error::Error> RegisterMemory for T {
125    fn register_range(&self, range: MemoryRange) -> Result<(), impl 'static + std::error::Error> {
126        (self)(range)
127    }
128}
129
130impl<T: RegisterMemory> MemoryRegistrar<T> {
131    pub fn new(
132        layout: &MemoryLayout,
133        registration_offset: u64,
134        granularity: u64,
135        register: T,
136    ) -> Self {
137        assert!(granularity.is_power_of_two());
138        assert!(granularity >= PAGE_SIZE);
139        let address_space_size = layout.ram().last().unwrap().range.end();
140
141        let mut ram: Vec<MemoryRange> = Vec::new();
142        for range in layout.ram().iter().map(|entry| entry.range) {
143            if let Some(previous) = ram.last_mut()
144                && previous.end() == range.start()
145            {
146                *previous = MemoryRange::new(previous.start()..range.end());
147            } else {
148                ram.push(range);
149            }
150        }
151
152        Self {
153            chunk_count: address_space_size.div_ceil(granularity),
154            registered: Bitmap::new(address_space_size, granularity),
155            state: Mutex::new(RegistrarState {
156                failed: Bitmap::new(address_space_size, granularity),
157            }),
158            register,
159            ram,
160            ram_with_node: layout.ram().to_vec(),
161            registration_offset,
162            granularity,
163        }
164    }
165
166    fn chunks(&self, range: MemoryRange) -> Range<u64> {
167        let start = range.start() / self.granularity;
168        let end = range.end().div_ceil(self.granularity);
169        start..end
170    }
171
172    fn register_range(&self, state: &mut RegistrarState, range: MemoryRange) -> Result<(), u64> {
173        let registered_range = MemoryRange::new(
174            self.registration_offset + range.start()..self.registration_offset + range.end(),
175        );
176        tracing::info!(CVM_ALLOWED, range = %registered_range, "registering memory");
177        if let Err(err) = self.register.register_range(registered_range) {
178            tracing::error!(
179                CVM_ALLOWED,
180                range = %registered_range,
181                registration_offset = self.registration_offset,
182                error = &err as &dyn std::error::Error,
183                "failed to register memory"
184            );
185            for chunk in self.chunks(range) {
186                state.failed.set_mut(chunk, true);
187            }
188            return Err(range.start());
189        }
190        Ok(())
191    }
192
193    pub fn register(&self, address: u64, len: u64) -> Result<(), u64> {
194        // Page align the requested range.
195        let requested_range = MemoryRange::new(
196            address & !(PAGE_SIZE - 1)..(address + len + (PAGE_SIZE - 1)) & !(PAGE_SIZE - 1),
197        );
198
199        // Check if the range is already registered.
200        'check_registered: {
201            for chunk in self.chunks(requested_range) {
202                if !self.registered.get(chunk) {
203                    break 'check_registered;
204                }
205            }
206            return Ok(());
207        }
208
209        // Register each chunk one at a time. We don't typically lock lots of
210        // memory at a time, so in practice there should only be one chunk
211        // anyway.
212        let mut state = self.state.lock();
213        for chunk in self.chunks(requested_range) {
214            if self.registered.get(chunk) {
215                continue;
216            }
217            if state.failed.get_mut(chunk) {
218                return Err(chunk * self.granularity);
219            }
220            // Register the full chunk, bounded by the RAM regions. This could
221            // be more efficient, but again, we expect there to only be one
222            // chunk in practice.
223            let full_range =
224                MemoryRange::new(chunk * self.granularity..(chunk + 1) * self.granularity);
225            for range in overlapping_ranges([full_range], self.ram.iter().copied()) {
226                self.register_range(&mut state, range)?;
227            }
228            self.registered.set(chunk, true);
229        }
230        Ok(())
231    }
232
233    /// Register every complete `alignment`-aligned RAM subrange in the address space.
234    ///
235    /// If `ignore_unaligned_ranges` is false, fail if any RAM remains outside
236    /// the aligned subranges.
237    pub fn register_all_aligned(
238        &self,
239        alignment: u64,
240        ignore_unaligned_ranges: bool,
241    ) -> Result<(), RegisterAllError> {
242        assert!(alignment.is_power_of_two());
243        assert!(alignment >= self.granularity);
244
245        let mut state = self.state.lock();
246        for &span in &self.ram {
247            let aligned_range = span.aligned_subrange(alignment);
248            let unaligned_edge = if ignore_unaligned_ranges {
249                None
250            } else if aligned_range.is_empty() {
251                Some(span)
252            } else if span.start() != aligned_range.start() {
253                Some(MemoryRange::new(span.start()..aligned_range.start()))
254            } else if span.end() != aligned_range.end() {
255                Some(MemoryRange::new(aligned_range.end()..span.end()))
256            } else {
257                None
258            };
259            if let Some(edge) = unaligned_edge {
260                let entry = self
261                    .ram_with_node
262                    .iter()
263                    .find(|entry| entry.range.overlaps(&edge))
264                    .expect("edge belongs to a RAM range");
265                return Err(RegisterAllError::UnalignedMemory {
266                    span,
267                    edge,
268                    vnode: entry.vnode,
269                    node_range: entry.range,
270                    alignment,
271                });
272            }
273
274            if aligned_range.is_empty() {
275                continue;
276            }
277
278            let mut unregistered_run_start = None;
279
280            for chunk in self.chunks(aligned_range) {
281                if state.failed.get_mut(chunk) {
282                    return Err(RegisterAllError::RegistrationFailed {
283                        address: chunk * self.granularity,
284                    });
285                }
286                if self.registered.get(chunk) {
287                    if let Some(start) = unregistered_run_start.take() {
288                        self.register_range(
289                            &mut state,
290                            MemoryRange::new(start..chunk * self.granularity),
291                        )
292                        .map_err(|address| RegisterAllError::RegistrationFailed { address })?;
293                        for registered_chunk in start / self.granularity..chunk {
294                            self.registered.set(registered_chunk, true);
295                        }
296                    }
297                } else {
298                    unregistered_run_start.get_or_insert(chunk * self.granularity);
299                }
300            }
301
302            if let Some(start) = unregistered_run_start {
303                let range = MemoryRange::new(start..aligned_range.end());
304                self.register_range(&mut state, range)
305                    .map_err(|address| RegisterAllError::RegistrationFailed { address })?;
306                for chunk in self.chunks(range) {
307                    self.registered.set(chunk, true);
308                }
309            }
310        }
311        Ok(())
312    }
313}
314
315#[cfg(test)]
316mod tests {
317    use super::MemoryRegistrar;
318    use super::RegisterAllError;
319    use memory_range::MemoryRange;
320    use std::cell::Cell;
321    use std::cell::RefCell;
322    use std::convert::Infallible;
323    use vm_topology::memory::MemoryLayout;
324    use vm_topology::memory::MemoryRangeWithNode;
325
326    const GRANULARITY: u64 = 1 << 30;
327
328    #[test]
329    fn test_registrar() {
330        let layout = MemoryLayout::new(
331            1 << 40,
332            &[
333                MemoryRange::new(0x10000..0x20000),
334                MemoryRange::new(1 << 40..2 << 40),
335            ],
336            &[],
337            &[],
338            None,
339        )
340        .unwrap();
341
342        let offset = 1 << 50;
343        let ranges = RefCell::new(Vec::new());
344        let registrar = MemoryRegistrar::new(&layout, offset, GRANULARITY, |range| {
345            println!("registering {:#x?}", range);
346            ranges.borrow_mut().push(range);
347            Ok::<_, Infallible>(())
348        });
349
350        for range in [
351            0x1000..0x8000,
352            0x20000..0x30000,
353            0x100000..0x200000,
354            1u64 << 33..(1u64 << 35) + 1,
355        ] {
356            registrar
357                .register(range.start, range.end - range.start)
358                .unwrap();
359        }
360
361        let mut expected = vec![
362            MemoryRange::new(offset..offset | 0x10000),
363            MemoryRange::new(offset | 0x20000..offset | GRANULARITY),
364        ];
365        expected.extend(
366            (1 << 33..(1 << 35) + GRANULARITY)
367                .step_by(GRANULARITY as usize)
368                .map(|start| MemoryRange::new(offset | start..offset | (start + GRANULARITY))),
369        );
370
371        let ranges = ranges.take();
372        assert_eq!(
373            ranges.as_slice(),
374            expected.as_slice(),
375            "ranges: {}\n\nexpected: {}",
376            ranges
377                .iter()
378                .map(|r| r.to_string())
379                .collect::<Vec<_>>()
380                .join("\n"),
381            expected
382                .iter()
383                .map(|r| r.to_string())
384                .collect::<Vec<_>>()
385                .join("\n")
386        );
387    }
388
389    #[test]
390    fn test_register_all_aligned_rejects_unaligned_edges_at_minimum_granularity() {
391        let layout = MemoryLayout::new_from_ranges(
392            &[MemoryRangeWithNode {
393                range: MemoryRange::new(0x10000..2 * GRANULARITY + 0x20000),
394                vnode: 7,
395            }],
396            &[],
397        )
398        .unwrap();
399
400        let ranges = RefCell::new(Vec::new());
401        let registrar = MemoryRegistrar::new(&layout, 0, GRANULARITY, |range| {
402            ranges.borrow_mut().push(range);
403            Ok::<_, Infallible>(())
404        });
405
406        assert_eq!(
407            registrar.register_all_aligned(GRANULARITY, false),
408            Err(RegisterAllError::UnalignedMemory {
409                span: MemoryRange::new(0x10000..2 * GRANULARITY + 0x20000),
410                edge: MemoryRange::new(0x10000..GRANULARITY),
411                vnode: 7,
412                node_range: MemoryRange::new(0x10000..2 * GRANULARITY + 0x20000),
413                alignment: GRANULARITY,
414            })
415        );
416
417        assert!(ranges.take().is_empty());
418    }
419
420    #[test]
421    fn test_register_all_aligned_reports_repro_suffix_and_numa_node() {
422        const PMD_GRANULARITY: u64 = 1 << 21;
423        let span = MemoryRange::new(0x80000000..0xd8150000);
424        let layout = MemoryLayout::new_from_ranges(
425            &[
426                MemoryRangeWithNode {
427                    range: MemoryRange::new(0x80000000..0xc0000000),
428                    vnode: 0,
429                },
430                MemoryRangeWithNode {
431                    range: MemoryRange::new(0xc0000000..0xd8150000),
432                    vnode: 1,
433                },
434            ],
435            &[],
436        )
437        .unwrap();
438
439        let registrar =
440            MemoryRegistrar::new(&layout, 0, PMD_GRANULARITY, |_| Ok::<_, Infallible>(()));
441
442        assert_eq!(
443            registrar.register_all_aligned(PMD_GRANULARITY, false),
444            Err(RegisterAllError::UnalignedMemory {
445                span,
446                edge: MemoryRange::new(0xd8000000..0xd8150000),
447                vnode: 1,
448                node_range: MemoryRange::new(0xc0000000..0xd8150000),
449                alignment: PMD_GRANULARITY,
450            })
451        );
452    }
453
454    #[test]
455    fn test_register_all_aligned_rejects_span_smaller_than_granularity() {
456        let span = MemoryRange::new(0x10000..0x20000);
457        let layout = MemoryLayout::new_from_ranges(
458            &[MemoryRangeWithNode {
459                range: span,
460                vnode: 2,
461            }],
462            &[],
463        )
464        .unwrap();
465
466        let registrar = MemoryRegistrar::new(&layout, 0, GRANULARITY, |_| Ok::<_, Infallible>(()));
467
468        assert_eq!(
469            registrar.register_all_aligned(GRANULARITY, false),
470            Err(RegisterAllError::UnalignedMemory {
471                span,
472                edge: span,
473                vnode: 2,
474                node_range: span,
475                alignment: GRANULARITY,
476            })
477        );
478    }
479
480    #[test]
481    fn test_register_all_aligned_uses_smaller_pages_for_edges() {
482        const SMALL_GRANULARITY: u64 = 1 << 21;
483
484        let layout = MemoryLayout::new_from_ranges(
485            &[MemoryRangeWithNode {
486                range: MemoryRange::new(SMALL_GRANULARITY..2 * GRANULARITY + SMALL_GRANULARITY),
487                vnode: 0,
488            }],
489            &[],
490        )
491        .unwrap();
492
493        let ranges = RefCell::new(Vec::new());
494        let registrar = MemoryRegistrar::new(&layout, 0, SMALL_GRANULARITY, |range| {
495            ranges.borrow_mut().push(range);
496            Ok::<_, Infallible>(())
497        });
498
499        registrar.register_all_aligned(GRANULARITY, true).unwrap();
500        registrar
501            .register_all_aligned(SMALL_GRANULARITY, false)
502            .unwrap();
503
504        assert_eq!(registrar.register(SMALL_GRANULARITY, 0x1000), Ok(()));
505        assert_eq!(registrar.register(2 * GRANULARITY, 0x1000), Ok(()));
506
507        assert_eq!(
508            ranges.take(),
509            [
510                MemoryRange::new(GRANULARITY..2 * GRANULARITY),
511                MemoryRange::new(SMALL_GRANULARITY..GRANULARITY),
512                MemoryRange::new(2 * GRANULARITY..2 * GRANULARITY + SMALL_GRANULARITY),
513            ]
514        );
515    }
516
517    #[test]
518    fn test_register_all_aligned_multiple_ranges_with_gap() {
519        const SMALL_GRANULARITY: u64 = 1 << 16;
520
521        // A mix of RAM ranges with unbacked gaps between them: one spanning a
522        // chunk boundary, one contained in a single chunk, one spanning
523        // several whole chunks, and one crossing a boundary with unaligned ends.
524        let layout = MemoryLayout::new_from_ranges(
525            &[
526                MemoryRangeWithNode {
527                    range: MemoryRange::new(0x10000..GRANULARITY + 0x20000),
528                    vnode: 0,
529                },
530                MemoryRangeWithNode {
531                    range: MemoryRange::new(3 * GRANULARITY + 0x10000..3 * GRANULARITY + 0x30000),
532                    vnode: 0,
533                },
534                MemoryRangeWithNode {
535                    range: MemoryRange::new(4 * GRANULARITY + 0x40000..4 * GRANULARITY + 0x50000),
536                    vnode: 0,
537                },
538                MemoryRangeWithNode {
539                    range: MemoryRange::new(5 * GRANULARITY..8 * GRANULARITY),
540                    vnode: 0,
541                },
542                MemoryRangeWithNode {
543                    range: MemoryRange::new(9 * GRANULARITY + 0x30000..10 * GRANULARITY + 0x10000),
544                    vnode: 0,
545                },
546            ],
547            &[],
548        )
549        .unwrap();
550
551        let ranges = RefCell::new(Vec::new());
552        let registrar = MemoryRegistrar::new(&layout, 0, SMALL_GRANULARITY, |range| {
553            ranges.borrow_mut().push(range);
554            Ok::<_, Infallible>(())
555        });
556
557        registrar.register_all_aligned(GRANULARITY, true).unwrap();
558
559        assert_eq!(
560            ranges.take(),
561            [MemoryRange::new(5 * GRANULARITY..8 * GRANULARITY)]
562        );
563    }
564
565    #[test]
566    fn test_register_all_aligned_coalesces_large_aligned_range() {
567        const SMALL_GRANULARITY: u64 = 1 << 21;
568
569        let layout = MemoryLayout::new_from_ranges(
570            &[MemoryRangeWithNode {
571                range: MemoryRange::new(0..3 * GRANULARITY),
572                vnode: 0,
573            }],
574            &[],
575        )
576        .unwrap();
577
578        let ranges = RefCell::new(Vec::new());
579        let registrar = MemoryRegistrar::new(&layout, 0, SMALL_GRANULARITY, |range| {
580            ranges.borrow_mut().push(range);
581            Ok::<_, Infallible>(())
582        });
583
584        registrar.register_all_aligned(GRANULARITY, false).unwrap();
585
586        assert_eq!(ranges.take(), [MemoryRange::new(0..3 * GRANULARITY)]);
587    }
588
589    #[test]
590    fn test_register_all_aligned_merges_adjacent_numa_ranges() {
591        let layout = MemoryLayout::new_from_ranges(
592            &[
593                MemoryRangeWithNode {
594                    range: MemoryRange::new(0..GRANULARITY / 2),
595                    vnode: 0,
596                },
597                MemoryRangeWithNode {
598                    range: MemoryRange::new(GRANULARITY / 2..GRANULARITY),
599                    vnode: 1,
600                },
601            ],
602            &[],
603        )
604        .unwrap();
605
606        let ranges = RefCell::new(Vec::new());
607        let registrar = MemoryRegistrar::new(&layout, 0, 1 << 21, |range| {
608            ranges.borrow_mut().push(range);
609            Ok::<_, Infallible>(())
610        });
611
612        registrar.register_all_aligned(GRANULARITY, false).unwrap();
613
614        assert_eq!(ranges.take(), [MemoryRange::new(0..GRANULARITY)]);
615    }
616
617    #[test]
618    fn test_failed_subrange_does_not_mark_chunk_registered() {
619        let layout = MemoryLayout::new_from_ranges(
620            &[
621                MemoryRangeWithNode {
622                    range: MemoryRange::new(0..0x10000),
623                    vnode: 0,
624                },
625                MemoryRangeWithNode {
626                    range: MemoryRange::new(0x20000..0x30000),
627                    vnode: 0,
628                },
629            ],
630            &[],
631        )
632        .unwrap();
633
634        let calls = Cell::new(0);
635        let registrar = MemoryRegistrar::new(&layout, 0, GRANULARITY, |_| {
636            let call = calls.get();
637            calls.set(call + 1);
638            if call == 0 {
639                Ok(())
640            } else {
641                Err(std::io::Error::other("registration failure"))
642            }
643        });
644
645        assert!(registrar.register(0, 1).is_err());
646        assert!(registrar.register(0, 1).is_err());
647        assert_eq!(calls.get(), 2);
648    }
649}
underhill_mem/registrar.rs

underhill_mem/
registrar.rs
