Merge remote-tracking branch 'brson/mainthread'

Conflicts:
	src/rt/rust_kernel.cpp
	src/rt/rust_scheduler.cpp
	src/rt/rust_scheduler.h

11 files changed, 248 insertions, 30 deletions

diff --git a/src/libcore/task.rs b/src/libcore/task.rs
index dc2dd5d179e..9bffb34da16 100644
--- a/src/libcore/task.rs
+++ b/src/libcore/task.rs
@@ -88,6 +88,13 @@ enum sched_mode {
     thread_per_task,
     #[doc = "Tasks are distributed among a fixed number of OS threads"]
     manual_threads(uint),
+    #[doc = "
+    Tasks are scheduled on the main OS thread
+
+    The main OS thread is the thread used to launch the runtime which,
+    in most cases, is the process's initial thread as created by the OS.
+    "]
+    osmain
 }
 
 #[doc = "
@@ -107,7 +114,7 @@ Scheduler configuration options
 "]
 type sched_opts = {
     mode: sched_mode,
-    native_stack_size: option<uint>,
+    native_stack_size: option<uint>
 };
 
 #[doc = "
@@ -525,9 +532,14 @@ fn spawn_raw(opts: task_opts, +f: fn~()) unsafe {
             }
             threads
           }
+          osmain { 0u /* Won't be used */ }
         };
 
-        let sched_id = rustrt::rust_new_sched(num_threads);
+        let sched_id = if opts.mode != osmain {
+            rustrt::rust_new_sched(num_threads)
+        } else {
+            rustrt::rust_osmain_sched_id()
+        };
         rustrt::rust_new_task_in_sched(sched_id)
     }
 
@@ -553,6 +565,7 @@ native mod rustrt {
 
     fn rust_task_is_unwinding(rt: *rust_task) -> bool;
     fn unsupervise();
+    fn rust_osmain_sched_id() -> sched_id;
 }
 
 
@@ -897,3 +910,23 @@ fn test_avoid_copying_the_body_unsupervise() {
         }
     }
 }
+
+#[test]
+fn test_osmain() {
+    let builder = task_builder();
+    let opts = {
+        sched: some({
+            mode: osmain,
+            native_stack_size: none
+        })
+        with get_opts(builder)
+    };
+    set_opts(builder, opts);
+
+    let po = comm::port();
+    let ch = comm::chan(po);
+    run(builder) {||
+        comm::send(ch, ());
+    }
+    comm::recv(po);
+}
diff --git a/src/rt/rust.cpp b/src/rt/rust.cpp
index 4f9cd3bf7fd..2fed06854cb 100644
--- a/src/rt/rust.cpp
+++ b/src/rt/rust.cpp
@@ -93,7 +93,7 @@ rust_start(uintptr_t main_fn, int argc, char **argv, void* crate_map) {
     root_task->start((spawn_fn)main_fn, NULL, args->args);
     root_task = NULL;
 
-    int ret = kernel->wait_for_exit();
+    int ret = kernel->run();
     delete args;
     delete kernel;
 
diff --git a/src/rt/rust_builtin.cpp b/src/rt/rust_builtin.cpp
index 79e1e880f57..0d1b67674a1 100644
--- a/src/rt/rust_builtin.cpp
+++ b/src/rt/rust_builtin.cpp
@@ -640,6 +640,12 @@ rust_dbg_call(dbg_callback cb, void *data) {
     return cb(data);
 }
 
+extern "C" CDECL rust_sched_id
+rust_osmain_sched_id() {
+    rust_task *task = rust_get_current_task();
+    return task->kernel->osmain_sched_id();
+}
+
 //
 // Local Variables:
 // mode: C++
diff --git a/src/rt/rust_kernel.cpp b/src/rt/rust_kernel.cpp
index e8e7e4af171..ff6b0e1056c 100644
--- a/src/rt/rust_kernel.cpp
+++ b/src/rt/rust_kernel.cpp
@@ -5,6 +5,7 @@
 #include "rust_port.h"
 #include "rust_util.h"
 #include "rust_scheduler.h"
+#include "rust_sched_launcher.h"
 
 #define KLOG_(...)                              \
     KLOG(this, kern, __VA_ARGS__)
@@ -19,8 +20,15 @@ rust_kernel::rust_kernel(rust_env *env) :
     rval(0),
     max_sched_id(0),
     sched_reaper(this),
+    osmain_driver(NULL),
     env(env)
 {
+    // Create the single threaded scheduler that will run on the platform's
+    // main thread
+    rust_manual_sched_launcher_factory launchfac;
+    osmain_scheduler = create_scheduler(&launchfac, 1, false);
+    osmain_driver = launchfac.get_driver();
+    sched_reaper.start();
 }
 
 void
@@ -60,24 +68,36 @@ void rust_kernel::free(void *mem) {
 
 rust_sched_id
 rust_kernel::create_scheduler(size_t num_threads) {
+    rust_thread_sched_launcher_factory launchfac;
+    return create_scheduler(&launchfac, num_threads, true);
+}
+
+rust_sched_id
+rust_kernel::create_scheduler(rust_sched_launcher_factory *launchfac,
+                              size_t num_threads, bool allow_exit) {
     rust_sched_id id;
     rust_scheduler *sched;
     {
         scoped_lock with(sched_lock);
-        // If this is the first scheduler then we need to launch
-        // the scheduler reaper.
-        bool start_reaper = sched_table.empty();
+
+        if (sched_table.size() == 1) {
+            // The OS main scheduler may not exit while there are other
+            // schedulers
+            KLOG_("Disallowing osmain scheduler to exit");
+            rust_scheduler *sched =
+                get_scheduler_by_id_nolock(osmain_scheduler);
+            assert(sched != NULL);
+            sched->disallow_exit();
+        }
+
         id = max_sched_id++;
         assert(id != INTPTR_MAX && "Hit the maximum scheduler id");
         sched = new (this, "rust_scheduler")
-            rust_scheduler(this, num_threads, id);
+            rust_scheduler(this, num_threads, id, allow_exit, launchfac);
         bool is_new = sched_table
             .insert(std::pair<rust_sched_id,
                               rust_scheduler*>(id, sched)).second;
         assert(is_new && "Reusing a sched id?");
-        if (start_reaper) {
-            sched_reaper.start();
-        }
     }
     sched->start_task_threads();
     return id;
@@ -86,6 +106,12 @@ rust_kernel::create_scheduler(size_t num_threads) {
 rust_scheduler *
 rust_kernel::get_scheduler_by_id(rust_sched_id id) {
     scoped_lock with(sched_lock);
+    return get_scheduler_by_id_nolock(id);
+}
+
+rust_scheduler *
+rust_kernel::get_scheduler_by_id_nolock(rust_sched_id id) {
+    sched_lock.must_have_lock();
     sched_map::iterator iter = sched_table.find(id);
     if (iter != sched_table.end()) {
         return iter->second;
@@ -117,6 +143,7 @@ rust_kernel::wait_for_schedulers()
     while (!sched_table.empty()) {
         while (!join_list.empty()) {
             rust_sched_id id = join_list.back();
+            KLOG_("Deleting scheduler %d", id);
             join_list.pop_back();
             sched_map::iterator iter = sched_table.find(id);
             assert(iter != sched_table.end());
@@ -124,6 +151,14 @@ rust_kernel::wait_for_schedulers()
             sched_table.erase(iter);
             sched->join_task_threads();
             delete sched;
+            if (sched_table.size() == 1) {
+                KLOG_("Allowing osmain scheduler to exit");
+                // It's only the osmain scheduler left. Tell it to exit
+                rust_scheduler *sched =
+                    get_scheduler_by_id_nolock(osmain_scheduler);
+                assert(sched != NULL);
+                sched->allow_exit();
+            }
         }
         if (!sched_table.empty()) {
             sched_lock.wait();
@@ -131,9 +166,12 @@ rust_kernel::wait_for_schedulers()
     }
 }
 
-/* Called on the main thread to wait for the kernel to exit */
+/* Called on the main thread to run the osmain scheduler to completion,
+   then wait for schedulers to exit */
 int
-rust_kernel::wait_for_exit() {
+rust_kernel::run() {
+    assert(osmain_driver != NULL);
+    osmain_driver->start_main_loop();
     sched_reaper.join();
     return rval;
 }
diff --git a/src/rt/rust_kernel.h b/src/rt/rust_kernel.h
index 0d96f78067f..6a7a7070c20 100644
--- a/src/rt/rust_kernel.h
+++ b/src/rt/rust_kernel.h
@@ -22,6 +22,9 @@ typedef intptr_t rust_port_id;
 
 typedef std::map<rust_sched_id, rust_scheduler*> sched_map;
 
+class rust_sched_driver;
+class rust_sched_launcher_factory;
+
 /**
  * A global object shared by all thread domains. Most of the data structures
  * in this class are synchronized since they are accessed from multiple
@@ -54,6 +57,13 @@ class rust_kernel {
     std::vector<rust_sched_id> join_list;
 
     rust_sched_reaper sched_reaper;
+    // The single-threaded scheduler that uses the main thread
+    rust_sched_id osmain_scheduler;
+    // Runs the single-threaded scheduler that executes tasks
+    // on the main thread
+    rust_sched_driver *osmain_driver;
+
+    rust_scheduler* get_scheduler_by_id_nolock(rust_sched_id id);
 
 public:
     struct rust_env *env;
@@ -71,11 +81,13 @@ public:
     void fail();
 
     rust_sched_id create_scheduler(size_t num_threads);
+    rust_sched_id create_scheduler(rust_sched_launcher_factory *launchfac,
+                                   size_t num_threads, bool allow_exit);
     rust_scheduler* get_scheduler_by_id(rust_sched_id id);
     // Called by a scheduler to indicate that it is terminating
     void release_scheduler_id(rust_sched_id id);
     void wait_for_schedulers();
-    int wait_for_exit();
+    int run();
 
 #ifdef __WIN32__
     void win32_require(LPCTSTR fn, BOOL ok);
@@ -88,6 +100,8 @@ public:
     void release_port_id(rust_port_id tid);
 
     void set_exit_status(int code);
+
+    rust_sched_id osmain_sched_id() { return osmain_scheduler; }
 };
 
 template <typename T> struct kernel_owned {
diff --git a/src/rt/rust_sched_launcher.cpp b/src/rt/rust_sched_launcher.cpp
index 3406835dfbb..bf0481b02f2 100644
--- a/src/rt/rust_sched_launcher.cpp
+++ b/src/rt/rust_sched_launcher.cpp
@@ -16,3 +16,21 @@ rust_thread_sched_launcher::rust_thread_sched_launcher(rust_scheduler *sched,
       rust_thread(SCHED_STACK_SIZE) {
 }
 
+rust_manual_sched_launcher::rust_manual_sched_launcher(rust_scheduler *sched,
+                                                       int id)
+    : rust_sched_launcher(sched, id) {
+}
+
+rust_sched_launcher *
+rust_thread_sched_launcher_factory::create(rust_scheduler *sched, int id) {
+    return new(sched->kernel, "rust_thread_sched_launcher")
+        rust_thread_sched_launcher(sched, id);
+}
+
+rust_sched_launcher *
+rust_manual_sched_launcher_factory::create(rust_scheduler *sched, int id) {
+    assert(launcher == NULL && "I can only track one sched_launcher");
+    launcher = new(sched->kernel, "rust_manual_sched_launcher")
+        rust_manual_sched_launcher(sched, id);
+    return launcher;
+}
diff --git a/src/rt/rust_sched_launcher.h b/src/rt/rust_sched_launcher.h
index 523a199401e..ace2cba2530 100644
--- a/src/rt/rust_sched_launcher.h
+++ b/src/rt/rust_sched_launcher.h
@@ -31,8 +31,42 @@ class rust_thread_sched_launcher
 public:
     rust_thread_sched_launcher(rust_scheduler *sched, int id);
     virtual void start() { rust_thread::start(); }
-    virtual void run() { driver.start_main_loop(); }
     virtual void join() { rust_thread::join(); }
+    virtual void run() { driver.start_main_loop(); }
+};
+
+class rust_manual_sched_launcher : public rust_sched_launcher {
+public:
+    rust_manual_sched_launcher(rust_scheduler *sched, int id);
+    virtual void start() { }
+    virtual void join() { }
+    rust_sched_driver *get_driver() { return &driver; };
+};
+
+class rust_sched_launcher_factory {
+public:
+    virtual ~rust_sched_launcher_factory() { }
+    virtual rust_sched_launcher *
+    create(rust_scheduler *sched, int id) = 0;
+};
+
+class rust_thread_sched_launcher_factory
+    : public rust_sched_launcher_factory {
+public:
+    virtual rust_sched_launcher *create(rust_scheduler *sched, int id);
+};
+
+class rust_manual_sched_launcher_factory
+    : public rust_sched_launcher_factory {
+private:
+    rust_manual_sched_launcher *launcher;
+public:
+    rust_manual_sched_launcher_factory() : launcher(NULL) { }
+    virtual rust_sched_launcher *create(rust_scheduler *sched, int id);
+    rust_sched_driver *get_driver() {
+        assert(launcher != NULL);
+        return launcher->get_driver();
+    }
 };
 
 #endif // RUST_SCHED_LAUNCHER_H
diff --git a/src/rt/rust_scheduler.cpp b/src/rt/rust_scheduler.cpp
index af52b48f985..1b5978b5cfe 100644
--- a/src/rt/rust_scheduler.cpp
+++ b/src/rt/rust_scheduler.cpp
@@ -7,15 +7,18 @@
 
 rust_scheduler::rust_scheduler(rust_kernel *kernel,
                                size_t num_threads,
-                               rust_sched_id id) :
+                               rust_sched_id id,
+                               bool allow_exit,
+                               rust_sched_launcher_factory *launchfac) :
     kernel(kernel),
     live_threads(num_threads),
     live_tasks(0),
-    num_threads(num_threads),
     cur_thread(0),
+    may_exit(allow_exit),
+    num_threads(num_threads),
     id(id)
 {
-    create_task_threads();
+    create_task_threads(launchfac);
 }
 
 rust_scheduler::~rust_scheduler() {
@@ -23,10 +26,9 @@ rust_scheduler::~rust_scheduler() {
 }
 
 rust_sched_launcher *
-rust_scheduler::create_task_thread(int id) {
-    rust_sched_launcher *thread =
-        new (kernel, "rust_thread_sched_launcher")
-        rust_thread_sched_launcher(this, id);
+rust_scheduler::create_task_thread(rust_sched_launcher_factory *launchfac,
+                                   int id) {
+    rust_sched_launcher *thread = launchfac->create(this, id);
     KLOG(kernel, kern, "created task thread: " PTR ", id: %d",
           thread, id);
     return thread;
@@ -39,11 +41,11 @@ rust_scheduler::destroy_task_thread(rust_sched_launcher *thread) {
 }
 
 void
-rust_scheduler::create_task_threads() {
+rust_scheduler::create_task_threads(rust_sched_launcher_factory *launchfac) {
     KLOG(kernel, kern, "Using %d scheduler threads.", num_threads);
 
     for(size_t i = 0; i < num_threads; ++i) {
-        threads.push(create_task_thread(i));
+        threads.push(create_task_thread(launchfac, i));
     }
 }
 
@@ -100,12 +102,11 @@ rust_scheduler::release_task() {
     {
         scoped_lock with(lock);
         live_tasks--;
-        if (live_tasks == 0) {
+        if (live_tasks == 0 && may_exit) {
             need_exit = true;
         }
     }
     if (need_exit) {
-        // There are no more tasks on this scheduler. Time to leave
         exit();
     }
 }
@@ -136,3 +137,22 @@ rust_scheduler::release_task_thread() {
         kernel->release_scheduler_id(id);
     }
 }
+
+void
+rust_scheduler::allow_exit() {
+    bool need_exit = false;
+    {
+        scoped_lock with(lock);
+        may_exit = true;
+        need_exit = live_tasks == 0;
+    }
+    if (need_exit) {
+        exit();
+    }
+}
+
+void
+rust_scheduler::disallow_exit() {
+    scoped_lock with(lock);
+    may_exit = false;
+}
diff --git a/src/rt/rust_scheduler.h b/src/rt/rust_scheduler.h
index 84f3df195a4..fec951f276c 100644
--- a/src/rt/rust_scheduler.h
+++ b/src/rt/rust_scheduler.h
@@ -6,36 +6,40 @@
 #include "rust_kernel.h"
 
 class rust_sched_launcher;
+class rust_sched_launcher_factory;
 
 class rust_scheduler : public kernel_owned<rust_scheduler> {
     // FIXME: Make these private
 public:
     rust_kernel *kernel;
 private:
-    // Protects live_threads and cur_thread increments
+    // Protects live_threads, live_tasks, cur_thread, may_exit
     lock_and_signal lock;
     // When this hits zero we'll tell the kernel to release us
     uintptr_t live_threads;
     // When this hits zero we'll tell the threads to exit
     uintptr_t live_tasks;
+    size_t cur_thread;
+    bool may_exit;
 
     array_list<rust_sched_launcher *> threads;
     const size_t num_threads;
-    size_t cur_thread;
 
     rust_sched_id id;
 
-    void create_task_threads();
+    void create_task_threads(rust_sched_launcher_factory *launchfac);
     void destroy_task_threads();
 
-    rust_sched_launcher *create_task_thread(int id);
+    rust_sched_launcher *
+    create_task_thread(rust_sched_launcher_factory *launchfac, int id);
     void destroy_task_thread(rust_sched_launcher *thread);
 
     void exit();
 
 public:
     rust_scheduler(rust_kernel *kernel, size_t num_threads,
-                   rust_sched_id id);
+                   rust_sched_id id, bool allow_exit,
+                   rust_sched_launcher_factory *launchfac);
     ~rust_scheduler();
 
     void start_task_threads();
@@ -51,6 +55,10 @@ public:
     void release_task_thread();
 
     rust_sched_id get_id() { return id; }
+    // Tells the scheduler that as soon as it runs out of tasks
+    // to run it should exit
+    void allow_exit();
+    void disallow_exit();
 };
 
 #endif /* RUST_SCHEDULER_H */
diff --git a/src/rt/rustrt.def.in b/src/rt/rustrt.def.in
index c9f46a45177..cdfb5ff67a7 100644
--- a/src/rt/rustrt.def.in
+++ b/src/rt/rustrt.def.in
@@ -97,3 +97,4 @@ rust_dbg_lock_unlock
 rust_dbg_lock_wait
 rust_dbg_lock_signal
 rust_dbg_call
+rust_osmain_sched_id
diff --git a/src/test/run-pass/osmain.rs b/src/test/run-pass/osmain.rs
new file mode 100644
index 00000000000..70f59ee2e1b
--- /dev/null
+++ b/src/test/run-pass/osmain.rs
@@ -0,0 +1,46 @@
+// Jump back and forth between the OS main thread and a new scheduler.
+// The OS main scheduler should continue to be available and not terminate
+// while it is not in use.
+
+fn main() {
+    run(100);
+}
+
+fn run(i: int) {
+
+    log(debug, i);
+
+    if i == 0 {
+        ret;
+    }
+
+    let builder = task::task_builder();
+    let opts = {
+        sched: some({
+            mode: task::osmain,
+            native_stack_size: none
+        })
+        with task::get_opts(builder)
+    };
+    task::set_opts(builder, opts);
+    task::unsupervise(builder);
+    task::run(builder) {||
+        task::yield();
+        let builder = task::task_builder();
+        let opts = {
+            sched: some({
+                mode: task::single_threaded,
+                native_stack_size: none
+            })
+            with task::get_opts(builder)
+        };
+        task::set_opts(builder, opts);
+        task::unsupervise(builder);
+        task::run(builder) {||
+            task::yield();
+            run(i - 1);
+            task::yield();
+        }
+        task::yield();
+    }
+}