Cleanup failure handling around rust_new_task_in_sched - closes #2668

3 files changed, 6 insertions, 2 deletions

diff --git a/src/rt/rust.cpp b/src/rt/rust.cpp
index 67b3bf84938..150156ddae9 100644
--- a/src/rt/rust.cpp
+++ b/src/rt/rust.cpp
@@ -95,6 +95,7 @@ rust_start(uintptr_t main_fn, int argc, char **argv, void* crate_map) {
     // Create the main scheduler and the main task
     rust_sched_id sched_id = kernel->create_scheduler(env->num_sched_threads);
     rust_scheduler *sched = kernel->get_scheduler_by_id(sched_id);
+    assert(sched != NULL);
     rust_task *root_task = sched->create_task(NULL, "main");
 
     // Build the command line arguments to pass to the root task
diff --git a/src/rt/rust_builtin.cpp b/src/rt/rust_builtin.cpp
index 04d11e020ea..a33d3cb90fe 100644
--- a/src/rt/rust_builtin.cpp
+++ b/src/rt/rust_builtin.cpp
@@ -630,7 +630,8 @@ extern "C" CDECL rust_task*
 rust_new_task_in_sched(rust_sched_id id) {
     rust_task *task = rust_get_current_task();
     rust_scheduler *sched = task->kernel->get_scheduler_by_id(id);
-    // FIXME (#2668): What if we didn't get the scheduler?
+    if (sched == NULL)
+        return NULL;
     return new_task_common(sched, task);
 }
 
diff --git a/src/rt/rust_kernel.cpp b/src/rt/rust_kernel.cpp
index b13b1490c0f..82fec98e86a 100644
--- a/src/rt/rust_kernel.cpp
+++ b/src/rt/rust_kernel.cpp
@@ -121,7 +121,9 @@ rust_kernel::get_scheduler_by_id(rust_sched_id id) {
 
 rust_scheduler *
 rust_kernel::get_scheduler_by_id_nolock(rust_sched_id id) {
-    assert(id != 0 && "invalid scheduler id");
+    if (id == 0) {
+        return NULL;
+    }
     sched_lock.must_have_lock();
     sched_map::iterator iter = sched_table.find(id);
     if (iter != sched_table.end()) {