rt: Run a single-threaded scheduler on the main thread

7 files changed, 57 insertions, 21 deletions

diff --git a/src/rt/rust.cpp b/src/rt/rust.cpp
index 00657f8ec0a..e328f81c391 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;
     delete srv;
diff --git a/src/rt/rust_kernel.cpp b/src/rt/rust_kernel.cpp
index 5062e4de3c3..d65837a47d6 100644
--- a/src/rt/rust_kernel.cpp
+++ b/src/rt/rust_kernel.cpp
@@ -1,6 +1,7 @@
 #include "rust_internal.h"
 #include "rust_util.h"
 #include "rust_scheduler.h"
+#include "rust_sched_launcher.h"
 
 #include <vector>
 
@@ -18,8 +19,15 @@ rust_kernel::rust_kernel(rust_srv *srv) :
     rval(0),
     max_sched_id(0),
     sched_reaper(this),
+    osmain_driver(NULL),
     env(srv->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
@@ -59,24 +67,25 @@ 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();
         id = max_sched_id++;
         K(srv, id != INTPTR_MAX, "Hit the maximum scheduler id");
         sched = new (this, "rust_scheduler")
-            rust_scheduler(this, srv, num_threads, id, true);
+            rust_scheduler(this, srv, num_threads, id, allow_exit, launchfac);
         bool is_new = sched_table
             .insert(std::pair<rust_sched_id,
                               rust_scheduler*>(id, sched)).second;
         A(this, is_new, "Reusing a sched id?");
-        if (start_reaper) {
-            sched_reaper.start();
-        }
     }
     sched->start_task_threads();
     return id;
@@ -123,6 +132,15 @@ 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");
+                sched_lock.unlock();
+                // It's only the osmain scheduler left. Tell it to exit
+                rust_scheduler *sched = get_scheduler_by_id(osmain_scheduler);
+                assert(sched != NULL);
+                sched_lock.lock();
+                sched->allow_exit();
+            }
         }
         if (!sched_table.empty()) {
             sched_lock.wait();
@@ -130,9 +148,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 78548917ee5..8c1d33c6a5e 100644
--- a/src/rt/rust_kernel.h
+++ b/src/rt/rust_kernel.h
@@ -13,6 +13,9 @@ class rust_scheduler;
 
 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
@@ -48,6 +51,11 @@ private:
     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;
 
 public:
 
@@ -66,11 +74,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);
diff --git a/src/rt/rust_sched_launcher.cpp b/src/rt/rust_sched_launcher.cpp
index dabae31f5b6..330299ec73c 100644
--- a/src/rt/rust_sched_launcher.cpp
+++ b/src/rt/rust_sched_launcher.cpp
@@ -32,6 +32,7 @@ rust_sched_launcher *
 rust_manual_sched_launcher_factory::create(rust_scheduler *sched, int id) {
     assert(launcher == NULL && "I can only track one sched_launcher");
     rust_srv *srv = sched->srv->clone();
-    return new(sched->kernel, "rust_manual_sched_launcher")
+    launcher = new(sched->kernel, "rust_manual_sched_launcher")
         rust_manual_sched_launcher(sched, srv, id);
+    return launcher;
 }
diff --git a/src/rt/rust_sched_launcher.h b/src/rt/rust_sched_launcher.h
index a68c1f500d6..f5db14a397e 100644
--- a/src/rt/rust_sched_launcher.h
+++ b/src/rt/rust_sched_launcher.h
@@ -36,10 +36,10 @@ public:
 
 class rust_manual_sched_launcher : public rust_sched_launcher {
 public:
-  rust_manual_sched_launcher(rust_scheduler *sched, rust_srv *srv, int id);
-  virtual void start() { }
-  virtual void join() { }
-  void start_main_loop() { driver.start_main_loop(); }
+    rust_manual_sched_launcher(rust_scheduler *sched, rust_srv *srv, int id);
+    virtual void start() { }
+    virtual void join() { }
+    rust_sched_driver *get_driver() { return &driver; };
 };
 
 class rust_sched_launcher_factory {
@@ -62,7 +62,10 @@ private:
 public:
     rust_manual_sched_launcher_factory() : launcher(NULL) { }
     virtual rust_sched_launcher *create(rust_scheduler *sched, int id);
-    rust_manual_sched_launcher *get_launcher() { return launcher; }
+    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 82eac1b1c93..b9f544b4549 100644
--- a/src/rt/rust_scheduler.cpp
+++ b/src/rt/rust_scheduler.cpp
@@ -6,7 +6,8 @@ rust_scheduler::rust_scheduler(rust_kernel *kernel,
                                rust_srv *srv,
                                size_t num_threads,
                                rust_sched_id id,
-                               bool allow_exit) :
+                               bool allow_exit,
+                               rust_sched_launcher_factory *launchfac) :
     kernel(kernel),
     srv(srv),
     env(srv->env),
@@ -17,8 +18,7 @@ rust_scheduler::rust_scheduler(rust_kernel *kernel,
     num_threads(num_threads),
     id(id)
 {
-    rust_thread_sched_launcher_factory launchfac;
-    create_task_threads(&launchfac);
+    create_task_threads(launchfac);
 }
 
 rust_scheduler::~rust_scheduler() {
diff --git a/src/rt/rust_scheduler.h b/src/rt/rust_scheduler.h
index 132fa9e6089..af501511832 100644
--- a/src/rt/rust_scheduler.h
+++ b/src/rt/rust_scheduler.h
@@ -38,7 +38,8 @@ private:
 
 public:
     rust_scheduler(rust_kernel *kernel, rust_srv *srv, size_t num_threads,
-                   rust_sched_id id, bool allow_exit);
+                   rust_sched_id id, bool allow_exit,
+                   rust_sched_launcher_factory *launchfac);
     ~rust_scheduler();
 
     void start_task_threads();