rust-kernel / src / allocator / fixed_size_block.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
//! Fixed-size block allocator (the ACTIVE allocator).
//!
//! Every request is rounded up to one of a few block sizes (8, 16, ... 2048
//! bytes). For each size we keep a linked list of free blocks; the "next"
//! pointer is stored *inside* the free block itself, so the lists need no
//! extra memory. alloc = pop from the list, dealloc = push onto it: O(1).
//! When a list is empty, or a request is bigger than 2048 bytes, we ask the
//! fallback `linked_list_allocator::Heap`.
//! Trade-off: up to ~50% wasted space per block (e.g. 33 bytes -> 64).

use super::Locked;
use alloc::alloc::{GlobalAlloc, Layout};
use core::{
    mem,
    ptr::{self, NonNull},
};

/// Block sizes. Each is a power of two, so it can double as the alignment.
const BLOCK_SIZES: &[usize] = &[8, 16, 32, 64, 128, 256, 512, 1024, 2048];

/// Index into BLOCK_SIZES of the smallest block that fits `layout`.
fn list_index(layout: &Layout) -> Option<usize> {
    let required_block_size = layout.size().max(layout.align());
    BLOCK_SIZES.iter().position(|&s| s >= required_block_size)
}

/// A free block. Lives inside the free memory it describes.
struct ListNode {
    next: Option<&'static mut ListNode>,
}

pub struct FixedSizeBlockAllocator {
    list_heads: [Option<&'static mut ListNode>; BLOCK_SIZES.len()],
    fallback_allocator: linked_list_allocator::Heap,
}

impl FixedSizeBlockAllocator {
    pub const fn new() -> Self {
        const EMPTY: Option<&'static mut ListNode> = None;
        FixedSizeBlockAllocator {
            list_heads: [EMPTY; BLOCK_SIZES.len()],
            fallback_allocator: linked_list_allocator::Heap::empty(),
        }
    }

    /// Unsafe: the range must be mapped and unused; call only once.
    pub unsafe fn init(&mut self, heap_start: usize, heap_size: usize) {
        unsafe { self.fallback_allocator.init(heap_start, heap_size) };
    }

    fn fallback_alloc(&mut self, layout: Layout) -> *mut u8 {
        match self.fallback_allocator.allocate_first_fit(layout) {
            Ok(ptr) => ptr.as_ptr(),
            Err(_) => ptr::null_mut(),
        }
    }
}

// The heap lock is taken with interrupts off: code that runs with
// interrupts off (the scheduler, wait queues, ...) may allocate, and if it
// found the lock held by a thread that was preempted mid-allocation, it
// would spin forever -- that thread can't run again until interrupts are on.
unsafe impl GlobalAlloc for Locked<FixedSizeBlockAllocator> {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        x86_64::instructions::interrupts::without_interrupts(|| unsafe { self.alloc_inner(layout) })
    }
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        x86_64::instructions::interrupts::without_interrupts(|| unsafe { self.dealloc_inner(ptr, layout) })
    }
}

impl Locked<FixedSizeBlockAllocator> {
    unsafe fn alloc_inner(&self, layout: Layout) -> *mut u8 {
        let mut allocator = self.lock();
        match list_index(&layout) {
            Some(index) => match allocator.list_heads[index].take() {
                // A free block of the right size exists: pop it off the list.
                Some(node) => {
                    allocator.list_heads[index] = node.next.take();
                    node as *mut ListNode as *mut u8
                }
                // List empty: carve a new block out of the fallback heap.
                None => {
                    let block_size = BLOCK_SIZES[index];
                    let layout = Layout::from_size_align(block_size, block_size).unwrap();
                    allocator.fallback_alloc(layout)
                }
            },
            None => allocator.fallback_alloc(layout), // too big for any block
        }
    }

    unsafe fn dealloc_inner(&self, ptr: *mut u8, layout: Layout) {
        let mut allocator = self.lock();
        match list_index(&layout) {
            Some(index) => {
                // Push the freed block onto the front of its size's list.
                let new_node = ListNode { next: allocator.list_heads[index].take() };
                assert!(mem::size_of::<ListNode>() <= BLOCK_SIZES[index]);
                assert!(mem::align_of::<ListNode>() <= BLOCK_SIZES[index]);
                let new_node_ptr = ptr as *mut ListNode;
                unsafe {
                    new_node_ptr.write(new_node);
                    allocator.list_heads[index] = Some(&mut *new_node_ptr);
                }
            }
            None => {
                let ptr = NonNull::new(ptr).unwrap();
                unsafe { allocator.fallback_allocator.deallocate(ptr, layout) };
            }
        }
    }
}