std::rt: Check exchange count on exit

2 files changed, 37 insertions, 6 deletions

diff --git a/src/libstd/rt/global_heap.rs b/src/libstd/rt/global_heap.rs
index ce7ff87b445..9f273fc8e65 100644
--- a/src/libstd/rt/global_heap.rs
+++ b/src/libstd/rt/global_heap.rs
@@ -14,7 +14,7 @@ use c_malloc = libc::malloc;
 use c_free = libc::free;
 use managed::raw::{BoxHeaderRepr, BoxRepr};
 use cast::transmute;
-use unstable::intrinsics::{atomic_xadd,atomic_xsub};
+use unstable::intrinsics::{atomic_xadd,atomic_xsub, atomic_load};
 use ptr::null;
 use intrinsic::TyDesc;
 
@@ -34,8 +34,7 @@ pub unsafe fn malloc(td: *TypeDesc, size: uint) -> *c_void {
     box.header.prev = null();
     box.header.next = null();
 
-    let exchange_count = &mut *exchange_count_ptr();
-    atomic_xadd(exchange_count, 1);
+    inc_count();
 
     return transmute(box);
 }
@@ -48,21 +47,47 @@ pub unsafe fn malloc_raw(size: uint) -> *c_void {
     if p.is_null() {
         fail!("Failure in malloc_raw: result ptr is null");
     }
+    inc_count();
     p
 }
 
 pub unsafe fn free(ptr: *c_void) {
-    let exchange_count = &mut *exchange_count_ptr();
-    atomic_xsub(exchange_count, 1);
-
     assert!(ptr.is_not_null());
+    dec_count();
     c_free(ptr);
 }
 ///Thin wrapper around libc::free, as with exchange_alloc::malloc_raw
 pub unsafe fn free_raw(ptr: *c_void) {
+    assert!(ptr.is_not_null());
+    dec_count();
     c_free(ptr);
 }
 
+fn inc_count() {
+    unsafe {
+        let exchange_count = &mut *exchange_count_ptr();
+        atomic_xadd(exchange_count, 1);
+    }
+}
+
+fn dec_count() {
+    unsafe {
+        let exchange_count = &mut *exchange_count_ptr();
+        atomic_xsub(exchange_count, 1);
+    }
+}
+
+pub fn cleanup() {
+    unsafe {
+        let count_ptr = exchange_count_ptr();
+        let allocations = atomic_load(&*count_ptr);
+        if allocations != 0 {
+            abort!("exchange heap not empty on exit\
+                    %i dangling allocations", allocations);
+        }
+    }
+}
+
 fn get_box_size(body_size: uint, body_align: uint) -> uint {
     let header_size = size_of::<BoxHeaderRepr>();
     // FIXME (#2699): This alignment calculation is suspicious. Is it right?
diff --git a/src/libstd/rt/mod.rs b/src/libstd/rt/mod.rs
index 10b5c78f99e..0a269aa8767 100644
--- a/src/libstd/rt/mod.rs
+++ b/src/libstd/rt/mod.rs
@@ -176,6 +176,8 @@ pub fn start(_argc: int, _argv: **u8, crate_map: *u8, main: ~fn()) -> int {
     sched.enqueue_task(main_task);
     sched.run();
 
+    cleanup();
+
     return 0;
 }
 
@@ -185,6 +187,10 @@ pub fn init(crate_map: *u8) {
     logging::init(crate_map);
 }
 
+pub fn cleanup() {
+    global_heap::cleanup();
+}
+
 /// Possible contexts in which Rust code may be executing.
 /// Different runtime services are available depending on context.
 /// Mostly used for determining if we're using the new scheduler