rt: Fix bugs in the osmain scheduler

src/rt/rust_kernel.cpp

@@ -78,6 +78,18 @@ rust_kernel::create_scheduler(rust_sched_launcher_factory *launchfac,
     rust_scheduler *sched;
     {
         scoped_lock with(sched_lock);
+
+        if (sched_table.size() == 1) {
+            // The OS main scheduler may not exit while there are other
+            // schedulers
+            KLOG_("Disallowing osmain scheduler to exit");
+            sched_lock.unlock();
+            rust_scheduler *sched = get_scheduler_by_id(osmain_scheduler);
+            assert(sched != NULL);
+            sched->disallow_exit();
+            sched_lock.lock();
+        }
+
         id = max_sched_id++;
         K(srv, id != INTPTR_MAX, "Hit the maximum scheduler id");
         sched = new (this, "rust_scheduler")
@@ -138,8 +150,8 @@ rust_kernel::wait_for_schedulers()
                 // 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();
+                sched_lock.lock();
             }
         }
         if (!sched_table.empty()) {

src/rt/rust_scheduler.cpp

@@ -152,3 +152,9 @@ rust_scheduler::allow_exit() {
         exit();
     }
 }
+
+void
+rust_scheduler::disallow_exit() {
+    scoped_lock with(lock);
+    may_exit = false;
+}

src/rt/rust_scheduler.h

@@ -58,6 +58,7 @@ public:
     // 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 */

src/test/run-pass/osmain.rs

@@ -3,7 +3,7 @@
 // while it is not in use.
 
 fn main() {
-    run(10);
+    run(100);
 }
 
 fn run(i: int) {