green: Properly wait for main before shutdown

There was a race in the code previously where schedulers would *immediately*
shut down after spawning the main task (because the global task count would
still be 0). This fixes the logic by blocking the sched pool task in receving on
a port instead of spawning a task into the pool to receive on a port.

The modifications necessary were to have a "simple task" running by the time the
code is executing, but this is a simple enough thing to implement and I forsee
this being necessary to have implemented in the future anyway.

1 files changed, 15 insertions, 12 deletions

diff --git a/src/libstd/rt/task.rs b/src/libstd/rt/task.rs
index c0e1086483d..765f0b427cd 100644
--- a/src/libstd/rt/task.rs
+++ b/src/libstd/rt/task.rs
@@ -292,6 +292,21 @@ impl Task {
     pub fn local_io<'a>(&'a mut self) -> Option<LocalIo<'a>> {
         self.imp.get_mut_ref().local_io()
     }
+
+    /// The main function of all rust executables will by default use this
+    /// function. This function will *block* the OS thread (hence the `unsafe`)
+    /// waiting for all known tasks to complete. Once this function has
+    /// returned, it is guaranteed that no more user-defined code is still
+    /// running.
+    pub unsafe fn wait_for_other_tasks(&mut self) {
+        TASK_COUNT.fetch_sub(1, SeqCst); // don't count ourselves
+        TASK_LOCK.lock();
+        while TASK_COUNT.load(SeqCst) > 0 {
+            TASK_LOCK.wait();
+        }
+        TASK_LOCK.unlock();
+        TASK_COUNT.fetch_add(1, SeqCst); // add ourselves back in
+    }
 }
 
 impl Drop for Task {
@@ -396,18 +411,6 @@ impl Drop for Death {
     }
 }
 
-/// The main function of all rust executables will by default use this function.
-/// This function will *block* the OS thread (hence the `unsafe`) waiting for
-/// all known tasks to complete. Once this function has returned, it is
-/// guaranteed that no more user-defined code is still running.
-pub unsafe fn wait_for_completion() {
-    TASK_LOCK.lock();
-    while TASK_COUNT.load(SeqCst) > 0 {
-        TASK_LOCK.wait();
-    }
-    TASK_LOCK.unlock();
-}
-
 #[cfg(test)]
 mod test {
     use super::*;