rt: Allow some schedulers to stay alive even without tasks to execute

3 files changed, 26 insertions, 8 deletions

diff --git a/src/rt/rust_kernel.cpp b/src/rt/rust_kernel.cpp
index c7625537f40..5062e4de3c3 100644
--- a/src/rt/rust_kernel.cpp
+++ b/src/rt/rust_kernel.cpp
@@ -69,7 +69,7 @@ rust_kernel::create_scheduler(size_t num_threads) {
         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);
+            rust_scheduler(this, srv, num_threads, id, true);
         bool is_new = sched_table
             .insert(std::pair<rust_sched_id,
                               rust_scheduler*>(id, sched)).second;
diff --git a/src/rt/rust_scheduler.cpp b/src/rt/rust_scheduler.cpp
index 40d62c76849..57b166208f9 100644
--- a/src/rt/rust_scheduler.cpp
+++ b/src/rt/rust_scheduler.cpp
@@ -5,14 +5,16 @@
 rust_scheduler::rust_scheduler(rust_kernel *kernel,
                                rust_srv *srv,
                                size_t num_threads,
-                               rust_sched_id id) :
+                               rust_sched_id id,
+                               bool allow_exit) :
     kernel(kernel),
     srv(srv),
     env(srv->env),
     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();
@@ -103,12 +105,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();
     }
 }
@@ -139,3 +140,16 @@ 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();
+    }
+}
diff --git a/src/rt/rust_scheduler.h b/src/rt/rust_scheduler.h
index 722a7e168fb..48a33c52ef6 100644
--- a/src/rt/rust_scheduler.h
+++ b/src/rt/rust_scheduler.h
@@ -12,16 +12,17 @@ public:
     rust_srv *srv;
     rust_env *env;
 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;
 
@@ -35,7 +36,7 @@ private:
 
 public:
     rust_scheduler(rust_kernel *kernel, rust_srv *srv, size_t num_threads,
-                   rust_sched_id id);
+                   rust_sched_id id, bool allow_exit);
     ~rust_scheduler();
 
     void start_task_threads();
@@ -51,6 +52,9 @@ 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();
 };
 
 #endif /* RUST_SCHEDULER_H */