Reducing the chances for race conditions in join.

4 files changed, 20 insertions, 3 deletions

diff --git a/src/rt/rust_internal.h b/src/rt/rust_internal.h
index 517e869ffee..5045b395477 100644
--- a/src/rt/rust_internal.h
+++ b/src/rt/rust_internal.h
@@ -115,7 +115,10 @@ static size_t const BUF_BYTES = 2048;
     private:                                                            \
     intptr_t ref_count;                                                 \
 public:                                                                 \
- void ref() { sync::increment(ref_count); }                             \
+ void ref() {                                                           \
+     intptr_t old = sync::increment(ref_count);                         \
+     assert(old > 0);                                                   \
+ }                                                                      \
  void deref() { if(0 == sync::decrement(ref_count)) { delete this; } }
 
 template <typename T> struct rc_base {
diff --git a/src/rt/rust_kernel.cpp b/src/rt/rust_kernel.cpp
index c11a5b69b73..ceddcf5b1c9 100644
--- a/src/rt/rust_kernel.cpp
+++ b/src/rt/rust_kernel.cpp
@@ -160,7 +160,17 @@ rust_kernel::get_task_by_id(rust_task_id id) {
     rust_task *task = NULL;
     // get leaves task unchanged if not found.
     task_table.get(id, &task);
-    if(task) task->ref();
+    if(task) {
+        if(task->get_ref_count() == 0) {
+            // this means the destructor is running, since the destructor
+            // grabs the kernel lock to unregister the task. Pretend this
+            // doesn't actually exist.
+            return NULL;
+        }
+        else {
+            task->ref();
+        }
+    }
     return task;
 }
 
diff --git a/src/rt/rust_task.h b/src/rt/rust_task.h
index 5d920c51567..c757a9e69e1 100644
--- a/src/rt/rust_task.h
+++ b/src/rt/rust_task.h
@@ -169,6 +169,10 @@ rust_task : public kernel_owned<rust_task>, rust_cond
     rust_port_id register_port(rust_port *port);
     void release_port(rust_port_id id);
     rust_port *get_port_by_id(rust_port_id id);
+
+    // Use this function sparingly. Depending on the ref count is generally
+    // not at all safe.
+    intptr_t get_ref_count() const { return ref_count; }
 };
 
 //
diff --git a/src/test/bench/task-perf-spawnalot.rs b/src/test/bench/task-perf-spawnalot.rs
index bb0455e0d21..b7d9e07022e 100644
--- a/src/test/bench/task-perf-spawnalot.rs
+++ b/src/test/bench/task-perf-spawnalot.rs
@@ -27,4 +27,4 @@ fn main(args: vec[str]) {
         task::_spawn(bind f(n));
         i += 1u;
     }
-}
\ No newline at end of file
+}