Rework at_exit to cleanup

Author: CDirkx

library/std/src/io/mod.rs

@@ -293,6 +293,10 @@ mod util;
 
 const DEFAULT_BUF_SIZE: usize = crate::sys_common::io::DEFAULT_BUF_SIZE;
 
+pub(crate) fn cleanup() {
+    stdio::cleanup()
+}
+
 struct Guard<'a> {
     buf: &'a mut Vec<u8>,
     len: usize,

library/std/src/io/stdio.rs

@@ -13,7 +13,6 @@ use crate::pin::Pin;
 use crate::sync::atomic::{AtomicBool, Ordering};
 use crate::sync::{Arc, Mutex, MutexGuard};
 use crate::sys::stdio;
-use crate::sys_common;
 use crate::sys_common::remutex::{ReentrantMutex, ReentrantMutexGuard};
 
 type LocalStream = Arc<Mutex<Vec<u8>>>;
@@ -508,6 +507,8 @@ pub struct StdoutLock<'a> {
     inner: ReentrantMutexGuard<'a, RefCell<LineWriter<StdoutRaw>>>,
 }
 
+static STDOUT: SyncOnceCell<ReentrantMutex<RefCell<LineWriter<StdoutRaw>>>> = SyncOnceCell::new();
+
 /// Constructs a new handle to the standard output of the current process.
 ///
 /// Each handle returned is a reference to a shared global buffer whose access
@@ -549,34 +550,28 @@ pub struct StdoutLock<'a> {
 /// ```
 #[stable(feature = "rust1", since = "1.0.0")]
 pub fn stdout() -> Stdout {
-    static INSTANCE: SyncOnceCell<ReentrantMutex<RefCell<LineWriter<StdoutRaw>>>> =
-        SyncOnceCell::new();
-
-    fn cleanup() {
-        if let Some(instance) = INSTANCE.get() {
-            // Flush the data and disable buffering during shutdown
-            // by replacing the line writer by one with zero
-            // buffering capacity.
-            // We use try_lock() instead of lock(), because someone
-            // might have leaked a StdoutLock, which would
-            // otherwise cause a deadlock here.
-            if let Some(lock) = Pin::static_ref(instance).try_lock() {
-                *lock.borrow_mut() = LineWriter::with_capacity(0, stdout_raw());
-            }
-        }
-    }
-
     Stdout {
-        inner: Pin::static_ref(&INSTANCE).get_or_init_pin(
-            || unsafe {
-                let _ = sys_common::at_exit(cleanup);
-                ReentrantMutex::new(RefCell::new(LineWriter::new(stdout_raw())))
-            },
+        inner: Pin::static_ref(&STDOUT).get_or_init_pin(
+            || unsafe { ReentrantMutex::new(RefCell::new(LineWriter::new(stdout_raw()))) },
             |mutex| unsafe { mutex.init() },
         ),
     }
 }
 
+pub fn cleanup() {
+    if let Some(instance) = STDOUT.get() {
+        // Flush the data and disable buffering during shutdown
+        // by replacing the line writer by one with zero
+        // buffering capacity.
+        // We use try_lock() instead of lock(), because someone
+        // might have leaked a StdoutLock, which would
+        // otherwise cause a deadlock here.
+        if let Some(lock) = Pin::static_ref(instance).try_lock() {
+            *lock.borrow_mut() = LineWriter::with_capacity(0, stdout_raw());
+        }
+    }
+}
+
 impl Stdout {
     /// Locks this handle to the standard output stream, returning a writable
     /// guard.

library/std/src/process.rs

@@ -1749,7 +1749,7 @@ impl Child {
 /// [platform-specific behavior]: #platform-specific-behavior
 #[stable(feature = "rust1", since = "1.0.0")]
 pub fn exit(code: i32) -> ! {
-    crate::sys_common::cleanup();
+    crate::sys_common::rt::cleanup();
     crate::sys::os::exit(code)
 }
 

library/std/src/rt.rs

@@ -26,33 +26,13 @@ fn lang_start_internal(
     argv: *const *const u8,
 ) -> isize {
     use crate::panic;
-    use crate::sys;
     use crate::sys_common;
-    use crate::sys_common::thread_info;
-    use crate::thread::Thread;
 
-    sys::init();
+    sys_common::rt::init(argc, argv);
+    let exit_code = panic::catch_unwind(main);
+    sys_common::rt::cleanup();
 
-    unsafe {
-        let main_guard = sys::thread::guard::init();
-        sys::stack_overflow::init();
-
-        // Next, set up the current Thread with the guard information we just
-        // created. Note that this isn't necessary in general for new threads,
-        // but we just do this to name the main thread and to give it correct
-        // info about the stack bounds.
-        let thread = Thread::new(Some("main".to_owned()));
-        thread_info::set(main_guard, thread);
-
-        // Store our args if necessary in a squirreled away location
-        sys::args::init(argc, argv);
-
-        // Let's run some code!
-        let exit_code = panic::catch_unwind(main);
-
-        sys_common::cleanup();
-        exit_code.unwrap_or(101) as isize
-    }
+    exit_code.unwrap_or(101) as isize
 }
 
 #[cfg(not(test))]

library/std/src/sys/hermit/mod.rs

@@ -96,11 +96,14 @@ pub extern "C" fn __rust_abort() {
     abort_internal();
 }
 
-#[cfg(not(test))]
-pub fn init() {
+// SAFETY: must be called only once during runtime initialization.
+pub unsafe fn init() {
     let _ = net::init();
 }
 
+// SAFETY: must be called only once during runtime cleanup.
+pub unsafe fn cleanup() {}
+
 #[cfg(not(test))]
 #[no_mangle]
 pub unsafe extern "C" fn runtime_entry(

library/std/src/sys/sgx/mod.rs

@@ -40,8 +40,11 @@ pub mod time;
 
 pub use crate::sys_common::os_str_bytes as os_str;
 
-#[cfg(not(test))]
-pub fn init() {}
+// SAFETY: must be called only once during runtime initialization.
+pub unsafe fn init() {}
+
+// SAFETY: must be called only once during runtime cleanup.
+pub unsafe fn cleanup() {}
 
 /// This function is used to implement functionality that simply doesn't exist.
 /// Programs relying on this functionality will need to deal with the error.

library/std/src/sys/unix/mod.rs

@@ -44,14 +44,12 @@ pub mod time;
 
 pub use crate::sys_common::os_str_bytes as os_str;
 
-#[cfg(not(test))]
-pub fn init() {
+// SAFETY: must be called only once during runtime initialization.
+pub unsafe fn init() {
     // The standard streams might be closed on application startup. To prevent
     // std::io::{stdin, stdout,stderr} objects from using other unrelated file
     // resources opened later, we reopen standards streams when they are closed.
-    unsafe {
-        sanitize_standard_fds();
-    }
+    sanitize_standard_fds();
 
     // By default, some platforms will send a *signal* when an EPIPE error
     // would otherwise be delivered. This runtime doesn't install a SIGPIPE
@@ -60,9 +58,7 @@ pub fn init() {
     //
     // Hence, we set SIGPIPE to ignore when the program starts up in order
     // to prevent this problem.
-    unsafe {
-        reset_sigpipe();
-    }
+    reset_sigpipe();
 
     cfg_if::cfg_if! {
         if #[cfg(miri)] {
@@ -129,6 +125,9 @@ pub fn init() {
     unsafe fn reset_sigpipe() {}
 }
 
+// SAFETY: must be called only once during runtime cleanup.
+pub unsafe fn cleanup() {}
+
 #[cfg(target_os = "android")]
 pub use crate::sys::android::signal;
 #[cfg(not(target_os = "android"))]

library/std/src/sys/unsupported/common.rs

@@ -10,8 +10,11 @@ pub use crate::sys_common::os_str_bytes as os_str;
 // spec definition?
 use crate::os::raw::c_char;
 
-#[cfg(not(test))]
-pub fn init() {}
+// SAFETY: must be called only once during runtime initialization.
+pub unsafe fn init() {}
+
+// SAFETY: must be called only once during runtime cleanup.
+pub unsafe fn cleanup() {}
 
 pub fn unsupported<T>() -> std_io::Result<T> {
     Err(unsupported_err())

library/std/src/sys/windows/mod.rs

@@ -49,8 +49,13 @@ cfg_if::cfg_if! {
     }
 }
 
-#[cfg(not(test))]
-pub fn init() {}
+// SAFETY: must be called only once during runtime initialization.
+pub unsafe fn init() {}
+
+// SAFETY: must be called only once during runtime cleanup.
+pub unsafe fn cleanup() {
+    net::cleanup();
+}
 
 pub fn decode_error_kind(errno: i32) -> ErrorKind {
     match errno as c::DWORD {

library/std/src/sys/windows/net.rs

@@ -9,7 +9,7 @@ use crate::sync::Once;
 use crate::sys;
 use crate::sys::c;
 use crate::sys_common::net;
-use crate::sys_common::{self, AsInner, FromInner, IntoInner};
+use crate::sys_common::{AsInner, FromInner, IntoInner};
 use crate::time::Duration;
 
 use libc::{c_int, c_long, c_ulong, c_void};
@@ -38,13 +38,15 @@ pub fn init() {
             &mut data,
         );
         assert_eq!(ret, 0);
-
-        let _ = sys_common::at_exit(|| {
-            c::WSACleanup();
-        });
     });
 }
 
+pub fn cleanup() {
+    unsafe {
+        c::WSACleanup();
+    }
+}
+
 /// Returns the last error from the Windows socket interface.
 fn last_error() -> io::Error {
     io::Error::from_raw_os_error(unsafe { c::WSAGetLastError() })