rt: Add rust_task::call_on_c_stack

3 files changed, 24 insertions, 7 deletions

diff --git a/src/rt/rust_task.cpp b/src/rt/rust_task.cpp
index 621eb1e54bc..571766c0c2a 100644
--- a/src/rt/rust_task.cpp
+++ b/src/rt/rust_task.cpp
@@ -194,7 +194,7 @@ void task_start_wrapper(spawn_args *a)
 
     // The cleanup work needs lots of stack
     cleanup_args ca = {a, threw_exception};
-    task->thread->c_context.call_and_change_stacks(&ca, (void*)cleanup_task);
+    task->call_on_c_stack(&ca, (void*)cleanup_task);
 
     task->ctx.next->swap(task->ctx);
 }
@@ -699,7 +699,18 @@ Returns true if we're currently running on the Rust stack
  */
 bool
 rust_task::on_rust_stack() {
-    return sp_in_stk_seg(get_sp(), stk);
+    uintptr_t sp = get_sp();
+    bool in_first_segment = sp_in_stk_seg(sp, stk);
+    if (in_first_segment) {
+        return true;
+    } else if (stk->next != NULL) {
+        // This happens only when calling the upcall to delete
+        // a stack segment
+        bool in_second_segment = sp_in_stk_seg(sp, stk->next);
+        return in_second_segment;
+    } else {
+        return false;
+    }
 }
 
 void
@@ -713,6 +724,12 @@ rust_task::config_notify(chan_handle chan) {
     notify_chan = chan;
 }
 
+void
+rust_task::call_on_c_stack(void *args, void *fn_ptr) {
+    I(thread, on_rust_stack());
+    thread->c_context.call_and_change_stacks(args, fn_ptr);
+}
+
 //
 // Local Variables:
 // mode: C++
diff --git a/src/rt/rust_task.h b/src/rt/rust_task.h
index 7f270f88826..c4939d208c1 100644
--- a/src/rt/rust_task.h
+++ b/src/rt/rust_task.h
@@ -181,6 +181,8 @@ public:
     void check_stack_canary();
 
     void config_notify(chan_handle chan);
+
+    void call_on_c_stack(void *args, void *fn_ptr);
 };
 
 //
diff --git a/src/rt/rust_upcall.cpp b/src/rt/rust_upcall.cpp
index c66b4370713..d12e7311303 100644
--- a/src/rt/rust_upcall.cpp
+++ b/src/rt/rust_upcall.cpp
@@ -47,8 +47,7 @@ inline void
 call_upcall_on_c_stack(void *args, void *fn_ptr) {
     check_stack_alignment();
     rust_task *task = rust_task_thread::get_task();
-    rust_task_thread *thread = task->thread;
-    thread->c_context.call_and_change_stacks(args, fn_ptr);
+    task->call_on_c_stack(args, fn_ptr);
 }
 
 extern "C" void record_sp(void *limit);
@@ -69,11 +68,10 @@ upcall_call_shim_on_c_stack(void *args, void *fn_ptr) {
     // stack.
     record_sp(0);
 
-    rust_task_thread *thread = task->thread;
     try {
-        thread->c_context.call_and_change_stacks(args, fn_ptr);
+        task->call_on_c_stack(args, fn_ptr);
     } catch (...) {
-        A(thread, false, "Native code threw an exception");
+        A(task->thread, false, "Native code threw an exception");
     }
 
     task = rust_task_thread::get_task();