std: Move libnative task count bookkeeping to std

When using tasks in Rust, the expectation is that the runtime does not exit
before all tasks have exited. This is enforced in libgreen through the
`SchedPool` type, and it is enforced in libnative through a `bookkeeping` module
and a global count/mutex pair. Unfortunately, this means that a process which
originates with libgreen will not wait for spawned native tasks.

In order to fix this problem, the bookkeeping module was moved from libnative to
libstd so the runtime itself can wait for native tasks to exit. Green tasks do
not manage themselves through this bookkeeping module, but native tasks will
continue to manage themselves through this module.

Closes #12684

2 files changed, 56 insertions, 0 deletions

diff --git a/src/libstd/rt/bookkeeping.rs b/src/libstd/rt/bookkeeping.rs
new file mode 100644
index 00000000000..5851a6a39c6
--- /dev/null
+++ b/src/libstd/rt/bookkeeping.rs
@@ -0,0 +1,52 @@
+// Copyright 2013-2014 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+//! Task bookkeeping
+//!
+//! This module keeps track of the number of running tasks so that entry points
+//! with libnative know when it's possible to exit the program (once all tasks
+//! have exited).
+//!
+//! The green counterpart for this is bookkeeping on sched pools, and it's up to
+//! each respective runtime to make sure that they call increment() and
+//! decrement() manually.
+
+#[experimental]; // this is a massive code smell
+#[doc(hidden)];
+
+use sync::atomics;
+use unstable::mutex::{StaticNativeMutex, NATIVE_MUTEX_INIT};
+
+static mut TASK_COUNT: atomics::AtomicUint = atomics::INIT_ATOMIC_UINT;
+static mut TASK_LOCK: StaticNativeMutex = NATIVE_MUTEX_INIT;
+
+pub fn increment() {
+    let _ = unsafe { TASK_COUNT.fetch_add(1, atomics::SeqCst) };
+}
+
+pub fn decrement() {
+    unsafe {
+        if TASK_COUNT.fetch_sub(1, atomics::SeqCst) == 1 {
+            let mut guard = TASK_LOCK.lock();
+            guard.signal();
+        }
+    }
+}
+
+/// Waits for all other native tasks in the system to exit. This is only used by
+/// the entry points of native programs
+pub fn wait_for_other_tasks() {
+    unsafe {
+        let mut guard = TASK_LOCK.lock();
+        while TASK_COUNT.load(atomics::SeqCst) > 0 {
+            guard.wait();
+        }
+    }
+}
diff --git a/src/libstd/rt/mod.rs b/src/libstd/rt/mod.rs
index 2f1a6989092..459bc061c56 100644
--- a/src/libstd/rt/mod.rs
+++ b/src/libstd/rt/mod.rs
@@ -128,6 +128,9 @@ pub mod args;
 // Support for running procedures when a program has exited.
 mod at_exit_imp;
 
+// Bookkeeping for task counts
+pub mod bookkeeping;
+
 // Stack overflow protection
 pub mod stack;
 
@@ -207,6 +210,7 @@ pub fn at_exit(f: proc()) {
 /// Invoking cleanup while portions of the runtime are still in use may cause
 /// undefined behavior.
 pub unsafe fn cleanup() {
+    bookkeeping::wait_for_other_tasks();
     at_exit_imp::run();
     args::cleanup();
     local_ptr::cleanup();