Demilade Sonuga's blog

All posts

Heap Allocator III

2023-03-24 · 34 min read

In this post, we're going to be testing our heap allocator.

So far, we haven't done one iota of automated testing. While that should be a problem in a real-life production software scenario, it isn't one here. But we'll still need it. The reason for this is that the heap allocator we created is going to be used as a foundation in our vector and box. If there's a bug here, when we're using vectors and boxes, it could easily manifest itself as mysterious errors.

With automated testing, we can verify that this unit is working well, so that in the case when such a mysterious error does pop up, we can be sure to an extent that the problem is not in this lower layer of the hierarchy.

Figuring Out How To Test It

Before we can get on with testing, we first need to do some thinking. How are we going to test a heap allocator reliably in such a way that whenever the test runs, we can be sure that it will pass only if the code is correct and it will fail if there is a bug? Dijstkra once said that automated testing can only prove the presence of bugs, but not the absence of them. So, what we should actually be asking is, how can we write tests such that they will show the presence of as many bugs as possible?

But to answer these questions, we first need to know what a bug in our allocator will look like. This all boils down to what exactly our allocator is supposed to be doing in certain scenarios. We already know what these scenarios look like from this post.

What's left to figure out is how to check that these scenarios hold up in our implementation.

To do this, we need memory that the allocator will manage. Perform operations, like the ones in the scenarios in that post and check to make sure that the outcome is what we expected.

Testing

First, let's create a dummy test function.

In main.rs

mod allocator;

In allocator.rs

#[cfg(test)]
mod tests {

#[test]
fn test_dummy() {
assert_eq!(1, 2);
}
}

Upon running, you'll get an error:

error[E0152]: found duplicate lang item `panic_impl`
   --> blasterball/src/main.rs:268:1
    |
268 | fn panic_handler(panic_info: &core::panic::PanicInfo) -> ! {
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    |
    = note: the lang item is first defined in crate `std` (which `test` depends on)

The reason we're having this error is that cargo test depends on the standard library. And since the standard library already has a panic handler implementation, we can't have ours. One thing we can do to resolve this is to delete our panic handler whenever we want to run our tests, but that will be tedious, deleting and pasting it back afterward. Instead, we can use a conditional compilation attribute #[cfg_attr()] to include code based on a condition.

In main.rs

// DELETED: #[panic_handler]
#[cfg_attr(not(test), panic_handler)] // NEW
fn panic_handler(panic_info: &core::panic::PanicInfo) -> ! {
let screen = get_screen().unwrap();
write!(screen, "{}", panic_info);
loop {}
}

The #[cfg_attr(not(test), panic_handler)] tells the compiler that whenever we're not testing, the function should be the panic handler. And whenever we're testing, the function should be skipped.

Running again will give a bunch of linker errors. If you look through them, you'll see a part like this:

= note: /usr/bin/ld: /usr/lib/gcc/x86_64-redhat-linux/12/../../../../lib64/Scrt1.o: in function `_start':
        (.text+0x1b): undefined reference to `main'
        collect2: error: ld returned 1 exit status

A main function is needed for the tests to run, but we marked our game with #![no_main]. To resolve this, another #[cfg_attr] will do.

In main.rs

#![no_std]
// DELETED: #![no_main]
#![cfg_attr(not(test), no_main)]
#![feature(abi_efiapi, abi_x86_interrupt)]
#![feature(panic_info_message)]

Running now will result in a test that fails as expected.

Moving on now, to test our allocator, we need memory to manage. To do this, we can use the vector provided by the standard library. At the beginning of this blog, I said that we won't use the standard library to build the game, but right now, we're not building it. We're just testing an aspect of it.

Anyways, with a vector, we can get memory from the OS's heap and then use that memory to test our allocator's management.

Firstly, we'll need to import Vec from the standard library. To do this, we'll also need another conditional compilation configuration.

// DELETED: #![no_std]
#![cfg_attr(not(test), no_std)] // NEW
#![cfg_attr(not(test), no_main)]
#![feature(abi_efiapi, abi_x86_interrupt)]
#![feature(panic_info_message)]

Now, a function to create an allocator managing an amount of memory:

#[cfg(test)]
mod tests {
use super::*;
use std::mem::ManuallyDrop;
use std::vec::Vec as StdVec;

/* DELETED:
#[test]
fn test_dummy() {
assert_eq!(1, 2);
}
*/


fn new_allocator(memory_size: usize) -> LinkedListAllocator {
// Allocate `memory_size` bytes from the OS's heap
// Using `ManuallyDrop` so the memory won't be freed after a return from this function
let mem: ManuallyDrop<StdVec<u8>> = ManuallyDrop::new(StdVec::with_capacity(memory_size));
let mem_ptr = mem.as_ptr() as *mut u8;
let dummy_node_ptr = Box::into_raw(Box::new(ListNode {
size: 0,
next: None
}));
let mut allocator = LinkedListAllocator {
// A dummy head, as expected by the allocator functions
head: dummy_node_ptr
};
// Adding the memory to manage
unsafe {
allocator.add_free_region(MemChunk {
start_addr: mem_ptr as usize,
size: memory_size
});
}
return allocator;
}
}

The mem is just a vector of bytes. Allocating a vector with a capacity of 4Kib simply results in an allocation of 4Kib memory.

At the end of a function's execution, structures with Drop implementations, like vectors, have their Drop implementations executed. When a vector is dropped, its associated memory is deallocated. So, trying to use that memory after the function's return will result in a use-after-free error, an anathema in Rust. To avoid these problems, we use ManuallyDrop to tell the compiler not to drop the vector after the function returns, so we are free to use the memory for testing.

For the same reasons of avoiding use-after-free errors, the Box::into_raw function is used on the dummy node that serves as the head to tell the compiler "I want to handle the memory management of this box myself".

The LinkedListAllocator is initialized with a dummy node because, as noted in the previous post, this is a requirement for our functions to work. Lastly, the memory we allocated is passed to add_free_region for the allocator to manage.

With this function in place, we can now get to testing.

Using this post for our scenarios, we start from the first: a heap of 10Kib from which 3Kib is allocated.

Before this 3Kib is allocated, the heap should have only one node and it should be a 10Kib node. After the allocation, the heap should still have just one node and it should be a 7Kib node.

#[cfg(test)]
mod tests {
use super::*;
use std::mem::ManuallyDrop;
use std::vec::Vec as StdVec;

// One kilobyte is 2^10 bytes
const ONE_KIB: usize = 2usize.pow(10);

#[test]
fn test_allocate_3Kib() {
unsafe {
let mut allocator = new_allocator(10 * ONE_KIB);
let first_region = (*allocator.head).next.unwrap();
assert_eq!((*first_region).size, 10 * ONE_KIB, "The first node should be 10Kib");
assert_eq!((*first_region).next, None, "There should be only one node");
let ptr = allocator.alloc(3 * ONE_KIB, 1);
assert!(ptr.is_some(), "Allocation should have been successful");
let first_region = (*allocator.head).next.unwrap();
assert_eq!((*first_region).size, 7 * ONE_KIB, "The first node should now be 7Kib");
assert_eq!((*first_region).next, None, "There should still be only one node");
}
}

// ...Others
}

In this test, an allocator is first created. Assertions are made to verify that the size of the single node before allocating is 10Kib and the size after is 7Kib. We call alloc with an alignment of 1 to make testing easier.

This covers the first scenario in the Heap Allocator I post.

According to the standards of some development methodologies like test-driven development, this testing that we've just done can't really be trusted because it has never failed. According to test-driven development, we should even be writing tests before writing the code we're testing and make sure that those tests fail as expected for the right reasons along the way.

Although we're not doing test-driven development, it will still be useful to see the tests fail in ways that would be expected. For instance, if the argument in assert_eq! was 8 * ONE_KIB instead of the 7Kib, then it should fail. You should mess around with the code and the tests to ensure they fail as you expect them to.

The next scenario is the same 3Kib request followed by a 4Kib then a 2Kib request.

#[test]
fn test_allocate_3_4_2Kib() {
unsafe {
let mut allocator = new_allocator(10 * ONE_KIB);
let first_region = (*allocator.head).next.unwrap();
assert_eq!((*first_region).size, 10 * ONE_KIB, "The first node should be 10Kib");
assert_eq!((*first_region).next, None, "There should be only one node");
let ptr_3 = allocator.alloc(3 * ONE_KIB, 1);
let ptr_4 = allocator.alloc(4 * ONE_KIB, 1);
let ptr_2 = allocator.alloc(2 * ONE_KIB, 1);
assert!(ptr_3.is_some(), "Allocation of 3Kib should have been successful");
assert!(ptr_4.is_some(), "Allocation of 4Kib should have been successful");
assert!(ptr_2.is_some(), "Allocation of 2Kib should have been successful");
let first_region = (*allocator.head).next.unwrap();
assert_eq!((*first_region).size, 1 * ONE_KIB, "The first node should now be 1Kib");
assert_eq!((*first_region).next, None, "There should still be only one node");
}
}

In this test, allocations of 3, 4 and 2Kibs are made. The result should be a single node of size 1Kib.

In the next scenario, the same allocations of 3, 4, and 2Kibs are made, but the 4Kib is deallocated. The result should be two nodes of 4 and 1Kibs, respectively.

#[test]
fn test_dealloc_3_4_2Kib1() {
unsafe {
let (mut allocator, four_kib_ptr) = setup_3_4_2_dealloc_test();
let first_region = (*allocator.head).next.unwrap();
assert_eq!((*first_region).size, 1 * ONE_KIB, "The first node should be 1Kib");
assert_eq!((*first_region).next, None, "There should still be only one node");
allocator.dealloc(four_kib_ptr, 4 * ONE_KIB);
let first_region = (*allocator.head).next.unwrap();
assert_eq!((*first_region).size, 4 * ONE_KIB, "The first node should now be 4Kib");
assert!((*first_region).next.is_some(), "There should be a next node");
let second_region = (*first_region).next.unwrap();
assert_eq!((*second_region).size, 1 * ONE_KIB, "The second node should be 1Kib");
assert_eq!((*second_region).next, None, "There should be only two nodes");
}
}

unsafe fn setup_3_4_2_dealloc_test1() -> (LinkedListAllocator, *mut u8) {
let mut allocator = new_allocator(10 * ONE_KIB);
let ptr_3 = allocator.alloc(3 * ONE_KIB, 1);
let ptr_4 = allocator.alloc(4 * ONE_KIB, 1).unwrap();
let ptr_2 = allocator.alloc(2 * ONE_KIB, 1);
(allocator, ptr_4)
}

In the test, the 3, 4 and 2Kibs are first allocated and it is verified that there is only one node with a size of 1Kib. After the deallocation of the 4Kib, it is verified that there are now two nodes: the first 4Kib and the next 1Kib, as expected.

In the next scenario, the same thing happens, but the after the deallocation of the 4Kib, the 3Kib, too, is freed. After this 3Kib is freed, due to of merging of nodes, the list will have two nodes: 7Kib and 1Kib, respectively.

#[test]
fn test_dealloc_3_4_2Kib2() {
unsafe {
let (mut allocator, three_kib_ptr) = setup_3_4_2_dealloc_test2();
let first_region = (*allocator.head).next.unwrap();
assert_eq!((*first_region).size, 4 * ONE_KIB, "First node should be 4Kib");
assert!((*first_region).next.is_some(), "There should be a second node");
let second_region = (*first_region).next.unwrap();
assert_eq!((*second_region).size, 1 * ONE_KIB, "The second node should be 1Kib");
assert!((*second_region).next.is_none(), "There should be only two nodes");

allocator.dealloc(three_kib_ptr, 3 * ONE_KIB);

let first_region = (*allocator.head).next.unwrap();
assert_eq!((*first_region).size, 7 * ONE_KIB, "The first node should be 7Kib");
assert!((*first_region).next.is_some(), "There should be a second node");
let second_region = (*first_region).next.unwrap();
assert_eq!((*second_region).size, 1 * ONE_KIB, "The second node should be 1Kib");
assert!((*second_region).next.is_none(), "There should be only two nodes");

}
}

unsafe fn setup_3_4_2_dealloc_test2() -> (LinkedListAllocator, *mut u8) {
let mut allocator = new_allocator(10 * ONE_KIB);
let ptr_3 = allocator.alloc(3 * ONE_KIB, 1).unwrap();
let ptr_4 = allocator.alloc(4 * ONE_KIB, 1).unwrap();
let ptr_2 = allocator.alloc(2 * ONE_KIB, 1);
allocator.dealloc(ptr_4, 4 * ONE_KIB);
(allocator, ptr_3)
}

Like in the previous scenario, the setup function is first used to get the preliminaries in place. Allocates the 3, 4, and 2Kib, then deallocates the 4Kib. The test function then deallocates the 3Kib and verifies that the first node is merged into a 7Kib region and the next node remains a 1Kib one.

The last scenario continues from the previous one and ends with the deallocation of the 2Kib chunk resulting in a single 10Kib node again.

#[test]
fn test_dealloc_3_4_2Kib3() {
unsafe {
let (mut allocator, two_kib_ptr) = setup_3_4_2_dealloc_test3();
let first_region = (*allocator.head).next.unwrap();
assert_eq!((*first_region).size, 7 * ONE_KIB, "First node should be 7Kib");
assert!((*first_region).next.is_some(), "There should be a second node");
let second_region = (*first_region).next.unwrap();
assert_eq!((*second_region).size, 1 * ONE_KIB, "The second node should be 1Kib");
assert!((*second_region).next.is_none(), "There should be only two nodes");

allocator.dealloc(two_kib_ptr, 2 * ONE_KIB);

let first_region = (*allocator.head).next.unwrap();
assert_eq!((*first_region).size, 10 * ONE_KIB, "The first node should be 10Kib");
assert!((*first_region).next.is_none(), "There should not be a second node");
}
}

unsafe fn setup_3_4_2_dealloc_test3() -> (LinkedListAllocator, *mut u8) {
let mut allocator = new_allocator(10 * ONE_KIB);
let ptr_3 = allocator.alloc(3 * ONE_KIB, 1).unwrap();
let ptr_4 = allocator.alloc(4 * ONE_KIB, 1).unwrap();
let ptr_2 = allocator.alloc(2 * ONE_KIB, 1).unwrap();
allocator.dealloc(ptr_4, 4 * ONE_KIB);
allocator.dealloc(ptr_3, 3 * ONE_KIB);
(allocator, ptr_2)
}

The tests should pass, but you should mess around with them. Make them fail predictably, just to verify that they are actually testing something.

These tests are not extensive at all. They don't even take alignment into consideration. But this was just a quick way of verifying that the allocator does its most basic functions reasonably.

Take Away

For the full code, go to the repo

In The Next Post

We're going to start creating a Box.