diff --git a/src/allocator.rs b/src/allocator.rs index d00f048..557c6f6 100644 --- a/src/allocator.rs +++ b/src/allocator.rs @@ -4,27 +4,29 @@ use prelude::*; -use core::{ops, mem, ptr}; +use core::{mem, ops}; -use {brk, tls, sys}; +use {brk, tls}; use bookkeeper::{self, Bookkeeper, Allocator}; use sync::Mutex; +/// Alias for the wrapper type of the thread-local variable holding the local allocator. +type ThreadLocalAllocator = MoveCell LocalAllocator, LocalAllocator>>; + /// The global default allocator. // TODO remove these filthy function pointers. static GLOBAL_ALLOCATOR: Mutex GlobalAllocator, GlobalAllocator>> = Mutex::new(LazyInit::new(global_init)); tls! { /// The thread-local allocator. - static THREAD_ALLOCATOR: MoveCell LocalAllocator, LocalAllocator>> = - MoveCell::new(LazyInit::new(local_init)); + static THREAD_ALLOCATOR: ThreadLocalAllocator = MoveCell::new(LazyInit::new(local_init)); } /// Initialize the global allocator. fn global_init() -> GlobalAllocator { // The initial acquired segment. let (aligner, initial_segment, excessive) = - brk::get(bookkeeper::EXTRA_ELEMENTS * 4, mem::align_of::()); + brk::get(4 * bookkeeper::EXTRA_ELEMENTS * mem::size_of::(), mem::align_of::()); // Initialize the new allocator. let mut res = GlobalAllocator { @@ -42,21 +44,38 @@ fn global_init() -> GlobalAllocator { /// Initialize the local allocator. fn local_init() -> LocalAllocator { + /// The destructor of the local allocator. + /// + /// This will simply free everything to the global allocator. + extern fn dtor(alloc: &ThreadLocalAllocator) { + // This is important! If we simply moved out of the reference, we would end up with another + // dtor could use the value after-free. Thus we replace it by the `Unreachable` state, + // meaning that any request will result in panic. + let alloc = alloc.replace(LazyInit::unreachable()); + + // Lock the global allocator. + // TODO dumb borrowck + let mut global_alloc = GLOBAL_ALLOCATOR.lock(); + let global_alloc = global_alloc.get(); + + // TODO, we know this is sorted, so we could abuse that fact to faster insertion in the + // global allocator. + + alloc.into_inner().inner.for_each(move |block| global_alloc.free(block)); + } + // The initial acquired segment. let initial_segment = GLOBAL_ALLOCATOR .lock() .get() - .alloc(bookkeeper::EXTRA_ELEMENTS * 4, mem::align_of::()); + .alloc(4 * bookkeeper::EXTRA_ELEMENTS * mem::size_of::(), mem::align_of::()); unsafe { - // Initialize the new allocator. - let mut res = LocalAllocator { - inner: Bookkeeper::new(Vec::from_raw_parts(initial_segment, 0)), - }; - // Attach the allocator to the current thread. - res.attach(); + THREAD_ALLOCATOR.register_thread_destructor(dtor).unwrap(); - res + LocalAllocator { + inner: Bookkeeper::new(Vec::from_raw_parts(initial_segment, 0)), + } } } @@ -64,26 +83,24 @@ fn local_init() -> LocalAllocator { /// /// This is simply to avoid repeating ourself, so we let this take care of the hairy stuff. fn get_allocator T>(f: F) -> T { - /// A dummy used as placeholder for the temporary initializer. - fn dummy() -> LocalAllocator { - unreachable!(); - } - // Get the thread allocator. - let thread_alloc = THREAD_ALLOCATOR.get(); - // Just dump some placeholding initializer in the place of the TLA. - let mut thread_alloc = thread_alloc.replace(LazyInit::new(dummy)); + THREAD_ALLOCATOR.with(|thread_alloc| { + // Just dump some placeholding initializer in the place of the TLA. + let mut thread_alloc_original = thread_alloc.replace(LazyInit::unreachable()); - // Call the closure involved. - let res = f(thread_alloc.get()); + // Call the closure involved. + let res = f(thread_alloc_original.get()); - // Put back the original allocator. - THREAD_ALLOCATOR.get().replace(thread_alloc); + // Put back the original allocator. + thread_alloc.replace(thread_alloc_original); - res + res + }) } /// Derives `Deref` and `DerefMut` to the `inner` field. +/// +/// This requires importing `core::ops`. macro_rules! derive_deref { ($imp:ty, $target:ty) => { impl ops::Deref for $imp { @@ -140,33 +157,6 @@ pub struct LocalAllocator { derive_deref!(LocalAllocator, Bookkeeper); -impl LocalAllocator { - /// 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) { - extern fn dtor(ptr: *mut LocalAllocator) { - let alloc = unsafe { ptr::read(ptr) }; - - // Lock the global allocator. - // TODO dumb borrowck - let mut global_alloc = GLOBAL_ALLOCATOR.lock(); - let global_alloc = global_alloc.get(); - - // Gotta' make sure no memleaks are here. - #[cfg(feature = "debug_tools")] - alloc.assert_no_leak(); - - // TODO, we know this is sorted, so we could abuse that fact to faster insertion in the - // global allocator. - - alloc.inner.for_each(move |block| global_alloc.free(block)); - } - - sys::register_thread_destructor(self as *mut LocalAllocator, dtor).unwrap(); - } -} - impl Allocator for LocalAllocator { #[inline] fn alloc_fresh(&mut self, size: usize, align: usize) -> Block { diff --git a/src/bookkeeper.rs b/src/bookkeeper.rs index 1ef1d72..4f3cf11 100644 --- a/src/bookkeeper.rs +++ b/src/bookkeeper.rs @@ -2,9 +2,8 @@ use prelude::*; -use core::marker::PhantomData; use core::ops::Range; -use core::{ptr, cmp, mem, ops}; +use core::{ptr, mem, ops}; /// Elements required _more_ than the length as capacity. /// @@ -38,8 +37,6 @@ pub struct Bookkeeper { /// These are **not** invariants: If these assumpptions are not held, it will simply act strange /// (e.g. logic bugs), but not memory unsafety. pool: Vec, - /// Is this currently reallocating? - reallocating: bool, } impl Bookkeeper { @@ -49,11 +46,13 @@ impl Bookkeeper { debug_assert!(vec.capacity() >= EXTRA_ELEMENTS, "Not enough initial capacity of the vector."); debug_assert!(vec.is_empty(), "Initial vector isn't empty."); - Bookkeeper { + let res = Bookkeeper { pool: vec, - // Be careful with this! - .. Bookkeeper::default() - } + }; + + res.check(); + + res } /// Perform a binary search to find the appropriate place where the block can be insert or is @@ -173,16 +172,6 @@ impl Bookkeeper { log!(self.pool, "Check OK!"); } } - - /// 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()); - } } /// An allocator. @@ -588,6 +577,9 @@ pub trait Allocator: ops::DerefMut { // Logging. log!(self.pool;self.pool.len(), "Pushing {:?}.", block); + // Some assertions... + debug_assert!(&block <= self.pool.last().unwrap(), "Pushing will make the list unsorted."); + // Short-circuit in case on empty block. if !block.is_empty() { // We will try to simply merge it with the last block. diff --git a/src/cell.rs b/src/cell.rs index 7d7b50e..83115c3 100644 --- a/src/cell.rs +++ b/src/cell.rs @@ -1,5 +1,3 @@ -use prelude::*; - use core::cell::UnsafeCell; use core::mem; diff --git a/src/fail.rs b/src/fail.rs index 340f737..ae3055b 100644 --- a/src/fail.rs +++ b/src/fail.rs @@ -37,7 +37,7 @@ fn default_oom_handler() -> ! { /// The rule of thumb is that this should be called, if and only if unwinding (which allocates) /// will hit the same error. pub fn oom() -> ! { - if let Some(handler) = THREAD_OOM_HANDLER.get().replace(None) { + if let Some(handler) = THREAD_OOM_HANDLER.with(|x| x.replace(None)) { // There is a local allocator available. handler(); } else { @@ -63,10 +63,13 @@ pub fn set_oom_handler(handler: fn() -> !) { /// This might panic if a thread OOM handler already exists. #[inline] pub fn set_thread_oom_handler(handler: fn() -> !) { - let mut thread_alloc = THREAD_OOM_HANDLER.get(); - let out = thread_alloc.replace(Some(handler)); + THREAD_OOM_HANDLER.with(|thread_oom| { + // Replace it with the new handler. + let res = thread_oom.replace(Some(handler)); - debug_assert!(out.is_none()); + // Make sure that it doesn't override another handler. + debug_assert!(res.is_none()); + }); } #[cfg(test)] @@ -88,6 +91,7 @@ mod test { #[should_panic] fn test_panic_thread_oom() { fn infinite() -> ! { + #[allow(empty_loop)] loop {} } fn panic() -> ! { @@ -95,7 +99,7 @@ mod test { } set_oom_handler(infinite); - set_thread_oom_handler(infinite); + set_thread_oom_handler(panic); oom(); } } diff --git a/src/lazy_init.rs b/src/lazy_init.rs index d09b71b..05f15c5 100644 --- a/src/lazy_init.rs +++ b/src/lazy_init.rs @@ -1,7 +1,5 @@ //! `LazyStatic` like initialization. -use core::{mem, ptr, intrinsics}; - /// The initialization state enum State { /// The data is uninitialized, initialization is pending. @@ -10,9 +8,14 @@ enum State { Uninitialized(F), /// The data is initialized, and ready for use. Initialized(T), + /// The data is unreachable, panick! + Unreachable, } /// A lazily initialized container. +/// +/// This container starts out simply containing an initializer (i.e., a function to construct the +/// value in question). When the value is requested, the initializer runs. pub struct LazyInit { /// The internal state. state: State, @@ -21,7 +24,8 @@ pub struct LazyInit { impl T, T> LazyInit { /// Create a new to-be-initialized container. /// - /// The closure will be executed when initialization is required. + /// The closure will be executed when initialization is required, and is guaranteed to be + /// executed at most once. #[inline] pub const fn new(init: F) -> LazyInit { LazyInit { @@ -29,21 +33,28 @@ impl T, T> LazyInit { } } + /// Create a lazy initializer with unreachable inner data. + /// + /// When access is tried, it will simply issue a panic. This is useful for e.g. temporarily + /// getting ownership. + pub const fn unreachable() -> LazyInit { + LazyInit { + state: State::Unreachable, + } + } + /// Get a mutable reference to the inner value. /// /// If it is uninitialize, it will be initialized and then returned. #[inline] pub fn get(&mut self) -> &mut T { - let mut inner; + let inner; - let res = match self.state { - State::Initialized(ref mut x) => { - return x; - }, - State::Uninitialized(ref mut f) => { - inner = f(); - }, - }; + match self.state { + State::Initialized(ref mut x) => return x, + State::Uninitialized(ref mut f) => inner = f(), + State::Unreachable => unreachable!(), + } self.state = State::Initialized(inner); @@ -54,6 +65,19 @@ impl T, T> LazyInit { unreachable!(); } } + + /// Get the inner of the container. + /// + /// This won't mutate the container itself, since it consumes it. The initializer will (if + /// necessary) be called and the result returned. If already initialized, the inner value will + /// be moved out and returned. + pub fn into_inner(self) -> T { + match self.state { + State::Initialized(x) => x, + State::Uninitialized(mut f) => f(), + State::Unreachable => unreachable!(), + } + } } #[cfg(test)] @@ -71,8 +95,9 @@ mod test { assert_eq!(*lazy.get(), 400); } + #[test] fn test_laziness() { - let mut is_called = Cell::new(false); + let is_called = Cell::new(false); let mut lazy = LazyInit::new(|| is_called.set(true)); assert!(!is_called.get()); lazy.get(); diff --git a/src/ptr.rs b/src/ptr.rs index a5d3915..64edb39 100644 --- a/src/ptr.rs +++ b/src/ptr.rs @@ -103,10 +103,10 @@ mod test { assert_eq!(**ptr.offset(1), b'b'); } - let mut x = ['a', 'b']; + let mut y = ['a', 'b']; unsafe { - let ptr = Pointer::new(&mut x[0] as *mut char); + let ptr = Pointer::new(&mut y[0] as *mut char); assert_eq!(**ptr.clone().cast::<[char; 1]>(), ['a']); assert_eq!(**ptr.offset(1), 'b'); } diff --git a/src/symbols.rs b/src/symbols.rs index 57cb4ae..3f82446 100644 --- a/src/symbols.rs +++ b/src/symbols.rs @@ -1,5 +1,8 @@ //! Rust allocation symbols. +// TODO remove this, this is a false positive. +#![allow(private_no_mangle_fns)] + use allocator; /// Rust allocation symbol. diff --git a/src/tls.rs b/src/tls.rs index ae06396..74833f3 100644 --- a/src/tls.rs +++ b/src/tls.rs @@ -1,63 +1,47 @@ -use prelude::*; +use core::{marker, mem}; -use core::{ops, marker}; +use sys; /// A thread-local container. -pub struct Cell { +pub struct Key { /// The inner data. inner: T, } -impl Cell { - /// Create a new `Cell` wrapper. +impl Key { + /// Create a new `Key` wrapper. /// /// # Safety /// /// This is invariant-breaking (assumes thread-safety) and thus unsafe. - pub const unsafe fn new(inner: T) -> Cell { - Cell { inner: inner } + pub const unsafe fn new(inner: T) -> Key { + Key { inner: inner } } - /// Get a reference to the inner value. + /// Obtain a reference temporarily. /// - /// Due to the variable being thread-local, one should never transfer it across thread - /// boundaries. The newtype returned ensures that. - pub fn get(&'static self) -> Ref { - Ref::new(&self.inner) - } -} - -unsafe impl marker::Sync for Cell {} - -/// A reference to a thread-local variable. -/// -/// The purpose of this is to block sending it across thread boundaries. -pub struct Ref { - inner: &'static T, -} - -impl Ref { - /// Create a new thread-bounded reference. + /// Due to [the lack of thread lifetimes](https://github.com/rust-lang/rfcs/pull/1705#issuecomment-238015901), we use a closure to make sure no leakage happens. /// - /// One might wonder why this is safe, and the answer is simple: this type doesn't guarantee - /// that the internal pointer is from the current thread, it just guarantees that _future - /// access_ through this struct is done in the current thread. - pub fn new(x: &'static T) -> Ref { - Ref { - inner: x, - } + /// Having a reference newtype would be unsound, due to the ability to leak a reference to + /// another thread. + #[inline] + pub fn with(&'static self, f: F) -> R + where F: FnOnce(&T) -> R { + f(&self.inner) } -} -impl ops::Deref for Ref { - type Target = T; - - fn deref(&self) -> &T { - self.inner + /// Register a TLS destructor on the current thread. + /// + /// Note that this has to be registered for every thread, it is needed for. + // TODO make this automatic on `Drop`.a + // TODO Is this sound? + #[inline] + pub fn register_thread_destructor(&'static self, dtor: extern fn(&T)) -> Result<(), ()> { + sys::register_thread_destructor(&self.inner as *const T as *mut T, unsafe { mem::transmute(dtor) }) } } -impl !Send for Ref {} +unsafe impl marker::Sync for Key {} /// Declare a thread-local static variable. /// @@ -72,6 +56,6 @@ macro_rules! tls { (#[$($attr:meta),*] static $name:ident: $ty:ty = $val:expr;) => { $(#[$attr])* #[thread_local] - static $name: tls::Cell<$ty> = unsafe { tls::Cell::new($val) }; + static $name: tls::Key<$ty> = unsafe { tls::Key::new($val) }; } } diff --git a/tests/util/mod.rs b/tests/util/mod.rs index 2c51643..2f36979 100644 --- a/tests/util/mod.rs +++ b/tests/util/mod.rs @@ -1,7 +1,5 @@ //! Test automation. -use ralloc; - use std::{thread, mem}; /// Magic trait for boxed `FnOnce`s.