user_driver/
page_allocator.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Allocator for pages within a pool.
5//!
6//! This is used for temporary allocations of per-queue DMA buffers, mainly for
7//! PRP lists.
8
9use crate::memory::MemoryBlock;
10use crate::memory::PAGE_SIZE;
11use crate::memory::PAGE_SIZE64;
12use guestmem::GuestMemory;
13use guestmem::GuestMemoryError;
14use guestmem::ranges::PagedRange;
15use inspect::Inspect;
16use parking_lot::Mutex;
17use std::sync::atomic::AtomicU8;
18
19#[derive(Inspect)]
20pub struct PageAllocator {
21    #[inspect(flatten)]
22    core: Mutex<PageAllocatorCore>,
23    #[inspect(skip)]
24    mem: MemoryBlock,
25    #[inspect(skip)]
26    event: event_listener::Event,
27    max: usize,
28}
29
30/// An error allocating pages from the page allocator.
31#[derive(Debug)]
32pub struct PageAllocationError {
33    pub requested: usize,
34    pub max: usize,
35}
36
37impl std::fmt::Debug for PageAllocator {
38    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
39        f.debug_struct("PageAllocator").finish()
40    }
41}
42
43impl PageAllocator {
44    pub fn new(mem: MemoryBlock) -> Self {
45        assert_eq!(mem.offset_in_page(), 0);
46        assert_eq!(mem.len() % PAGE_SIZE, 0);
47        let count = mem.len() / PAGE_SIZE;
48        Self {
49            core: Mutex::new(PageAllocatorCore::new(count)),
50            mem,
51            event: Default::default(),
52            max: count,
53        }
54    }
55
56    pub async fn alloc_pages(&self, n: usize) -> Result<ScopedPages<'_>, PageAllocationError> {
57        // A single page must be left over for the PRP list, so one request may
58        // not use all pages.
59        if self.max < n + 1 {
60            return Err(PageAllocationError {
61                requested: n,
62                max: self.max - 1,
63            });
64        }
65        let mut core = loop {
66            let listener = {
67                let core = self.core.lock();
68                if core.remaining() >= n {
69                    break core;
70                }
71                // Fairness is pretty bad with this approach--small allocations
72                // could easily prevent a large allocation from ever succeeding.
73                // But we don't really have this use case right now, so this is OK.
74                self.event.listen()
75            };
76            listener.await;
77        };
78
79        let pfns = self.mem.pfns();
80        let pages = (0..n)
81            .map(|_| {
82                let n = core.alloc().unwrap();
83                ScopedPage {
84                    page_index: n,
85                    physical_address: pfns[n] * PAGE_SIZE64,
86                }
87            })
88            .collect();
89        Ok(ScopedPages { alloc: self, pages })
90    }
91
92    pub async fn alloc_bytes(&self, n: usize) -> Result<ScopedPages<'_>, PageAllocationError> {
93        self.alloc_pages(n.div_ceil(PAGE_SIZE)).await
94    }
95}
96
97#[derive(Inspect)]
98struct PageAllocatorCore {
99    #[inspect(with = "|x| x.len()")]
100    free: Vec<usize>,
101}
102
103impl PageAllocatorCore {
104    fn new(count: usize) -> Self {
105        let free = (0..count).rev().collect();
106        Self { free }
107    }
108
109    fn remaining(&self) -> usize {
110        self.free.len()
111    }
112
113    fn alloc(&mut self) -> Option<usize> {
114        self.free.pop()
115    }
116
117    fn free(&mut self, n: usize) {
118        self.free.push(n);
119    }
120}
121
122pub struct ScopedPages<'a> {
123    alloc: &'a PageAllocator,
124    pages: Vec<ScopedPage>,
125}
126
127struct ScopedPage {
128    page_index: usize,
129    physical_address: u64,
130}
131
132impl ScopedPages<'_> {
133    pub fn page_count(&self) -> usize {
134        self.pages.len()
135    }
136
137    pub fn physical_address(&self, index: usize) -> u64 {
138        self.pages[index].physical_address
139    }
140
141    pub fn page_as_slice(&self, index: usize) -> &[AtomicU8] {
142        &self.alloc.mem.as_slice()[self.pages[index].page_index * PAGE_SIZE..][..PAGE_SIZE]
143    }
144
145    pub fn read(&self, data: &mut [u8]) {
146        assert!(data.len() <= self.pages.len() * PAGE_SIZE);
147        for (chunk, page) in data.chunks_mut(PAGE_SIZE).zip(&self.pages) {
148            self.alloc.mem.read_at(page.page_index * PAGE_SIZE, chunk);
149        }
150    }
151
152    pub fn copy_to_guest_memory(
153        &self,
154        guest_memory: &GuestMemory,
155        mem: PagedRange<'_>,
156    ) -> Result<(), GuestMemoryError> {
157        let mut remaining = mem.len();
158        for (i, page) in self.pages.iter().enumerate() {
159            let len = PAGE_SIZE.min(remaining);
160            remaining -= len;
161            guest_memory.write_range_from_atomic(
162                &mem.subrange(i * PAGE_SIZE, len),
163                &self.alloc.mem.as_slice()[page.page_index * PAGE_SIZE..][..len],
164            )?;
165        }
166        Ok(())
167    }
168
169    pub fn write(&self, data: &[u8]) {
170        assert!(data.len() <= self.pages.len() * PAGE_SIZE);
171        for (chunk, page) in data.chunks(PAGE_SIZE).zip(&self.pages) {
172            self.alloc.mem.write_at(page.page_index * PAGE_SIZE, chunk);
173        }
174    }
175
176    pub fn copy_from_guest_memory(
177        &self,
178        guest_memory: &GuestMemory,
179        mem: PagedRange<'_>,
180    ) -> Result<(), GuestMemoryError> {
181        let mut remaining = mem.len();
182        for (i, page) in self.pages.iter().enumerate() {
183            let len = PAGE_SIZE.min(remaining);
184            remaining -= len;
185            guest_memory.read_range_to_atomic(
186                &mem.subrange(i * PAGE_SIZE, len),
187                &self.alloc.mem.as_slice()[page.page_index * PAGE_SIZE..][..len],
188            )?;
189        }
190        Ok(())
191    }
192}
193
194impl Drop for ScopedPages<'_> {
195    fn drop(&mut self) {
196        let n = self.pages.len();
197        {
198            let mut core = self.alloc.core.lock();
199            for page in self.pages.drain(..) {
200                core.free(page.page_index);
201            }
202        }
203        self.alloc.event.notify_additional(n);
204    }
205}
user_driver/page_allocator.rs

user_driver/
page_allocator.rs
