auto merge of #11212 : alexcrichton/rust/local-task-count, r=brson

For libgreen, bookeeping should not be global but rather on a per-pool basis.
Inside libnative, it's known that there must be a global counter with a
mutex/cvar.

The benefit of taking this strategy is to remove this functionality from libstd
to allow fine-grained control of it through libnative/libgreen. Notably, helper
threads in libnative can manually decrement the global count so they don't count
towards the global count of threads. Also, the shutdown process of *all* sched
pools is now dependent on the number of tasks in the pool being 0 rather than
this only being a hardcoded solution for the initial sched pool in libgreen.

This involved adding a Local::try_take() method on the Local trait in order for
the channel wakeup to work inside of libgreen. The channel send was happening
from a SchedTask when there is no Task available in TLS, and now this is
possible to work (remote wakeups are always possible, just a little slower).

3 files changed, 54 insertions, 8 deletions

diff --git a/src/libnative/bookeeping.rs b/src/libnative/bookeeping.rs
new file mode 100644
index 00000000000..ca40c1a1958
--- /dev/null
+++ b/src/libnative/bookeeping.rs
@@ -0,0 +1,49 @@
+// Copyright 2013 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.
+
+//! 1:1 Task bookeeping
+//!
+//! This module keeps track of the number of running 1:1 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 bookeeping on sched pools.
+
+use std::sync::atomics;
+use std::unstable::mutex::{Mutex, MUTEX_INIT};
+
+static mut TASK_COUNT: atomics::AtomicUint = atomics::INIT_ATOMIC_UINT;
+static mut TASK_LOCK: Mutex = MUTEX_INIT;
+
+pub fn increment() {
+    unsafe { TASK_COUNT.fetch_add(1, atomics::SeqCst); }
+}
+
+pub fn decrement() {
+    unsafe {
+        if TASK_COUNT.fetch_sub(1, atomics::SeqCst) == 1 {
+            TASK_LOCK.lock();
+            TASK_LOCK.signal();
+            TASK_LOCK.unlock();
+        }
+    }
+}
+
+/// 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 {
+        TASK_LOCK.lock();
+        while TASK_COUNT.load(atomics::SeqCst) > 0 {
+            TASK_LOCK.wait();
+        }
+        TASK_LOCK.unlock();
+    }
+}
diff --git a/src/libnative/lib.rs b/src/libnative/lib.rs
index d92127cfbb2..498945a04cb 100644
--- a/src/libnative/lib.rs
+++ b/src/libnative/lib.rs
@@ -27,14 +27,12 @@
 //    answer is that you don't need them)
 
 use std::os;
-use std::rt::local::Local;
-use std::rt::task::Task;
 use std::rt;
 
+mod bookeeping;
 pub mod io;
 pub mod task;
 
-
 // XXX: this should not exist here
 #[cfg(stage0)]
 #[lang = "start"]
@@ -83,11 +81,7 @@ pub fn start(argc: int, argv: **u8, main: proc()) -> int {
 /// This function has all of the same details as `start` except for a different
 /// number of arguments.
 pub fn run(main: proc()) -> int {
-    // Run the main procedure and then wait for everything to finish
     main();
-    unsafe {
-        let mut task = Local::borrow(None::<Task>);
-        task.get().wait_for_other_tasks();
-    }
+    bookeeping::wait_for_other_tasks();
     os::get_exit_status()
 }
diff --git a/src/libnative/task.rs b/src/libnative/task.rs
index 8f2dff42404..c4d3f651777 100644
--- a/src/libnative/task.rs
+++ b/src/libnative/task.rs
@@ -27,6 +27,7 @@ use std::unstable::stack;
 
 use io;
 use task;
+use bookeeping;
 
 /// Creates a new Task which is ready to execute as a 1:1 task.
 pub fn new() -> ~Task {
@@ -79,8 +80,10 @@ pub fn spawn_opts(opts: TaskOpts, f: proc()) {
             stack::record_stack_bounds(my_stack - stack + 1024, my_stack);
         }
 
+        bookeeping::increment();
         let mut f = Some(f);
         task.run(|| { f.take_unwrap()() });
+        bookeeping::decrement();
     })
 }