Rewrite io::lazy

rust-lang · pitdicker · Feb 21, 2019 · Feb 22, 2019 · Mar 4, 2019 · Feb 28, 2019
commit 271d990bf602b629007a8b8eb659197079358219
diff --git a/src/libstd/io/lazy.rs b/src/libstd/io/lazy.rs
@@ -1,64 +1,77 @@
-use crate::cell::Cell;
-use crate::ptr;
-use crate::sync::Arc;
+use crate::cell::UnsafeCell;
+use crate::sync::atomic::AtomicUsize;
+use crate::sync::atomic::Ordering;
 use crate::sys_common;
 use crate::sys_common::mutex::Mutex;
 
+/// Helper for lazy initialization of a static, with a destructor that runs when the main (Rust)
+/// thread exits.
+///
+/// Currently used only inside the standard library, by the stdio types.
+///
+/// # Safety
+/// - `UnsafeCell`: We only create a mutable reference during initialization and during the shutdown
+///   phase. At both times there can't exist any other references.
+/// - Destruction. The `Drop` implementation of `T` should not access references to anything except
+///   itself, they are not guaranteed to exist. It should also not rely on other machinery of the
+///   standard library to be available.
+/// - Initialization. The `init` function for `get` should not call `get` itself,  to prevent
+///   infinite recursion and acquiring the guard mutex reentrantly.
+/// - We use the `Mutex` from `sys::common` because it has a `const` constructor. It currently has
+///   UB when acquired reentrantly without calling `init`.
 pub struct Lazy<T> {
-    // We never call `lock.init()`, so it is UB to attempt to acquire this mutex reentrantly!
-    lock: Mutex,
-    ptr: Cell<*mut Arc<T>>,
+    guard: Mutex, // Only used to protect initialization.
+    status: AtomicUsize,
+    data: UnsafeCell<Option<T>>,
 }
 
-#[inline]
-const fn done<T>() -> *mut Arc<T> { 1_usize as *mut _ }
-
 unsafe impl<T> Sync for Lazy<T> {}
 
+const UNINITIALIZED: usize = 0;
+const SHUTDOWN: usize = 1;
+const AVAILABLE: usize = 2;
+
 impl<T> Lazy<T> {
     pub const fn new() -> Lazy<T> {
         Lazy {
-            lock: Mutex::new(),
-            ptr: Cell::new(ptr::null_mut()),
+            guard: Mutex::new(),
+            status: AtomicUsize::new(UNINITIALIZED),
+            data: UnsafeCell::new(None),
         }
     }
 }
 
 impl<T: Send + Sync + 'static> Lazy<T> {
-    /// Safety: `init` must not call `get` on the variable that is being
-    /// initialized.
-    pub unsafe fn get(&'static self, init: fn() -> Arc<T>) -> Option<Arc<T>> {
-        let _guard = self.lock.lock();
-        let ptr = self.ptr.get();
-        if ptr.is_null() {
-            Some(self.init(init))
-        } else if ptr == done() {
-            None
-        } else {
-            Some((*ptr).clone())
-        }
-    }
+    pub unsafe fn get(&'static self, init: fn() -> T) -> Option<&T> {
+        match self.status.load(Ordering::Acquire) {
+            UNINITIALIZED => {
+                let _guard = self.guard.lock();
+                // Double-check to make sure this `Lazy` didn't get initialized by another
+                // thread in the small window before we acquired the mutex.
+                if self.status.load(Ordering::Relaxed) != UNINITIALIZED {
+                    return self.get(init);
+                }
+
+                // Register an `at_exit` handler that drops `data` when the main thread exits.
+                let registered = sys_common::at_exit(move || {
+                    *self.data.get() = None; // `T` gets dropped here
+                    self.status.store(SHUTDOWN, Ordering::Release);
+                });
+                if registered.is_err() {
+                    // Registering the handler will only fail if we are already in the shutdown
+                    // phase. In that case don't attempt to initialize.
+                    self.status.store(SHUTDOWN, Ordering::Release);
+                    return None;
+                }
+
+                // Run the initializer of `T`.
+                *self.data.get() = Some(init());
+                self.status.store(AVAILABLE, Ordering::Release);
 
-    // Must only be called with `lock` held
-    unsafe fn init(&'static self, init: fn() -> Arc<T>) -> Arc<T> {
-        // If we successfully register an at exit handler, then we cache the
-        // `Arc` allocation in our own internal box (it will get deallocated by
-        // the at exit handler). Otherwise we just return the freshly allocated
-        // `Arc`.
-        let registered = sys_common::at_exit(move || {
-            let ptr = {
-                let _guard = self.lock.lock();
-                self.ptr.replace(done())
-            };
-            drop(Box::from_raw(ptr))
-        });
-        // This could reentrantly call `init` again, which is a problem
-        // because our `lock` allows reentrancy!
-        // That's why `get` is unsafe and requires the caller to ensure no reentrancy happens.
-        let ret = init();
-        if registered.is_ok() {
-            self.ptr.set(Box::into_raw(Box::new(ret.clone())));
+                (*self.data.get()).as_ref()
+            },
+            SHUTDOWN => None,
+            _ => (*self.data.get()).as_ref(),
         }
-        ret
     }
 }
diff --git a/src/libstd/io/stdio.rs b/src/libstd/io/stdio.rs
@@ -4,7 +4,7 @@ use crate::cell::RefCell;
 use crate::fmt;
 use crate::io::lazy::Lazy;
 use crate::io::{self, Initializer, BufReader, LineWriter};
-use crate::sync::{Arc, Mutex, MutexGuard};
+use crate::sync::{Mutex, MutexGuard};
 use crate::sys::stdio;
 use crate::sys_common::remutex::{ReentrantMutex, ReentrantMutexGuard};
 use crate::thread::LocalKey;
@@ -138,7 +138,7 @@ fn handle_ebadf<T>(r: io::Result<T>, default: T) -> io::Result<T> {
 /// an error.
 #[stable(feature = "rust1", since = "1.0.0")]
 pub struct Stdin {
-    inner: Arc<Mutex<BufReader<Maybe<StdinRaw>>>>,
+    inner: &'static Mutex<BufReader<Maybe<StdinRaw>>>,
 }
 
 /// A locked reference to the `Stdin` handle.
@@ -202,21 +202,19 @@ pub struct StdinLock<'a> {
 /// ```
 #[stable(feature = "rust1", since = "1.0.0")]
 pub fn stdin() -> Stdin {
-    static INSTANCE: Lazy<Mutex<BufReader<Maybe<StdinRaw>>>> = Lazy::new();
-    return Stdin {
-        inner: unsafe {
-            INSTANCE.get(stdin_init).expect("cannot access stdin during shutdown")
-        },
-    };
-
-    fn stdin_init() -> Arc<Mutex<BufReader<Maybe<StdinRaw>>>> {
-        // This must not reentrantly access `INSTANCE`
+    static STDIN: Lazy<Mutex<BufReader<Maybe<StdinRaw>>>> = Lazy::new();
+    fn stdin_init() -> Mutex<BufReader<Maybe<StdinRaw>>> {
         let stdin = match stdin_raw() {
             Ok(stdin) => Maybe::Real(stdin),
             _ => Maybe::Fake
         };
+        Mutex::new(BufReader::with_capacity(stdio::STDIN_BUF_SIZE, stdin))
+    }
 
-        Arc::new(Mutex::new(BufReader::with_capacity(stdio::STDIN_BUF_SIZE, stdin)))
+    Stdin {
+        inner: unsafe {
+            STDIN.get(stdin_init).expect("cannot access stdin during shutdown")
+        },
     }
 }
 
@@ -355,7 +353,7 @@ pub struct Stdout {
     // FIXME: this should be LineWriter or BufWriter depending on the state of
     //        stdout (tty or not). Note that if this is not line buffered it
     //        should also flush-on-panic or some form of flush-on-abort.
-    inner: Arc<ReentrantMutex<RefCell<LineWriter<Maybe<StdoutRaw>>>>>,
+    inner: &'static ReentrantMutex<RefCell<LineWriter<Maybe<StdoutRaw>>>>,
 }
 
 /// A locked reference to the `Stdout` handle.
@@ -418,20 +416,19 @@ pub struct StdoutLock<'a> {
 /// ```
 #[stable(feature = "rust1", since = "1.0.0")]
 pub fn stdout() -> Stdout {
-    static INSTANCE: Lazy<ReentrantMutex<RefCell<LineWriter<Maybe<StdoutRaw>>>>> = Lazy::new();
-    return Stdout {
-        inner: unsafe {
-            INSTANCE.get(stdout_init).expect("cannot access stdout during shutdown")
-        },
-    };
-
-    fn stdout_init() -> Arc<ReentrantMutex<RefCell<LineWriter<Maybe<StdoutRaw>>>>> {
-        // This must not reentrantly access `INSTANCE`
+    static STDOUT: Lazy<ReentrantMutex<RefCell<LineWriter<Maybe<StdoutRaw>>>>> = Lazy::new();
+    fn stdout_init() -> ReentrantMutex<RefCell<LineWriter<Maybe<StdoutRaw>>>> {
         let stdout = match stdout_raw() {
             Ok(stdout) => Maybe::Real(stdout),
             _ => Maybe::Fake,
         };
-        Arc::new(ReentrantMutex::new(RefCell::new(LineWriter::new(stdout))))
+        ReentrantMutex::new(RefCell::new(LineWriter::new(stdout)))
+    }
+
+    Stdout {
+        inner: unsafe {
+            STDOUT.get(stdout_init).expect("cannot access stdout during shutdown")
+        },
     }
 }
 
@@ -513,7 +510,7 @@ impl fmt::Debug for StdoutLock<'_> {
 /// an error.
 #[stable(feature = "rust1", since = "1.0.0")]
 pub struct Stderr {
-    inner: Arc<ReentrantMutex<RefCell<Maybe<StderrRaw>>>>,
+    inner: &'static ReentrantMutex<RefCell<Maybe<StderrRaw>>>,
 }
 
 /// A locked reference to the `Stderr` handle.
@@ -571,20 +568,19 @@ pub struct StderrLock<'a> {
 /// ```
 #[stable(feature = "rust1", since = "1.0.0")]
 pub fn stderr() -> Stderr {
-    static INSTANCE: Lazy<ReentrantMutex<RefCell<Maybe<StderrRaw>>>> = Lazy::new();
-    return Stderr {
-        inner: unsafe {
-            INSTANCE.get(stderr_init).expect("cannot access stderr during shutdown")
-        },
-    };
-
-    fn stderr_init() -> Arc<ReentrantMutex<RefCell<Maybe<StderrRaw>>>> {
-        // This must not reentrantly access `INSTANCE`
+    static STDERR: Lazy<ReentrantMutex<RefCell<Maybe<StderrRaw>>>> = Lazy::new();
+    fn stderr_init() -> ReentrantMutex<RefCell<Maybe<StderrRaw>>> {
         let stderr = match stderr_raw() {
             Ok(stderr) => Maybe::Real(stderr),
             _ => Maybe::Fake,
         };
-        Arc::new(ReentrantMutex::new(RefCell::new(stderr)))
+        ReentrantMutex::new(RefCell::new(stderr))
+    }
+
+    Stderr {
+        inner: unsafe {
+            STDERR.get(stderr_init).expect("cannot access stderr during shutdown")
+        },
     }
 }