auto merge of #11187 : alexcrichton/rust/once, r=brson

Rationale can be found in the first commit, but this is basically the same thing as `pthread_once`

2 files changed, 150 insertions, 13 deletions

diff --git a/src/libstd/rt/local_ptr.rs b/src/libstd/rt/local_ptr.rs
index 42cce272e44..f13691a7bfe 100644
--- a/src/libstd/rt/local_ptr.rs
+++ b/src/libstd/rt/local_ptr.rs
@@ -160,30 +160,20 @@ pub mod native {
     use option::{Option, Some, None};
     use ptr;
     use tls = rt::thread_local_storage;
-    use unstable::mutex::{Mutex, MUTEX_INIT};
 
-    static mut LOCK: Mutex = MUTEX_INIT;
-    static mut INITIALIZED: bool = false;
     static mut RT_TLS_KEY: tls::Key = -1;
 
     /// Initialize the TLS key. Other ops will fail if this isn't executed
     /// first.
     pub fn init() {
         unsafe {
-            LOCK.lock();
-            if !INITIALIZED {
-                tls::create(&mut RT_TLS_KEY);
-                INITIALIZED = true;
-            }
-            LOCK.unlock();
+            tls::create(&mut RT_TLS_KEY);
         }
     }
 
     pub unsafe fn cleanup() {
-        rtassert!(INITIALIZED);
+        rtassert!(RT_TLS_KEY != -1);
         tls::destroy(RT_TLS_KEY);
-        LOCK.destroy();
-        INITIALIZED = false;
     }
 
     /// Give a pointer to thread-local storage.
diff --git a/src/libstd/unstable/mutex.rs b/src/libstd/unstable/mutex.rs
index 7ab8a83d04b..36682b85418 100644
--- a/src/libstd/unstable/mutex.rs
+++ b/src/libstd/unstable/mutex.rs
@@ -47,6 +47,7 @@
 
 #[allow(non_camel_case_types)];
 
+use int;
 use libc::c_void;
 use sync::atomics;
 
@@ -315,10 +316,156 @@ mod imp {
     }
 }
 
+/// A type which can be used to run a one-time global initialization. This type
+/// is *unsafe* to use because it is built on top of the `Mutex` in this module.
+/// It does not know whether the currently running task is in a green or native
+/// context, and a blocking mutex should *not* be used under normal
+/// circumstances on a green task.
+///
+/// Despite its unsafety, it is often useful to have a one-time initialization
+/// routine run for FFI bindings or related external functionality. This type
+/// can only be statically constructed with the `ONCE_INIT` value.
+///
+/// # Example
+///
+/// ```rust
+/// use std::unstable::mutex::{Once, ONCE_INIT};
+///
+/// static mut START: Once = ONCE_INIT;
+/// unsafe {
+///     START.doit(|| {
+///         // run initialization here
+///     });
+/// }
+/// ```
+pub struct Once {
+    priv mutex: Mutex,
+    priv cnt: atomics::AtomicInt,
+    priv lock_cnt: atomics::AtomicInt,
+}
+
+/// Initialization value for static `Once` values.
+pub static ONCE_INIT: Once = Once {
+    mutex: MUTEX_INIT,
+    cnt: atomics::INIT_ATOMIC_INT,
+    lock_cnt: atomics::INIT_ATOMIC_INT,
+};
+
+impl Once {
+    /// Perform an initialization routine once and only once. The given closure
+    /// will be executed if this is the first time `doit` has been called, and
+    /// otherwise the routine will *not* be invoked.
+    ///
+    /// This method will block the calling *os thread* if another initialization
+    /// routine is currently running.
+    ///
+    /// When this function returns, it is guaranteed that some initialization
+    /// has run and completed (it may not be the closure specified).
+    pub fn doit(&mut self, f: ||) {
+        // Implementation-wise, this would seem like a fairly trivial primitive.
+        // The stickler part is where our mutexes currently require an
+        // allocation, and usage of a `Once` should't leak this allocation.
+        //
+        // This means that there must be a deterministic destroyer of the mutex
+        // contained within (because it's not needed after the initialization
+        // has run).
+        //
+        // The general scheme here is to gate all future threads once
+        // initialization has completed with a "very negative" count, and to
+        // allow through threads to lock the mutex if they see a non negative
+        // count. For all threads grabbing the mutex, exactly one of them should
+        // be responsible for unlocking the mutex, and this should only be done
+        // once everyone else is done with the mutex.
+        //
+        // This atomicity is achieved by swapping a very negative value into the
+        // shared count when the initialization routine has completed. This will
+        // read the number of threads which will at some point attempt to
+        // acquire the mutex. This count is then squirreled away in a separate
+        // variable, and the last person on the way out of the mutex is then
+        // responsible for destroying the mutex.
+        //
+        // It is crucial that the negative value is swapped in *after* the
+        // initialization routine has completed because otherwise new threads
+        // calling `doit` will return immediately before the initialization has
+        // completed.
+
+        let prev = self.cnt.fetch_add(1, atomics::SeqCst);
+        if prev < 0 {
+            // Make sure we never overflow, we'll never have int::min_value
+            // simultaneous calls to `doit` to make this value go back to 0
+            self.cnt.store(int::min_value, atomics::SeqCst);
+            return
+        }
+
+        // If the count is negative, then someone else finished the job,
+        // otherwise we run the job and record how many people will try to grab
+        // this lock
+        unsafe { self.mutex.lock() }
+        if self.cnt.load(atomics::SeqCst) > 0 {
+            f();
+            let prev = self.cnt.swap(int::min_value, atomics::SeqCst);
+            self.lock_cnt.store(prev, atomics::SeqCst);
+        }
+        unsafe { self.mutex.unlock() }
+
+        // Last one out cleans up after everyone else, no leaks!
+        if self.lock_cnt.fetch_add(-1, atomics::SeqCst) == 1 {
+            unsafe { self.mutex.destroy() }
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
-    use super::{Mutex, MUTEX_INIT};
+    use prelude::*;
+
     use rt::thread::Thread;
+    use super::{ONCE_INIT, Once, Mutex, MUTEX_INIT};
+    use task;
+
+    #[test]
+    fn smoke_once() {
+        static mut o: Once = ONCE_INIT;
+        let mut a = 0;
+        unsafe { o.doit(|| a += 1); }
+        assert_eq!(a, 1);
+        unsafe { o.doit(|| a += 1); }
+        assert_eq!(a, 1);
+    }
+
+    #[test]
+    fn stampede_once() {
+        static mut o: Once = ONCE_INIT;
+        static mut run: bool = false;
+
+        let (p, c) = SharedChan::new();
+        for _ in range(0, 10) {
+            let c = c.clone();
+            do spawn {
+                for _ in range(0, 4) { task::deschedule() }
+                unsafe {
+                    o.doit(|| {
+                        assert!(!run);
+                        run = true;
+                    });
+                    assert!(run);
+                }
+                c.send(());
+            }
+        }
+
+        unsafe {
+            o.doit(|| {
+                assert!(!run);
+                run = true;
+            });
+            assert!(run);
+        }
+
+        for _ in range(0, 10) {
+            p.recv();
+        }
+    }
 
     #[test]
     fn somke_lock() {