libstd: add an RAII utility for sys_common::mutex::Mutex

Signed-off-by: NODA, Kai <nodakai@gmail.com>

3 files changed, 39 insertions, 17 deletions

diff --git a/src/libstd/sys_common/at_exit_imp.rs b/src/libstd/sys_common/at_exit_imp.rs
index 26da51c9825..d268d9ad6f9 100644
--- a/src/libstd/sys_common/at_exit_imp.rs
+++ b/src/libstd/sys_common/at_exit_imp.rs
@@ -14,6 +14,7 @@
 
 use boxed::FnBox;
 use ptr;
+use mem;
 use sys_common::mutex::Mutex;
 
 type Queue = Vec<Box<FnBox()>>;
@@ -25,6 +26,8 @@ type Queue = Vec<Box<FnBox()>>;
 static LOCK: Mutex = Mutex::new();
 static mut QUEUE: *mut Queue = ptr::null_mut();
 
+const DONE: *mut Queue = 1_usize as *mut _;
+
 // The maximum number of times the cleanup routines will be run. While running
 // the at_exit closures new ones may be registered, and this count is the number
 // of times the new closures will be allowed to register successfully. After
@@ -35,7 +38,7 @@ unsafe fn init() -> bool {
     if QUEUE.is_null() {
         let state: Box<Queue> = box Vec::new();
         QUEUE = Box::into_raw(state);
-    } else if QUEUE as usize == 1 {
+    } else if QUEUE == DONE {
         // can't re-init after a cleanup
         return false
     }
@@ -44,18 +47,18 @@ unsafe fn init() -> bool {
 }
 
 pub fn cleanup() {
-    for i in 0..ITERS {
+    for i in 1..=ITERS {
         unsafe {
-            LOCK.lock();
-            let queue = QUEUE;
-            QUEUE = if i == ITERS - 1 {1} else {0} as *mut _;
-            LOCK.unlock();
+            let queue = {
+                let _guard = LOCK.lock();
+                mem::replace(&mut QUEUE, if i == ITERS { DONE } else { ptr::null_mut() })
+            };
 
             // make sure we're not recursively cleaning up
-            assert!(queue as usize != 1);
+            assert!(queue != DONE);
 
             // If we never called init, not need to cleanup!
-            if queue as usize != 0 {
+            if !queue.is_null() {
                 let queue: Box<Queue> = Box::from_raw(queue);
                 for to_run in *queue {
                     to_run();
@@ -66,15 +69,13 @@ pub fn cleanup() {
 }
 
 pub fn push(f: Box<FnBox()>) -> bool {
-    let mut ret = true;
     unsafe {
-        LOCK.lock();
+        let _guard = LOCK.lock();
         if init() {
             (*QUEUE).push(f);
+            true
         } else {
-            ret = false;
+            false
         }
-        LOCK.unlock();
     }
-    ret
 }
diff --git a/src/libstd/sys_common/mutex.rs b/src/libstd/sys_common/mutex.rs
index d1a738770d3..608355b7d70 100644
--- a/src/libstd/sys_common/mutex.rs
+++ b/src/libstd/sys_common/mutex.rs
@@ -37,7 +37,15 @@ impl Mutex {
     /// Behavior is undefined if the mutex has been moved between this and any
     /// previous function call.
     #[inline]
-    pub unsafe fn lock(&self) { self.0.lock() }
+    pub unsafe fn raw_lock(&self) { self.0.lock() }
+
+    /// Calls raw_lock() and then returns an RAII guard to guarantee the mutex
+    /// will be unlocked.
+    #[inline]
+    pub unsafe fn lock(&self) -> MutexGuard {
+        self.raw_lock();
+        MutexGuard(&self.0)
+    }
 
     /// Attempts to lock the mutex without blocking, returning whether it was
     /// successfully acquired or not.
@@ -51,8 +59,11 @@ impl Mutex {
     ///
     /// Behavior is undefined if the current thread does not actually hold the
     /// mutex.
+    ///
+    /// Consider switching from the pair of raw_lock() and raw_unlock() to
+    /// lock() whenever possible.
     #[inline]
-    pub unsafe fn unlock(&self) { self.0.unlock() }
+    pub unsafe fn raw_unlock(&self) { self.0.unlock() }
 
     /// Deallocates all resources associated with this mutex.
     ///
@@ -64,3 +75,14 @@ impl Mutex {
 
 // not meant to be exported to the outside world, just the containing module
 pub fn raw(mutex: &Mutex) -> &imp::Mutex { &mutex.0 }
+
+#[must_use]
+/// A simple RAII utility for the above Mutex without the poisoning semantics.
+pub struct MutexGuard<'a>(&'a imp::Mutex);
+
+impl<'a> Drop for MutexGuard<'a> {
+    #[inline]
+    fn drop(&mut self) {
+        unsafe { self.0.unlock(); }
+    }
+}
diff --git a/src/libstd/sys_common/thread_local.rs b/src/libstd/sys_common/thread_local.rs
index d0d6224de0a..75f6b9ac7fd 100644
--- a/src/libstd/sys_common/thread_local.rs
+++ b/src/libstd/sys_common/thread_local.rs
@@ -162,13 +162,12 @@ impl StaticKey {
         // we just simplify the whole branch.
         if imp::requires_synchronized_create() {
             static INIT_LOCK: Mutex = Mutex::new();
-            INIT_LOCK.lock();
+            let _guard = INIT_LOCK.lock();
             let mut key = self.key.load(Ordering::SeqCst);
             if key == 0 {
                 key = imp::create(self.dtor) as usize;
                 self.key.store(key, Ordering::SeqCst);
             }
-            INIT_LOCK.unlock();
             rtassert!(key != 0);
             return key
         }