Fix bug related to strange behavior of Arc

Apparently, atomically reference counted objects would allow destructors to run after the allocator had been deinitialized. To fix this, we make sure the allocator is never disabled, but simply will fallback to the global alloctor when deinitialized. This also fixes an important bug in TLS keys (i.e. the order of the thread dtors were unknown).

We introduce a few other smaller changes as well:

- Use uniform convention for todos. Add a colon after the TODO, and start with capital letter, and end with punctuation.

- Fix broken tests which would free a larger buffer than allocated.

- Remove debug_tools in favor of planned Valgrind compatibility.

- Remove `LazyInit`'s unreachable state.

- Replace `get_allocator` with a macro (due to the lack of generic closures).

Cargo.toml

@@ -37,11 +37,10 @@ default = ["allocator", "clippy"]
 # ---
 alloc_id = []
 allocator = []
-debug_tools = []
 log = ["write", "alloc_id"]
 no_log_lock = ["log"]
 security = []
-testing = ["log", "write", "debug_tools"]
+testing = ["log", "write"]
 unsafe_no_brk_lock = []
 unsafe_no_mutex_lock = []
 write = []

README.md

@@ -80,55 +80,6 @@ fn main() {
 }
 ```
 
-### Debug check: double free
-
-Ooh, this one is a cool one. `ralloc` detects various memory bugs when compiled
-with the `debug_tools` feature. These checks include double free checks:
-
-```rust
-extern crate ralloc;
-
-fn main() {
-    // We start by allocating some stuff.
-    let a = Box::new(500u32);
-    // Then we memcpy the pointer (this is UB).
-    let b = unsafe { Box::from_raw(&*a as *mut u32) };
-    // Now both destructors are called. First a, then b, which is a double
-    // free. Luckily, `ralloc` provides a nice message for you, when in debug
-    // tools mode:
-    //    Assertion failed: Double free.
-
-    // Setting RUST_BACKTRACE allows you to get a stack backtrace, so that you
-    // can find where the double free occurs.
-}
-```
-
-### Debug check: memory leaks.
-
-TODO Temporarily disabled.
-
-`ralloc` got memleak superpowers too! Enable `debug_tools` and do:
-
-```rust
-extern crate ralloc;
-
-use std::{mem, spawn};
-
-fn main() {
-    thread::spawn(|| {
-        {
-            // We start by allocating some stuff.
-            let a = Box::new(500u32);
-            // We then leak `a`.
-            mem::forget(a);
-        }
-        // The box is now leaked, and the destructor won't be called.
-
-        // When this thread exits, the program will panic.
-    });
-}
-```
-
 ### Partial deallocation
 
 Many allocators limits deallocations to be allocated block, that is, you cannot

src/allocator.rs

@@ -10,15 +10,15 @@ use {brk, tls, sync};
 use bookkeeper::{self, Bookkeeper, Allocator};
 
 /// Alias for the wrapper type of the thread-local variable holding the local allocator.
-type ThreadLocalAllocator = MoveCell<LazyInit<fn() -> LocalAllocator, LocalAllocator>>;
+type ThreadLocalAllocator = MoveCell<Option<LazyInit<fn() -> LocalAllocator, LocalAllocator>>>;
 
 /// The global default allocator.
-// TODO remove these filthy function pointers.
+// TODO: Remove these filthy function pointers.
 static GLOBAL_ALLOCATOR: sync::Mutex<LazyInit<fn() -> GlobalAllocator, GlobalAllocator>> =
     sync::Mutex::new(LazyInit::new(global_init));
 tls! {
     /// The thread-local allocator.
-    static THREAD_ALLOCATOR: ThreadLocalAllocator = MoveCell::new(LazyInit::new(local_init));
+    static THREAD_ALLOCATOR: ThreadLocalAllocator = MoveCell::new(Some(LazyInit::new(local_init)));
 }
 
 /// Initialize the global allocator.
@@ -47,17 +47,19 @@ fn local_init() -> LocalAllocator {
     ///
     /// 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());
+        // This is important! The thread destructors guarantee no other, and thus one could use the
+        // allocator _after_ this destructor have been finished. In fact, this is a real problem,
+        // and happens when using `Arc` and terminating the main thread, for this reason we place
+        // `None` as a permanent marker indicating that the allocator is deinitialized. After such
+        // a state is in place, all allocation calls will be redirected to the global allocator,
+        // which is of course still usable at this moment.
+        let alloc = alloc.replace(None).expect("Thread-local allocator is already freed.");
 
         // 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
+        // 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));
@@ -82,21 +84,45 @@ fn local_init() -> LocalAllocator {
 
 /// Temporarily get the allocator.
 ///
-/// This is simply to avoid repeating ourself, so we let this take care of the hairy stuff.
-fn get_allocator<T, F: FnOnce(&mut LocalAllocator) -> T>(f: F) -> T {
-    // Get the thread allocator.
-    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_original.get());
-
-        // Put back the original allocator.
-        thread_alloc.replace(thread_alloc_original);
-
-        res
-    })
+/// This is simply to avoid repeating ourself, so we let this take care of the hairy stuff:
+///
+/// 1. Initialize the allocator if needed.
+/// 2. If the allocator is not yet initialized, fallback to the global allocator.
+/// 3. Unlock/move temporarily out of reference.
+///
+/// This is a macro due to the lack of generic closure, which makes it impossible to have one
+/// closure for both cases (global and local).
+// TODO: Instead of falling back to the global allocator, the thread dtor should be set such that
+// it run after the TLS keys that might be declared.
+macro_rules! get_allocator {
+    (|$v:ident| $b:expr) => {
+        // Get the thread allocator.
+        THREAD_ALLOCATOR.with(|thread_alloc| {
+            // Just dump some placeholding initializer in the place of the TLA.
+            if let Some(mut thread_alloc_original) = thread_alloc.replace(None) {
+                let res = {
+                    // Call the closure involved.
+                    let $v = thread_alloc_original.get();
+                    $b
+                };
+
+                // Put back the original allocator.
+                thread_alloc.replace(Some(thread_alloc_original));
+
+                res
+            } else {
+                // The local allocator seems to have been deinitialized, for this reason we fallback to
+                // the global allocator.
+
+                // Lock the global allocator.
+                let mut guard = GLOBAL_ALLOCATOR.lock();
+
+                // Call the block in question.
+                let $v = guard.get();
+                $b
+            }
+        })
+    }
 }
 
 /// Derives `Deref` and `DerefMut` to the `inner` field.
@@ -174,9 +200,7 @@ impl Allocator for LocalAllocator {
 /// The OOM handler handles out-of-memory conditions.
 #[inline]
 pub fn alloc(size: usize, align: usize) -> *mut u8 {
-    get_allocator(|alloc| {
-        *Pointer::from(alloc.alloc(size, align))
-    })
+    get_allocator!(|alloc| *Pointer::from(alloc.alloc(size, align)))
 }
 
 /// Free a buffer.
@@ -201,9 +225,7 @@ pub fn alloc(size: usize, align: usize) -> *mut u8 {
 /// Secondly, freeing an used buffer can introduce use-after-free.
 #[inline]
 pub unsafe fn free(ptr: *mut u8, size: usize) {
-    get_allocator(|alloc| {
-        alloc.free(Block::from_raw_parts(Pointer::new(ptr), size))
-    });
+    get_allocator!(|alloc| alloc.free(Block::from_raw_parts(Pointer::new(ptr), size)))
 }
 
 /// Reallocate memory.
@@ -226,7 +248,7 @@ pub unsafe fn free(ptr: *mut u8, size: usize) {
 /// this is marked unsafe.
 #[inline]
 pub unsafe fn realloc(ptr: *mut u8, old_size: usize, size: usize, align: usize) -> *mut u8 {
-    get_allocator(|alloc| {
+    get_allocator!(|alloc| {
         *Pointer::from(alloc.realloc(
             Block::from_raw_parts(Pointer::new(ptr), old_size),
             size,
@@ -246,7 +268,7 @@ pub unsafe fn realloc(ptr: *mut u8, old_size: usize, size: usize, align: usize)
 /// Due to being able to shrink (and thus free) the buffer, this is marked unsafe.
 #[inline]
 pub unsafe fn realloc_inplace(ptr: *mut u8, old_size: usize, size: usize) -> Result<(), ()> {
-    get_allocator(|alloc| {
+    get_allocator!(|alloc| {
         if alloc.realloc_inplace(
             Block::from_raw_parts(Pointer::new(ptr), old_size),
             size

src/block.rs

@@ -3,7 +3,7 @@
 //! Blocks are the main unit for the memory bookkeeping. A block is a simple construct with a
 //! `Pointer` pointer and a size. Occupied (non-free) blocks are represented by a zero-sized block.
 
-// TODO check the allow(cast_possible_wrap)s again.
+// TODO: Check the allow(cast_possible_wrap)s again.
 
 use prelude::*;
 

src/bookkeeper.rs

@@ -56,7 +56,7 @@ pub struct Bookkeeper {
     ///
     /// This is used to avoid unbounded metacircular reallocation (reservation).
     ///
-    // TODO find a replacement for this "hack".
+    // TODO: Find a replacement for this "hack".
     reserving: bool,
     /// The allocator ID.
     ///
@@ -72,7 +72,7 @@ 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.");
 
-        // TODO, when added use expr field attributes
+        // TODO: When added use expr field attributes.
         #[cfg(feature = "alloc_id")]
         let res = Bookkeeper {
             pool: vec,
@@ -98,7 +98,6 @@ impl Bookkeeper {
     /// 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, "Searching (exact) for {:?}.", block);
 
@@ -121,7 +120,6 @@ impl Bookkeeper {
     /// It is guaranteed that no block left to the returned value, satisfy the above condition.
     #[inline]
     fn find_bound(&mut self, block: &Block) -> Range<usize> {
-        // TODO optimize this function.
         // Logging.
         log!(self, "Searching (bounds) for {:?}.", block);
 
@@ -288,8 +286,6 @@ pub trait Allocator: ops::DerefMut<Target = Bookkeeper> {
     ///
     /// A block representing the marked area is then returned.
     fn alloc(&mut self, size: usize, align: usize) -> Block {
-        // TODO: scan more intelligently.
-
         // Logging.
         log!(self, "Allocating {} bytes with alignment {}.", size, align);
 
@@ -546,7 +542,6 @@ pub trait Allocator: ops::DerefMut<Target = Bookkeeper> {
                 // Run a consistency check.
                 self.check();
 
-                // TODO, drop excessive space
                 return Ok(res);
             }
         }

src/brk.rs

@@ -18,7 +18,7 @@ use core::cmp;
 /// The return value is always greater than or equals to the argument.
 #[inline]
 fn canonicalize_space(min: usize) -> usize {
-    // TODO tweak this.
+    // TODO: Tweak this.
     /// The BRK multiplier.
     ///
     /// The factor determining the linear dependence between the minimum segment, and the acquired

src/cell.rs

@@ -7,7 +7,7 @@ use core::mem;
 ///
 /// This allows you to take ownership and replace the internal data with a new value. The
 /// functionality is similar to the one provided by [RFC #1659](https://github.com/rust-lang/rfcs/pull/1659).
-// TODO use that rfc ^
+// TODO: Use the features provided by the RFC.
 pub struct MoveCell<T> {
     /// The inner data.
     inner: UnsafeCell<T>,

src/fail.rs

@@ -68,7 +68,7 @@ pub fn set_thread_oom_handler(handler: fn() -> !) {
         let res = thread_oom.replace(Some(handler));
 
         // Make sure that it doesn't override another handler.
-        // TODO Make this a warning
+        // TODO: Make this a warning.
         debug_assert!(res.is_none(), "Overriding the old handler. Is this intentional?");
     });
 }

src/lazy_init.rs

@@ -8,8 +8,6 @@ enum State<F, T> {
     Uninitialized(F),
     /// The data is initialized, and ready for use.
     Initialized(T),
-    /// The data is unreachable, panick!
-    Unreachable,
 }
 
 /// A lazily initialized container.
@@ -33,16 +31,6 @@ impl<F: FnMut() -> T, T> LazyInit<F, T> {
         }
     }
 
-    /// 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<F, T> {
-        LazyInit {
-            state: State::Unreachable,
-        }
-    }
-
     /// Get a mutable reference to the inner value.
     ///
     /// If it is uninitialize, it will be initialized and then returned.
@@ -53,7 +41,6 @@ impl<F: FnMut() -> T, T> LazyInit<F, T> {
         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);
@@ -61,7 +48,7 @@ impl<F: FnMut() -> T, T> LazyInit<F, T> {
         if let State::Initialized(ref mut x) = self.state {
             x
         } else {
-            // TODO find a better way.
+            // TODO: Find a better way to deal with this case.
             unreachable!();
         }
     }
@@ -75,7 +62,6 @@ impl<F: FnMut() -> T, T> LazyInit<F, T> {
         match self.state {
             State::Initialized(x) => x,
             State::Uninitialized(mut f) => f(),
-            State::Unreachable => unreachable!(),
         }
     }
 }
@@ -103,15 +89,4 @@ mod test {
         lazy.get();
         assert!(is_called.get());
     }
-
-    #[test]
-    #[should_panic]
-    #[allow(unused_assignments)]
-    fn test_unreachable() {
-        let mut a = LazyInit::new(|| 2);
-
-        a = LazyInit::unreachable();
-
-        a.get();
-    }
 }

src/log.rs

@@ -60,7 +60,7 @@ pub mod internal {
         /// formatter if the underlying condition is true.
         ///
         /// For example, a plain position cursor will write `"|"` when `n == self.pos`.
-        // TODO use an iterator instead.
+        // TODO: Use an iterator instead.
         fn at(&self, f: &mut fmt::Formatter, n: usize) -> fmt::Result;
 
         /// The after hook.
@@ -186,7 +186,7 @@ pub mod internal {
 
     impl<'a, T: Cursor> fmt::Debug for BlockLogger<'a, T> {
         fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-            // TODO handle alignment etc.
+            // TODO: Handle alignment etc.
 
             for (n, i) in self.blocks.iter().enumerate() {
                 self.cur.at(f, n)?;

src/prelude.rs

@@ -1,6 +1,6 @@
 //! Frequently used imports.
 
-// TODO remove all this?
+// TODO: Reconsider this. Is this an anti-pattern?
 
 pub use block::Block;
 pub use cell::MoveCell;

src/symbols.rs

@@ -1,6 +1,6 @@
 //! Rust allocation symbols.
 
-// TODO remove this, this is a false positive.
+// TODO: Remove this, this is a false positive.
 #![allow(private_no_mangle_fns)]
 
 use allocator;

src/sys.rs

@@ -45,7 +45,7 @@ pub fn yield_now() {
 ///
 /// The argument to the destructor is a pointer to the so-called "load", which is the data
 /// shipped with the destructor.
-// TODO I haven't figured out a safe general solution yet. Libstd relies on devirtualization,
+// TODO: I haven't figured out a safe general solution yet. Libstd relies on devirtualization,
 // which, when missed, can make it quite expensive.
 pub fn register_thread_destructor<T>(load: *mut T, dtor: extern fn(*mut T)) -> Result<(), ()> {
     // Check if thread dtors are supported.
@@ -64,7 +64,7 @@ pub fn register_thread_destructor<T>(load: *mut T, dtor: extern fn(*mut T)) -> R
 /// Write text to the log.
 ///
 /// The log target is defined by the `shim` crate.
-// TODO find a better way to silence the warning than this attribute
+// TODO: Find a better way to silence the warning than this attribute.
 #[allow(dead_code)]
 pub fn log(s: &str) -> Result<(), ()> {
     if shim::log(s) == -1 { Err(()) } else { Ok(()) }

src/tls.rs

@@ -37,8 +37,7 @@ impl<T: 'static> Key<T> {
     /// 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?
+    // TODO: Make this automatic on `Drop`.
     #[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) })

src/vec.rs

@@ -131,7 +131,7 @@ impl<T: Leak> Vec<T> {
     }
 }
 
-// TODO remove this in favour of `derive` when rust-lang/rust#35263 is fixed.
+// TODO: Remove this in favour of `derive` when rust-lang/rust#35263 is fixed.
 impl<T: Leak> Default for Vec<T> {
     fn default() -> Vec<T> {
         Vec {

tests/arc.rs

@@ -5,7 +5,8 @@ extern crate ralloc;
 use std::sync::Arc;
 use std::thread;
 
-fn main() {
+#[test]
+fn test_arc() {
     let numbers: Vec<_> = (0..100).collect();
     let shared_numbers = Arc::new(numbers);
 

tests/partial_free.rs

@@ -16,7 +16,7 @@ fn partial_free() {
             });
 
             util::acid(|| {
-                ralloc::free(buf.offset(8), 75);
+                ralloc::free(buf.offset(8), 55);
                 *buf = 5;
             });