wasmer_wasix/state/linker/
memory_allocator.rs

1use std::cmp::{Ord, PartialEq, PartialOrd};
2
3use tracing::trace;
4use wasmer::{AsStoreMut, Memory, MemoryError};
5
6const WASM_PAGE_SIZE: u32 = wasmer::WASM_PAGE_SIZE as u32;
7
8struct AllocatedPage {
9    // The base_ptr is mutable, and will move forward as memory is allocated from the page.
10    base_ptr: u32,
11
12    // The amount of memory remaining until the end of the allocated region. Despite the
13    // name of this struct, the region does not have to be only one page.
14    remaining: u32,
15}
16
17// Used to allocate and manage memory for dynamic modules that are loaded in or
18// out, since each module may request a specific amount of memory to be allocated
19// for it before starting it up.
20// TODO: Only supports Memory32, should implement proper Memory64 support
21pub(super) struct MemoryAllocator {
22    allocated_pages: Vec<AllocatedPage>,
23}
24
25impl MemoryAllocator {
26    pub fn new() -> Self {
27        Self {
28            allocated_pages: vec![],
29        }
30    }
31
32    pub fn allocate(
33        &mut self,
34        memory: &Memory,
35        store: &mut impl AsStoreMut,
36        size: u32,
37        alignment: u32,
38    ) -> Result<u32, MemoryError> {
39        match self.allocate_in_existing_pages(size, alignment) {
40            Some(base_ptr) => Ok(base_ptr),
41            None => self.allocate_new_page(memory, store, size),
42        }
43    }
44
45    // Finds a page which has enough free memory for the request, and allocates in it.
46    // Returns the address of the allocated region if one was found.
47    fn allocate_in_existing_pages(&mut self, size: u32, alignment: u32) -> Option<u32> {
48        // A type to hold intermediate search results. The idea is to allocate on the page
49        // that has the least amount of free space, so we can later satisfy larger allocation
50        // requests without having to allocate entire new pages.
51        struct CandidatePage {
52            index: usize,
53            base_ptr: u32,
54            to_add: u32,
55            remaining_free: u32,
56        }
57
58        impl PartialEq for CandidatePage {
59            fn eq(&self, other: &Self) -> bool {
60                self.remaining_free == other.remaining_free
61            }
62        }
63
64        impl Eq for CandidatePage {}
65
66        impl PartialOrd for CandidatePage {
67            fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
68                Some(self.cmp(other))
69            }
70        }
71
72        impl Ord for CandidatePage {
73            fn cmp(&self, other: &Self) -> std::cmp::Ordering {
74                self.remaining_free.cmp(&other.remaining_free)
75            }
76        }
77
78        let mut candidates = std::collections::BinaryHeap::new();
79
80        for (index, page) in self.allocated_pages.iter().enumerate() {
81            // Offset for proper alignment
82            let offset = if page.base_ptr % alignment == 0 {
83                0
84            } else {
85                alignment - (page.base_ptr % alignment)
86            };
87
88            if page.remaining >= offset + size {
89                candidates.push(std::cmp::Reverse(CandidatePage {
90                    index,
91                    base_ptr: page.base_ptr + offset,
92                    to_add: offset + size,
93                    remaining_free: page.remaining - offset - size,
94                }));
95            }
96        }
97
98        candidates.pop().map(|elected| {
99            let page = &mut self.allocated_pages[elected.0.index];
100
101            trace!(
102                free = page.remaining,
103                base_ptr = elected.0.base_ptr,
104                "Found existing memory page with sufficient space"
105            );
106
107            page.base_ptr += elected.0.to_add;
108            page.remaining -= elected.0.to_add;
109            elected.0.base_ptr
110        })
111    }
112
113    fn allocate_new_page(
114        &mut self,
115        memory: &Memory,
116        store: &mut impl AsStoreMut,
117        size: u32,
118    ) -> Result<u32, MemoryError> {
119        // No need to account for alignment here, as pages are already 64k-aligned
120        let to_grow = size.div_ceil(WASM_PAGE_SIZE);
121        let pages = memory.grow(store, to_grow)?;
122
123        let base_ptr = pages.0 * WASM_PAGE_SIZE;
124        let total_allocated = to_grow * WASM_PAGE_SIZE;
125
126        // The initial size bytes are already allocated, rest goes into the list
127        if total_allocated > size {
128            self.allocated_pages.push(AllocatedPage {
129                base_ptr: base_ptr + size,
130                remaining: total_allocated - size,
131            });
132        }
133
134        trace!(
135            page_count = to_grow,
136            size, base_ptr, "Allocated new memory page(s) to accommodate requested memory"
137        );
138
139        Ok(base_ptr)
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::{MemoryAllocator, WASM_PAGE_SIZE};
146    use wasmer::{Engine, Memory, Store};
147
148    #[test]
149    fn test_memory_allocator() {
150        let engine = Engine::default();
151        let mut store = Store::new(engine);
152        let memory = Memory::new(
153            &mut store,
154            wasmer::MemoryType {
155                minimum: wasmer::Pages(2),
156                maximum: None,
157                shared: true,
158            },
159        )
160        .unwrap();
161        let mut allocator = MemoryAllocator::new();
162
163        // Small allocation in new page
164        let addr = allocator.allocate(&memory, &mut store, 24, 4).unwrap();
165        assert_eq!(addr, 2 * WASM_PAGE_SIZE);
166        assert_eq!(memory.grow(&mut store, 0).unwrap().0, 3);
167
168        // Small allocation in existing page
169        let addr = allocator.allocate(&memory, &mut store, 16, 4).unwrap();
170        assert_eq!(addr, 2 * WASM_PAGE_SIZE + 24);
171
172        // Small allocation in existing page, with bigger alignment
173        let addr = allocator.allocate(&memory, &mut store, 64, 32).unwrap();
174        assert_eq!(addr, 2 * WASM_PAGE_SIZE + 64);
175        // Should still have 3 pages
176        assert_eq!(memory.grow(&mut store, 0).unwrap().0, 3);
177
178        // Big allocation in new pages
179        let addr = allocator
180            .allocate(&memory, &mut store, 2 * WASM_PAGE_SIZE + 256, 1024)
181            .unwrap();
182        assert_eq!(addr, WASM_PAGE_SIZE * 3);
183        assert_eq!(memory.grow(&mut store, 0).unwrap().0, 6);
184
185        // Small allocation with multiple empty pages
186        // page 2 has 128 bytes allocated, page 5 has 256, allocation should go
187        // to page 5 (we should allocate from the page with the least free space)
188        let addr = allocator
189            .allocate(&memory, &mut store, 1024 * 63, 64)
190            .unwrap();
191        assert_eq!(addr, 5 * WASM_PAGE_SIZE + 256);
192
193        // Another small allocation, but this time it won't fit on page 5
194        let addr = allocator.allocate(&memory, &mut store, 4096, 512).unwrap();
195        assert_eq!(addr, 2 * WASM_PAGE_SIZE + 512);
196    }
197}
wasmer_wasix/state/linker/memory_allocator.rs

wasmer_wasix/state/linker/
memory_allocator.rs
