core: Add task::unkillable

4 files changed, 30 insertions, 1 deletions

diff --git a/src/rt/rust_builtin.cpp b/src/rt/rust_builtin.cpp
index 4ba6c299c62..8f5c3c6fa14 100644
--- a/src/rt/rust_builtin.cpp
+++ b/src/rt/rust_builtin.cpp
@@ -805,6 +805,18 @@ rust_global_env_chan_ptr() {
     return task->kernel->get_global_env_chan();
 }
 
+extern "C" void
+rust_task_inhibit_kill() {
+    rust_task *task = rust_get_current_task();
+    task->inhibit_kill();
+}
+
+extern "C" void
+rust_task_allow_kill() {
+    rust_task *task = rust_get_current_task();
+    task->allow_kill();
+}
+
 //
 // Local Variables:
 // mode: C++
diff --git a/src/rt/rust_task.cpp b/src/rt/rust_task.cpp
index d7c62780fbb..823937443b5 100644
--- a/src/rt/rust_task.cpp
+++ b/src/rt/rust_task.cpp
@@ -37,6 +37,7 @@ rust_task::rust_task(rust_sched_loop *sched_loop, rust_task_state state,
     cond_name("none"),
     killed(false),
     reentered_rust_stack(false),
+    disallow_kill(false),
     c_stack(NULL),
     next_c_sp(0),
     next_rust_sp(0),
@@ -211,7 +212,7 @@ rust_task::must_fail_from_being_killed() {
 bool
 rust_task::must_fail_from_being_killed_unlocked() {
     kill_lock.must_have_lock();
-    return killed && !reentered_rust_stack;
+    return killed && !reentered_rust_stack && !disallow_kill;
 }
 
 // Only run this on the rust stack
@@ -645,6 +646,16 @@ rust_task::on_rust_stack() {
     }
 }
 
+void
+rust_task::inhibit_kill() {
+    disallow_kill = true;
+}
+
+void
+rust_task::allow_kill() {
+    disallow_kill = false;
+}
+
 //
 // Local Variables:
 // mode: C++
diff --git a/src/rt/rust_task.h b/src/rt/rust_task.h
index 5a841f00f59..1867c8f4ed8 100644
--- a/src/rt/rust_task.h
+++ b/src/rt/rust_task.h
@@ -155,6 +155,7 @@ private:
     bool killed;
     // Indicates that we've called back into Rust from C
     bool reentered_rust_stack;
+    bool disallow_kill;
 
     // The stack used for running C code, borrowed from the scheduler thread
     stk_seg *c_stack;
@@ -268,6 +269,9 @@ public:
     const char *get_cond_name() { return cond_name; }
 
     void cleanup_after_turn();
+
+    void inhibit_kill();
+    void allow_kill();
 };
 
 // FIXME: It would be really nice to be able to get rid of this.
diff --git a/src/rt/rustrt.def.in b/src/rt/rustrt.def.in
index ae6d0521bef..345cd5f7a2e 100644
--- a/src/rt/rustrt.def.in
+++ b/src/rt/rustrt.def.in
@@ -152,3 +152,5 @@ rust_global_env_chan_ptr
 rust_port_take
 rust_port_drop
 rust_port_task
+rust_task_inhibit_kill
+rust_task_allow_kill
\ No newline at end of file