//! Memory bookkeeping. use prelude::*; use brk; use vec::Vec; use core::marker::PhantomData; use core::ops::Range; use core::{ptr, cmp, mem}; /// The memory bookkeeper. /// /// This is the main component of ralloc. Its job is to keep track of the free blocks in a /// structured manner, such that allocation, reallocation, and deallocation are all efficient. /// Particularly, it keeps a list of blocks, commonly called the "block pool". This list is kept. /// Entries in the block pool can be "empty", meaning that you can overwrite the entry without /// breaking consistency. /// /// Only making use of only [`alloc`](#method.alloc), [`free`](#method.free), /// [`realloc`](#method.realloc) (and following their respective assumptions) guarantee that no /// buffer overrun, arithmetic overflow, panic, or otherwise unexpected crash will happen. pub struct Bookkeeper { /// The internal block pool. /// /// # Guarantees /// /// Certain guarantees are made: /// /// 1. The list is always sorted with respect to the block's pointers. /// 2. No two consecutive or empty block delimited blocks are adjacent, except if the right /// block is empty. /// 3. There are no trailing empty blocks. /// /// These are invariants assuming that only the public methods are used. pool: Vec, /// The number of bytes currently allocated. #[cfg(feature = "debug_tools")] allocated: usize, /// The "breaker", i.e. the fresh allocator. /// /// This has as job as acquiring new memory through some external source (e.g. BRK or the /// global allocator). breaker: PhantomData, } impl Bookkeeper { /// Create a new, empty block pool. /// /// This will make no allocations or BRKs. #[inline] #[cfg(feature = "debug_tools")] pub const fn new() -> Bookkeeper { Bookkeeper { pool: Vec::new(), allocated: 0, } } #[inline] #[cfg(not(feature = "debug_tools"))] pub const fn new() -> Bookkeeper { Bookkeeper { pool: Vec::new(), } } /// Allocate a chunk of memory. /// /// This function takes a size and an alignment. From these a fitting block is found, to which /// a pointer is returned. The block returned is guaranteed to be aligned to `align`. /// /// # Example /// /// We start with our initial segment. /// /// ```notrust /// Address space /// I---------------------------------I /// B /// l /// k /// s /// ``` /// /// We then split it at the aligner, which is used for making sure that /// the pointer is aligned properly. /// /// ```notrust /// Address space /// I------I /// B ^ I--------------------------I /// l al /// k /// s /// ``` /// /// We then use the remaining block, but leave the excessive space. /// /// ```notrust /// Address space /// I------I /// B I--------I /// l \_________________/ /// k our allocated block. /// s /// ``` /// /// A block representing the marked area is then returned. pub fn alloc(&mut self, size: usize, align: usize) -> Block { // TODO: scan more intelligently. // Logging. log!(self.pool, "Allocating {} bytes with alignment {}.", size, align); if let Some((n, b)) = self.pool.iter_mut().enumerate().filter_map(|(n, i)| { if i.size() >= size { // Try to split at the aligner. i.align(align).and_then(|(mut a, mut b)| { if b.size() >= size { // Override the old block. *i = a; Some((n, b)) } else { // Put the split block back together and place it back in its spot. a.merge_right(&mut b).unwrap(); *i = a; None } }) } else { None } }).next() { if self.pool[n].is_empty() { // For empty alignment invariant. let _ = self.remove_at(n); } let (res, excessive) = b.split(size); // Mark the excessive space as free. // There are many corner cases that make knowing where to insert it difficult // so we search instead. let bound = self.find_bound(&excessive); self.free_bound(bound, excessive); // Check consistency. self.check(); debug_assert!(res.aligned_to(align), "Alignment failed."); debug_assert!(res.size() == size, "Requested space does not match with the returned \ block."); self.leave(res) } else { // No fitting block found. Allocate a new block. let res = self.alloc_fresh(size, align); // "Leave" the allocator. self.leave(res) } } /// Free a memory block. /// /// After this have been called, no guarantees are made about the passed pointer. If it want /// to, it could begin shooting laser beams. /// /// Freeing an invalid block will drop all future guarantees about this bookkeeper. /// /// # Example /// /// ```notrust /// Address space /// I------I /// B I--------I /// l \_________________/ /// k the used block we want to deallocate. /// s /// ``` /// /// If the blocks are adjacent, we merge them: /// /// ```notrust /// Address space /// I------I /// B I-----------------I /// l I--------I /// k /// s /// ``` /// /// This gives us: /// /// ```notrust /// Address space /// I------------------------I /// B I--------I /// l /// k /// s /// ``` /// /// And we're done. If it cannot be done, we insert the block, while keeping the list sorted. /// See [`insert`](#method.insert) for details. #[inline] pub fn free(&mut self, block: Block) { // Just logging for the unlucky people debugging this shit. No problem. log!(self.pool, "Freeing {:?}...", block); // "Enter" the allocator. let block = self.enter(block); self.reserve_more(1); // Binary search for the block. let bound = self.find_bound(&block); // Free the given block. self.free_bound(bound, block); } /// Reallocate memory. /// /// If necessary (inplace reallocation is not possible or feasible) it will allocate a new /// buffer, fill it with the contents of the old buffer, and deallocate the replaced buffer. /// /// The following guarantees are made: /// /// 1. The returned block is valid and aligned to `align`. /// 2. The returned block contains the same data byte-for-byte as the original buffer. /// /// The data will be truncated if `new_size` is smaller than `block`'s size. /// /// # Example /// /// We will first try to perform an in-place reallocation, and if that fails, we will use /// memmove. /// /// ```notrust /// Address space /// I------I /// B \~~~~~~~~~~~~~~~~~~~~~/ /// l needed /// k /// s /// ``` /// /// We simply find the block next to our initial block. If this block is free and have /// sufficient size, we will simply merge it into our initial block, and leave the excessive /// space as free. If these conditions are not met, we have to allocate a new list, and then /// deallocate the old one, after which we use memmove to copy the data over to the newly /// allocated list. pub fn realloc(&mut self, block: Block, new_size: usize, align: usize) -> Block { // Find the index bound. let ind = self.find_bound(&block); // Logging. log!(self.pool;ind, "Reallocating {:?} to size {} with align {}...", block, new_size, align); // "Leave" the allocator. let block = self.enter(block); // Try to do an inplace reallocation. match self.realloc_inplace_bound(ind, block, new_size) { Ok(block) => self.leave(block), Err(block) => { // Reallocation cannot be done inplace. // Allocate a new block with the same size. let mut res = self.alloc(new_size, align); // Copy the old data to the new location. block.copy_to(&mut res); // Free the old block. // Allocation may have moved insertion so we search again. let bound = self.find_bound(&block); self.free_bound(bound, block); // Check consistency. self.check(); debug_assert!(res.aligned_to(align), "Alignment failed."); debug_assert!(res.size() >= new_size, "Requested space does not match with the \ returned block."); // Leave the allocator. self.leave(res) }, } } /// Extend/shrink the buffer inplace. /// /// This will try to extend the buffer without copying, if the new size is larger than the old /// one. If not, truncate the block and place it back to the pool. /// /// On failure, return `Err(Block)` with the old _intact_ block. Shrinking cannot fail. /// /// This shouldn't be used when the index of insertion is known, since this performs an binary /// search to find the blocks index. When you know the index use /// [`realloc_inplace_bound`](#method.realloc_inplace_bound.html). #[inline] pub fn realloc_inplace(&mut self, block: Block, new_size: usize) -> Result { // Logging. log!(self.pool, "Reallocating {:?} inplace to {}...", block, new_size); // Find the bounds of given block. let bound = self.find_bound(&block); // Go for it! let res = self.realloc_inplace_bound(bound, block, new_size); // Check consistency. debug_assert!(res.as_ref().ok().map_or(true, |x| x.size() == new_size), "Requested space \ does not match with the returned block."); res } /// Reallocate a block on a know index bound inplace. /// /// See [`realloc_inplace`](#method.realloc_inplace.html) for more information. fn realloc_inplace_bound(&mut self, ind: Range, mut block: Block, new_size: usize) -> Result { // Logging. log!(self.pool;ind, "Try inplace reallocating {:?} to size {}.", block, new_size); /// Assertions... debug_assert!(self.find(&block) == ind.start, "Block is not inserted at the appropriate \ index."); if new_size <= block.size() { // Shrink the block. log!(self.pool;ind, "Shrinking {:?}.", block); // Split the block in two segments, the main segment and the excessive segment. let (block, excessive) = block.split(new_size); // Free the excessive segment. self.free_bound(ind, excessive); // Make some assertions to avoid dumb bugs. debug_assert!(block.size() == new_size, "Block wasn't shrinked properly."); // Run a consistency check. self.check(); return Ok(block); // We check if `ind` is the end of the array. } else { let mut mergable = false; if let Some(entry) = self.pool.get_mut(ind.end) { mergable = entry.size() + block.size() >= new_size && block.left_to(entry); } // Note that we are sure that no segments in the array are adjacent (unless they have size // 0). This way we know that we will, at maximum, need one and only one block for extending // the current block. if mergable { // Logging... log!(self.pool;ind, "Merging {:?} to the right.", block); // We'll merge it with the block at the end of the range. block.merge_right(&mut self.remove_at(ind.end)).unwrap(); // Merge succeeded. // Place the excessive block back. let (res, excessive) = block.split(new_size); // Remove_at may have shortened the vector. if ind.start == self.pool.len() { self.push_no_reserve(excessive); } else if !excessive.is_empty() { self.pool[ind.start] = excessive; } // Block will still not be adjacent, due to `excessive` being guaranteed to not be // adjacent to the next block. // Run a consistency check. self.check(); // TODO, drop excessive space return Ok(res); } } Err(block) } /// Free a block placed in some index bound. /// /// This will at maximum insert one element. /// /// See [`free`](#method.free) for more information. #[inline] fn free_bound(&mut self, ind: Range, mut block: Block) { // Logging. log!(self.pool;ind, "Freeing {:?}.", block); // Short circuit in case of empty block. if block.is_empty() { return; } // When compiled with `security`, we zero this block. block.sec_zero(); if ind.start == self.pool.len() { self.push_no_reserve(block); return; } // Assertions... debug_assert!(self.find(&block) == ind.start, "Block is not inserted at the appropriate \ index."); // Try to merge it with the block to the right. if ind.end < self.pool.len() && block.left_to(&self.pool[ind.end]) { block.merge_right(&mut self.remove_at(ind.end)).unwrap(); // The merging succeeded. We proceed to try to close in the possible gap. if ind.start != 0 && self.pool[ind.start - 1].merge_right(&mut block).is_ok() { self.check(); return; } // Dammit, let's try to merge left. } else if ind.start != 0 && self.pool[ind.start - 1].merge_right(&mut block).is_ok() { // Check consistency. self.check(); return; } // Well, it failed, so we insert it the old-fashioned way. self.insert(ind.start, block); // Check consistency. self.check(); } /// Allocate _fresh_ space. /// /// "Fresh" means that the space is allocated through the breaker. /// /// The returned pointer is guaranteed to be aligned to `align`. fn alloc_fresh(&mut self, size: usize, align: usize) -> Block { // Logging. log!(self.pool, "Fresh allocation of size {} with alignment {}.", size, align); // Break it to me! let res = B::alloc_fresh(size, align); // Check consistency. self.check(); res } /// Push two blocks to the block pool. /// /// This will append the blocks to the end of the block pool (and merge if possible). Make sure /// that these blocks has a value higher than any of the elements in the list, to keep it /// sorted. /// /// This guarantees linearity so that the blocks will be adjacent. #[inline] fn double_push(&mut self, block_a: Block, block_b: Block) { // Logging. log!(self.pool;self.pool.len(), "Pushing {:?} and {:?}.", block_a, block_b); // Catch stupid bug... debug_assert!(block_a <= block_b, "The first pushed block is not lower or equal to the second."); // Reserve extra elements. self.reserve_more(2); self.push_no_reserve(block_a); self.push_no_reserve(block_b); } /// Push an element without reserving. fn push_no_reserve(&mut self, mut block: Block) { // Logging. log!(self.pool;self.pool.len(), "Pushing {:?}.", block); // Short-circuit in case on empty block. if !block.is_empty() { // We will try to simply merge it with the last block. if let Some(x) = self.pool.last_mut() { if x.merge_right(&mut block).is_ok() { return; } } // Merging failed. Note that trailing empty blocks are not allowed, hence the last block is // the only non-empty candidate which may be adjacent to `block`. // We push. let res = self.pool.push(block); // Make some assertions. debug_assert!(res.is_ok(), "Push failed (buffer full)."); } self.check(); } /// Perform a binary search to find the appropriate place where the block can be insert or is /// located. /// /// It is guaranteed that no block left to the returned value, satisfy the above condition. #[inline] fn find(&mut self, block: &Block) -> usize { // TODO optimize this function. // Logging. log!(self.pool, "Searching (exact) for {:?}.", block); let ind = match self.pool.binary_search(block) { Ok(x) | Err(x) => x, }; let len = self.pool.len(); // Move left. ind - self.pool.iter_mut() .rev() .skip(len - ind) .take_while(|x| x.is_empty()) .count() } /// Perform a binary search to find the appropriate bound where the block can be insert or is /// located. /// /// It is guaranteed that no block left to the returned value, satisfy the above condition. #[inline] fn find_bound(&mut self, block: &Block) -> Range { // TODO optimize this function. // Logging. log!(self.pool, "Searching (bounds) for {:?}.", block); let mut left_ind = match self.pool.binary_search(block) { Ok(x) | Err(x) => x, }; let len = self.pool.len(); // Move left. left_ind -= self.pool.iter_mut() .rev() .skip(len - left_ind) .take_while(|x| x.is_empty()) .count(); let mut right_ind = match self.pool.binary_search(&block.empty_right()) { Ok(x) | Err(x) => x, }; // Move right. right_ind += self.pool.iter() .skip(right_ind) .take_while(|x| x.is_empty()) .count(); left_ind..right_ind } /// Insert a block entry at some index. /// /// If the space is non-empty, the elements will be pushed filling out the empty gaps to the /// right. If all places to the right is occupied, it will reserve additional space to the /// block pool. /// /// # Panics /// /// Panics on when `ind` is greater than the block pool's length. /// /// # Example /// /// We want to insert the block denoted by the tildes into our list. Perform a binary search to /// find where insertion is appropriate. /// /// ```notrust /// Address space /// I------I /// B < here I--------I /// l I------------I /// k /// s I---I /// I~~~~~~~~~~I /// ``` /// /// We keep pushing the blocks to the right to the next entry until a empty entry is reached: /// /// ```notrust /// Address space /// I------I /// B < here I--------I <~ this one cannot move down, due to being blocked. /// l /// k I------------I <~ thus we have moved this one down. /// s I---I /// I~~~~~~~~~~I /// ``` /// /// Repeating yields: /// /// ```notrust /// Address space /// I------I /// B < here /// l I--------I <~ this one cannot move down, due to being blocked. /// k I------------I <~ thus we have moved this one down. /// s I---I /// I~~~~~~~~~~I /// ``` /// /// Now an empty space is left out, meaning that we can insert the block: /// /// ```notrust /// Address space /// I------I /// B I----------I /// l I--------I /// k I------------I /// s I---I /// ``` /// /// The insertion is now completed. #[inline] fn insert(&mut self, ind: usize, block: Block) { // Logging. log!(self.pool;ind, "Inserting block {:?}...", block); // Bound check. assert!(self.pool.len() >= ind, "Insertion out of bounds."); // Some assertions... debug_assert!(self.pool.len() <= ind || block <= self.pool[ind], "Inserting at {} will make \ the list unsorted.", ind); debug_assert!(self.find(&block) == ind, "Block is not inserted at the appropriate index."); debug_assert!(!block.is_empty(), "Inserting an empty block."); // Find the next gap, where a used block were. let n = { // The element we search for. let elem = self.pool .iter() .skip(ind) .enumerate() .filter(|&(_, x)| x.is_empty()) .next() .map(|(n, _)| n); elem.unwrap_or_else(|| { // Reserve capacity. self.reserve_more(1); // We default to the end of the pool. self.pool.len() - ind }) }; // Log the operation. log!(self.pool;ind, "Moving {} blocks to the right.", n); unsafe { // TODO clean this mess up. if ind + n == self.pool.len() { // We will move a block into reserved memory but outside of the vec's bounds. For // that reason, we push an uninitialized element to extend the length, which will // be assigned in the memcpy. let res = self.pool.push(mem::uninitialized()); // Just some assertions... debug_assert!(res.is_ok(), "Push failed (buffer full)."); } // Memmove the elements. ptr::copy(self.pool.get_unchecked(ind) as *const Block, self.pool.get_unchecked_mut(ind + 1) as *mut Block, n); // Set the element. *self.pool.get_unchecked_mut(ind) = block; } // Check consistency. self.check(); } /// Remove a block. fn remove_at(&mut self, ind: usize) -> Block { // Logging. log!(self.pool;ind, "Removing block."); if ind == self.pool.len() - 1 { let res = self.pool[ind].pop(); // Make sure there are no trailing empty blocks. let new_len = self.pool.len() - self.pool.iter().rev().take_while(|x| x.is_empty()).count(); // Truncate the vector. self.pool.truncate(new_len); res } else { // Calculate the upper and lower bound let empty = self.pool[ind + 1].empty_left(); let empty2 = empty.empty_left(); // Replace the block at `ind` with the left empty block from `ind + 1`. let res = mem::replace(&mut self.pool[ind], empty); // Iterate over the pool from `ind` and down. let skip = self.pool.len() - ind; for place in self.pool.iter_mut().rev().skip(skip).take_while(|x| x.is_empty()) { // Empty the blocks. *place = empty2.empty_left(); } res } } /// Reserve space for the block pool. /// /// This will ensure the capacity is at least `needed` greater than the current length, /// potentially reallocating the block pool. fn reserve_more(&mut self, extra: usize) { // Logging. log!(bk.pool;bk.pool.len(), "Reserving {} past {}, currently has capacity {}.", extra, bk.pool.len(), bk.pool.capacity()); let needed = bk.pool.len() + extra; if needed > bk.pool.capacity() { B::realloc_pool(self, needed); // Check consistency. bk.check(); } } /// Leave the allocator. /// /// A block should be "registered" through this function when it leaves the allocated (e.g., is /// returned), these are used to keep track of the current heap usage, and memory leaks. #[inline] fn leave(&mut self, block: Block) -> Block { // Update the number of bytes allocated. #[cfg(feature = "debug_tools")] { self.allocated += block.size(); } block } /// Enter the allocator. /// /// A block should be "registered" through this function when it enters the allocated (e.g., is /// given as argument), these are used to keep track of the current heap usage, and memory /// leaks. #[inline] fn enter(&mut self, block: Block) -> Block { // Update the number of bytes allocated. #[cfg(feature = "debug_tools")] { self.allocated -= block.size(); } block } /// Perform consistency checks. /// /// This will check for the following conditions: /// /// 1. The list is sorted. /// 2. No blocks are adjacent. /// /// This is NOOP in release mode. fn check(&self) { if cfg!(debug_assertions) { // Logging. log!(self.pool, "Checking..."); // Reverse iterator over the blocks. let mut it = self.pool.iter().enumerate().rev(); if let Some((_, x)) = it.next() { // Make sure there are no leading empty blocks. assert!(!x.is_empty()); let mut next = x; for (n, i) in it { // Check if sorted. assert!(next >= i, "The block pool is not sorted at index, {} ({:?} < {:?}).", n, next, i); // Make sure no blocks are adjacent. assert!(!i.left_to(next) || i.is_empty(), "Adjacent blocks at index, {} ({:?} and \ {:?})", n, i, next); // Make sure an empty block has the same address as its right neighbor. assert!(!i.is_empty() || i == next, "Empty block not adjacent to right neighbor \ at index {} ({:?} and {:?})", n, i, next); // Set the variable tracking the previous block. next = i; } // Check for trailing empty blocks. assert!(!self.pool.last().unwrap().is_empty(), "Trailing empty blocks."); } // Logging... log!(self.pool, "Check OK!"); } } /// Attach this allocator to the current thread. /// /// This will make sure this allocator's data is freed to the pub unsafe fn attach(&mut self) { fn dtor(ptr: *mut Bookkeeper) { let alloc = *ptr; // Lock the global allocator. let global_alloc = allocator::GLOBAL_ALLOCATOR.lock(); // TODO, we know this is sorted, so we could abuse that fact to faster insertion in the // global allocator. // Free everything in the allocator. while let Some(i) = alloc.pool.pop() { global_alloc.free(i); } // Deallocate the vector itself. global_alloc.free(Block::from(alloc.pool)); // Gotta' make sure no memleaks are here. #[cfg(feature = "debug_tools")] alloc.assert_no_leak(); } sys::register_thread_destructor(self as *mut Bookkeeper, dtor).unwrap(); } /// Check for memory leaks. /// /// This will ake sure that all the allocated blocks have been freed. #[cfg(feature = "debug_tools")] fn assert_no_leak(&self) { assert!(self.allocated == self.pool.capacity() * mem::size_of::(), "Not all blocks \ freed. Total allocated space is {} ({} free blocks).", self.allocated, self.pool.len()); } } trait Breaker { /// Allocate _fresh_ space. /// /// "Fresh" means that the space is allocated through the breaker. /// /// The returned pointer is guaranteed to be aligned to `align`. fn alloc_fresh(bk: &mut Bookkeeper, size: usize, align: usize) -> Block; /// Realloate the block pool to some specified capacity. fn realloc_pool(bk: &mut Bookkeeper, cap: usize); } /// SBRK fresh allocator. /// /// This will extend the data segment whenever new memory is needed. Since this includes leaving /// userspace, this shouldn't be used when other allocators are available (i.e. the bookkeeper is /// local). struct Sbrk; impl Breaker for Sbrk { #[inline] fn alloc_fresh(bk: &mut Bookkeeper, size: usize, align: usize) -> Block { // Obtain what you need. let (alignment_block, res, excessive) = brk::get(size, align); // Add it to the list. This will not change the order, since the pointer is higher than all // the previous blocks. bk.double_push(alignment_block, excessive); res } #[inline] fn realloc_pool(bk: &mut Bookkeeper, extra: usize) { // TODO allow BRK-free non-inplace reservations. // TODO Enable inplace reallocation in this position. // Reallocate the block pool. // Make a fresh allocation. let size = (needed + cmp::min(bk.pool.capacity(), 200 + bk.pool.capacity() / 2) // We add: + 1 // block for the alignment block. + 1 // block for the freed vector. + 1 // block for the excessive space. ) * mem::size_of::(); let (alignment_block, alloc, excessive) = brk::get(size, mem::align_of::()); // Refill the pool. let old = bk.pool.refill(alloc); // Double push the alignment block and the excessive space linearly (note that it is in // fact in the end of the pool, due to BRK _extending_ the segment). bk.double_push(alignment_block, excessive); // Free the old vector. bk.free(old); } } /// Allocate fresh memory from the global allocator. struct GlobalAllocator; impl Breaker for GlobalAllocator { #[inline] fn alloc_fresh(bk: &mut Bookkeeper, size: usize, align: usize) -> Block { /// Canonicalize the requested space. /// /// We request excessive space to the upstream allocator to avoid repeated requests and /// lock contentions. #[inline] fn canonicalize_space(min: usize) -> usize { // TODO tweak this. // To avoid having mega-allocations allocate way to much space, we // have a maximal extra space limit. if min > 8192 { min } else { // To avoid paying for short-living or little-allocating threads, we have no minimum. // Instead we multiply. min * 4 // This won't overflow due to the conditition of this branch. } } // Get the block from the global allocator. let (res, excessive) = allocator::GLOBAL_ALLOCATOR.lock() .alloc(canonicalize_space(size), align) .split(size); // Free the excessive space to the current allocator. Note that you cannot simply push // (which is the case for SBRK), due the block not necessarily being above all the other // blocks in the pool. For this reason, we let `free` handle the search and so on. bk.free(excessive); res } #[inline] fn realloc_pool(bk: &mut Bookkeeper, extra: usize) { // TODO allow BRK-free non-inplace reservations. // TODO Enable inplace reallocation in this position. // Reallocate the block pool. // Make a fresh allocation. let size = (needed + cmp::min(bk.pool.capacity(), 200 + bk.pool.capacity() / 2) // We add: + 1 // block for the alignment block. + 1 // block for the freed vector. + 1 // block for the excessive space. ) * mem::size_of::(); let (alignment_block, alloc, excessive) = brk::get(size, mem::align_of::()); // Refill the pool. let old = bk.pool.refill(alloc); // Double push the alignment block and the excessive space linearly (note that it is in // fact in the end of the pool, due to BRK _extending_ the segment). bk.double_push(alignment_block, excessive); // Free the old vector. bk.free(old); } }