rt: Export a scheduler API

4 files changed, 36 insertions, 3 deletions

diff --git a/src/rt/rust_builtin.cpp b/src/rt/rust_builtin.cpp
index 76f80192e89..53093190efe 100644
--- a/src/rt/rust_builtin.cpp
+++ b/src/rt/rust_builtin.cpp
@@ -431,16 +431,43 @@ nano_time(uint64_t *ns) {
     *ns = t.time_ns();
 }
 
+extern "C" CDECL rust_sched_id
+rust_get_sched_id() {
+    rust_task *task = rust_task_thread::get_task();
+    return task->sched->get_id();
+}
+
+extern "C" CDECL rust_sched_id
+rust_new_sched(size_t threads) {
+    rust_task *task = rust_task_thread::get_task();
+    A(task->thread, threads > 0,
+      "Can't create a scheduler with no threads, silly!");
+    return task->kernel->create_scheduler(threads);
+}
+
 extern "C" CDECL rust_task_id
 get_task_id() {
     rust_task *task = rust_task_thread::get_task();
     return task->user.id;
 }
 
+static rust_task_id
+new_task_common(rust_scheduler *sched, rust_task *parent) {
+    return sched->create_task(parent, NULL);
+}
+
 extern "C" CDECL rust_task_id
 new_task() {
     rust_task *task = rust_task_thread::get_task();
-    return task->sched->create_task(task, NULL);
+    return new_task_common(task->sched, task);
+}
+
+extern "C" CDECL rust_task_id
+rust_new_task_in_sched(rust_sched_id id) {
+    rust_task *task = rust_task_thread::get_task();
+    rust_scheduler *sched = task->kernel->get_scheduler_by_id(id);
+    // FIXME: What if we didn't get the scheduler?
+    return new_task_common(sched, task);
 }
 
 extern "C" CDECL void
diff --git a/src/rt/rust_kernel.cpp b/src/rt/rust_kernel.cpp
index b534c7f5414..d014877c278 100644
--- a/src/rt/rust_kernel.cpp
+++ b/src/rt/rust_kernel.cpp
@@ -59,10 +59,11 @@ void rust_kernel::free(void *mem) {
 rust_sched_id
 rust_kernel::create_scheduler(size_t num_threads) {
     I(this, !sched_lock.lock_held_by_current_thread());
+    rust_sched_id id;
     rust_scheduler *sched;
     {
         scoped_lock with(sched_lock);
-        rust_sched_id id = max_sched_id++;
+        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);
@@ -72,7 +73,7 @@ rust_kernel::create_scheduler(size_t num_threads) {
         live_schedulers++;
     }
     sched->start_task_threads();
-    return 0;
+    return id;
 }
 
 rust_scheduler *
diff --git a/src/rt/rust_scheduler.h b/src/rt/rust_scheduler.h
index 533f773ee35..5a931cad348 100644
--- a/src/rt/rust_scheduler.h
+++ b/src/rt/rust_scheduler.h
@@ -49,6 +49,8 @@ public:
     // Called by each thread when it terminates. When all threads
     // terminate the scheduler does as well.
     void release_task_thread();
+
+    rust_sched_id get_id() { return id; }
 };
 
 #endif /* RUST_SCHEDULER_H */
diff --git a/src/rt/rustrt.def.in b/src/rt/rustrt.def.in
index 81f8a9c25f7..e51d47065f6 100644
--- a/src/rt/rustrt.def.in
+++ b/src/rt/rustrt.def.in
@@ -24,6 +24,9 @@ rand_free
 rand_new
 rand_next
 refcount
+rust_get_sched_id
+rust_new_sched
+rust_new_task_in_sched
 rust_path_is_dir
 rust_path_exists
 rust_getcwd