From 7277cd71984e6a09bc2f8a8a828a5e213b485d00 Mon Sep 17 00:00:00 2001
From: Brian Anderson <banderson@mozilla.com>
Date: Tue, 15 May 2012 11:34:52 -0700
Subject: [PATCH] core: Add task::unkillable

---
 src/libcore/task.rs     | 63 ++++++++++++++++++++++++++++++++++++++++-
 src/rt/rust_builtin.cpp | 12 ++++++++
 src/rt/rust_task.cpp    | 13 ++++++++-
 src/rt/rust_task.h      |  4 +++
 src/rt/rustrt.def.in    |  2 ++
 5 files changed, 92 insertions(+), 2 deletions(-)

diff --git a/src/libcore/task.rs b/src/libcore/task.rs
index 2c91b0de0119..b9a07a7625b3 100644
--- a/src/libcore/task.rs
+++ b/src/libcore/task.rs
@@ -51,7 +51,7 @@
 export yield;
 export failing;
 export get_task;
-
+export unkillable;
 
 /* Data types */
 
@@ -467,6 +467,29 @@ fn get_task() -> task {
     task(rustrt::get_task_id())
 }
 
+#[doc = "
+Temporarily make the task unkillable
+
+# Example
+
+    task::unkillable {||
+        // detach / yield / destroy must all be called together
+        rustrt::rust_port_detach(po);
+        // This must not result in the current task being killed
+        task::yield();
+        rustrt::rust_port_destroy(po);
+    }
+
+"]
+unsafe fn unkillable(f: fn()) {
+    resource allow_failure(_i: ()) {
+        rustrt::rust_task_allow_kill();
+    }
+    let _allow_failure = allow_failure(());
+    rustrt::rust_task_inhibit_kill();
+    f();
+}
+
 
 /* Internal */
 
@@ -566,6 +589,8 @@ fn rust_task_config_notify(
     fn rust_task_is_unwinding(rt: *rust_task) -> bool;
     fn unsupervise();
     fn rust_osmain_sched_id() -> sched_id;
+    fn rust_task_inhibit_kill();
+    fn rust_task_allow_kill();
 }
 
 
@@ -930,3 +955,39 @@ fn test_osmain() {
     }
     comm::recv(po);
 }
+
+#[test]
+#[ignore(cfg(target_os = "win32"))]
+#[should_fail]
+fn test_unkillable() unsafe {
+    import comm::methods;
+    let po = comm::port();
+    let ch = po.chan();
+
+    // We want to do this after failing
+    spawn {||
+        iter::repeat(10u, yield);
+        ch.send(());
+    }
+
+    spawn {||
+        yield();
+        // We want to fail after the unkillable task
+        // blocks on recv
+        fail;
+    }
+
+    unkillable {||
+        let p = ~0;
+        let pp: *uint = unsafe::reinterpret_cast(p);
+        unsafe::forget(p);
+
+        // If we are killed here then the box will leak
+        po.recv();
+
+        let _p: ~int = unsafe::reinterpret_cast(pp);
+    }
+
+    // Now we can be killed
+    po.recv();
+}
diff --git a/src/rt/rust_builtin.cpp b/src/rt/rust_builtin.cpp
index 4ba6c299c62d..8f5c3c6fa14a 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 d7c62780fbb1..823937443b5d 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 5a841f00f591..1867c8f4ed81 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 ae6d0521bef8..345cd5f7a2e7 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