Skip to content

Commit

Permalink
[fix] bug related with base alignment
Browse files Browse the repository at this point in the history
  • Loading branch information
@hky1999
hky1999 committed Dec 11, 2024
1 parent adb3c90 commit 8dcb9b6
Showing 1 changed file with 194 additions and 6 deletions.
200 changes: 194 additions & 6 deletions src/bitmap.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,13 @@ use bitmap_allocator::BitAlloc;

use crate::{AllocError, AllocResult, BaseAllocator, PageAllocator};

const MAX_ALIGN_1GB: usize = 0x4000_0000;

cfg_if::cfg_if! {
if #[cfg(feature = "page-alloc-1t")] {
if #[cfg(test)] {
/// Use 4GB memory for testing.
type BitAllocUsed = bitmap_allocator::BitAlloc1M;
} else if #[cfg(feature = "page-alloc-1t")] {
/// Support max 256M * PAGE_SIZE = 1TB memory (assume that PAGE_SIZE = 4KB).
type BitAllocUsed = bitmap_allocator::BitAlloc256M;
} else if #[cfg(feature = "page-alloc-64g")] {
Expand Down Expand Up @@ -48,11 +53,20 @@ impl<const PAGE_SIZE: usize> BitmapPageAllocator<PAGE_SIZE> {
impl<const PAGE_SIZE: usize> BaseAllocator for BitmapPageAllocator<PAGE_SIZE> {
fn init(&mut self, start: usize, size: usize) {
assert!(PAGE_SIZE.is_power_of_two());
let end = super::align_down(start + size, PAGE_SIZE);
let start = super::align_up(start, PAGE_SIZE);
self.base = start;

// Range for real: [align_up(start, PAGE_SIZE), align_down(start + size, PAGE_SIZE))
let start = crate::align_up(start, PAGE_SIZE);
let end = crate::align_down(start + size, PAGE_SIZE);
self.total_pages = (end - start) / PAGE_SIZE;
self.inner.insert(0..self.total_pages);

// Calculate the base offset stored in the real [`BitAlloc`] instance.
self.base = crate::align_down(start, MAX_ALIGN_1GB);

// Range in bitmap: [start - self.base, start - self.base + total_pages * PAGE_SIZE)
let start = start - self.base;
let start_idx = start / PAGE_SIZE;

self.inner.insert(start_idx..start_idx + self.total_pages);
}

fn add_memory(&mut self, _start: usize, _size: usize) -> AllocResult {
Expand All @@ -64,6 +78,10 @@ impl<const PAGE_SIZE: usize> PageAllocator for BitmapPageAllocator<PAGE_SIZE> {
const PAGE_SIZE: usize = PAGE_SIZE;

fn alloc_pages(&mut self, num_pages: usize, align_pow2: usize) -> AllocResult<usize> {
assert!(
align_pow2 <= MAX_ALIGN_1GB,
"allocator does not support align > 1GB"
);
if align_pow2 % PAGE_SIZE != 0 {
return Err(AllocError::InvalidParam);
}
Expand Down Expand Up @@ -91,6 +109,10 @@ impl<const PAGE_SIZE: usize> PageAllocator for BitmapPageAllocator<PAGE_SIZE> {
num_pages: usize,
align_pow2: usize,
) -> AllocResult<usize> {
assert!(
align_pow2 <= MAX_ALIGN_1GB,
"allocator does not support align > 1GB"
);
if align_pow2 % PAGE_SIZE != 0 {
return Err(AllocError::InvalidParam);
}
Expand Down Expand Up @@ -142,4 +164,170 @@ impl<const PAGE_SIZE: usize> PageAllocator for BitmapPageAllocator<PAGE_SIZE> {
}
}

impl<const PAGE_SIZE: usize> BitmapPageAllocator<PAGE_SIZE> {}
#[cfg(test)]
mod tests {
use super::*;

const PAGE_SIZE: usize = 4096;

#[test]
fn test_bitmap_page_allocator_one_page() {
let mut allocator = BitmapPageAllocator::<PAGE_SIZE>::new();
allocator.init(PAGE_SIZE, PAGE_SIZE);

assert_eq!(allocator.total_pages(), 1);
assert_eq!(allocator.used_pages(), 0);
assert_eq!(allocator.available_pages(), 1);

let addr = allocator.alloc_pages(1, PAGE_SIZE).unwrap();
assert_eq!(addr, 0x1000);
assert_eq!(allocator.used_pages(), 1);
assert_eq!(allocator.available_pages(), 0);

allocator.dealloc_pages(addr, 1);
assert_eq!(allocator.used_pages(), 0);
assert_eq!(allocator.available_pages(), 1);

let addr = allocator.alloc_pages(1, PAGE_SIZE).unwrap();
assert_eq!(addr, 0x1000);
assert_eq!(allocator.used_pages(), 1);
assert_eq!(allocator.available_pages(), 0);
}

#[test]
fn test_bitmap_page_allocator_size_2g() {
const SIZE_1G: usize = 1024 * 1024 * 1024;
const SIZE_2G: usize = 2 * SIZE_1G;

const TEST_BASE_ADDR: usize = SIZE_1G + PAGE_SIZE;

let mut allocator = BitmapPageAllocator::<PAGE_SIZE>::new();
allocator.init(TEST_BASE_ADDR, SIZE_2G);

let mut num_pages = 1;
// Test allocation and deallocation of 1, 10, 100, 1000 pages.
while num_pages <= 1000 {
assert_eq!(allocator.total_pages(), SIZE_2G / PAGE_SIZE);
assert_eq!(allocator.used_pages(), 0);
assert_eq!(allocator.available_pages(), SIZE_2G / PAGE_SIZE);

let addr = allocator.alloc_pages(num_pages, PAGE_SIZE).unwrap();
assert_eq!(addr, TEST_BASE_ADDR);
assert_eq!(allocator.used_pages(), num_pages);
assert_eq!(allocator.available_pages(), SIZE_2G / PAGE_SIZE - num_pages);

allocator.dealloc_pages(addr, num_pages);
assert_eq!(allocator.used_pages(), 0);
assert_eq!(allocator.available_pages(), SIZE_2G / PAGE_SIZE);

num_pages *= 10;
}

// Test allocation and deallocation of 1, 10, 100 pages with alignment.
num_pages = 1;
let mut align = PAGE_SIZE;
while align <= MAX_ALIGN_1GB {
assert_eq!(allocator.total_pages(), SIZE_2G / PAGE_SIZE);
assert_eq!(allocator.used_pages(), 0);
assert_eq!(allocator.available_pages(), SIZE_2G / PAGE_SIZE);

let addr = allocator.alloc_pages(num_pages, align).unwrap();
assert_eq!(addr, crate::align_up(TEST_BASE_ADDR, align));
assert_eq!(allocator.used_pages(), num_pages);
assert_eq!(allocator.available_pages(), SIZE_2G / PAGE_SIZE - num_pages);

allocator.dealloc_pages(addr, num_pages);
assert_eq!(allocator.used_pages(), 0);
assert_eq!(allocator.available_pages(), SIZE_2G / PAGE_SIZE);

num_pages *= 10;
align <<= 9;
}

num_pages = 1;
align = PAGE_SIZE;
let mut i = 0;
let mut addrs = [(0, 0); 3];
let mut used_pages = 0;

// Test allocation of 1, 10, 100 pages with alignment.
while i < 3 {
assert_eq!(allocator.total_pages(), SIZE_2G / PAGE_SIZE);
assert_eq!(allocator.used_pages(), used_pages);
assert_eq!(
allocator.available_pages(),
SIZE_2G / PAGE_SIZE - used_pages
);

let addr = allocator.alloc_pages(num_pages, align).unwrap();
assert!(crate::is_aligned(addr, align));

addrs[i] = (addr, num_pages);

used_pages += num_pages;
assert_eq!(allocator.used_pages(), used_pages);
assert_eq!(
allocator.available_pages(),
SIZE_2G / PAGE_SIZE - used_pages
);

num_pages *= 10;
align <<= 9;
i += 1;
}

i = 0;
// Test deallocation of 1, 10, 100 pages.
while i < 3 {
let addr = addrs[i].0;
let num_pages = addrs[i].1;
allocator.dealloc_pages(addr, num_pages);

used_pages -= num_pages;
assert_eq!(allocator.used_pages(), used_pages);
assert_eq!(
allocator.available_pages(),
SIZE_2G / PAGE_SIZE - used_pages
);
i += 1;
}

assert_eq!(allocator.used_pages(), 0);
assert_eq!(allocator.available_pages(), SIZE_2G / PAGE_SIZE);

// Test allocation of 1, 10, 100 pages with alignment at a specific address.
num_pages = 1;
align = PAGE_SIZE;
i = 0;
used_pages = 0;
let mut test_addr_base = TEST_BASE_ADDR;

while i < 3 {
assert_eq!(allocator.total_pages(), SIZE_2G / PAGE_SIZE);
assert_eq!(allocator.used_pages(), used_pages);
assert_eq!(
allocator.available_pages(),
SIZE_2G / PAGE_SIZE - used_pages
);

let addr = allocator
.alloc_pages_at(test_addr_base, num_pages, align)
.unwrap();
assert_eq!(addr, test_addr_base);

used_pages += num_pages;
assert_eq!(allocator.used_pages(), used_pages);
assert_eq!(
allocator.available_pages(),
SIZE_2G / PAGE_SIZE - used_pages
);

num_pages *= 10;
align <<= 9;

test_addr_base = crate::align_up(test_addr_base + num_pages * PAGE_SIZE, align);

i += 1;
}
}
}

0 comments on commit 8dcb9b6

Please sign in to comment.