Properly ref counting to fix valgrind issues on linux.

6 files changed, 38 insertions, 15 deletions

diff --git a/src/rt/memory_region.cpp b/src/rt/memory_region.cpp
index d9755c55888..7b38f65caaf 100644
--- a/src/rt/memory_region.cpp
+++ b/src/rt/memory_region.cpp
@@ -4,7 +4,7 @@
 // NB: please do not commit code with this uncommented. It's
 // hugely expensive and should only be used as a last resort.
 //
-#define TRACK_ALLOCATIONS
+// #define TRACK_ALLOCATIONS
 
 #define MAGIC 0xbadc0ffe
 
diff --git a/src/rt/rust.cpp b/src/rt/rust.cpp
index a94eb6c1a24..cc2f86da4b2 100644
--- a/src/rt/rust.cpp
+++ b/src/rt/rust.cpp
@@ -108,6 +108,8 @@ rust_start(uintptr_t main_fn, int argc, char **argv, void* crate_map) {
     }
 
     root_task->start(main_fn, (uintptr_t)args->args);
+    root_task->deref();
+    root_task = NULL;
 
     int ret = kernel->start_task_threads();
     delete args;
diff --git a/src/rt/rust_builtin.cpp b/src/rt/rust_builtin.cpp
index 77a4e8d2b09..ef236247396 100644
--- a/src/rt/rust_builtin.cpp
+++ b/src/rt/rust_builtin.cpp
@@ -303,8 +303,10 @@ task_join(rust_task *task, rust_task_id tid) {
         join_task->lock.unlock();
     }
     if (!join_task->failed) {
+        join_task->deref();
         return 0;
     } else {
+        join_task->deref();
         return -1;
     }
 }
@@ -728,16 +730,27 @@ get_task_context(rust_task *task, rust_task_id id) {
     return regs;
 }
 
+extern "C" CDECL void
+drop_task(rust_task *task, rust_task_id tid) {
+    rust_task *target = task->kernel->get_task_by_id(tid);
+    if(target) {
+        target->deref();
+        // Deref twice because get_task_by_id does once.
+        target->deref();
+    }
+}
+
 extern "C" CDECL rust_task *
 get_task_pointer(rust_task *task, rust_task_id id) {
-    return task->kernel->get_task_by_id(id);
+    rust_task *t = task->kernel->get_task_by_id(id);
+    return t;
 }
 
 extern "C" CDECL void
 start_task(rust_task *task, rust_task_id id) {
     rust_task * target = task->kernel->get_task_by_id(id);
-
     target->start();
+    target->deref();
 }
 
 extern "C" void *task_trampoline asm("task_trampoline");
@@ -754,6 +767,7 @@ migrate_alloc(rust_task *task, void *alloc, rust_task_id tid) {
     if(target) {
         task->local_region.release_alloc(alloc);
         target->local_region.claim_alloc(alloc);
+        target->deref();
     }
     else {
         // We couldn't find the target. Maybe we should just free?
@@ -849,6 +863,7 @@ chan_id_send(rust_task *task, type_desc *t, rust_task_id target_task_id,
         if(port) {
             port->remote_chan->send(sptr);
         }
+        target_task->deref();
     }
 }
 
diff --git a/src/rt/rust_kernel.cpp b/src/rt/rust_kernel.cpp
index 89c06885683..c11a5b69b73 100644
--- a/src/rt/rust_kernel.cpp
+++ b/src/rt/rust_kernel.cpp
@@ -160,6 +160,7 @@ 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();
     return task;
 }
 
diff --git a/src/rt/rust_upcall.cpp b/src/rt/rust_upcall.cpp
index 99d888cc3b1..dc406a83ea0 100644
--- a/src/rt/rust_upcall.cpp
+++ b/src/rt/rust_upcall.cpp
@@ -133,7 +133,9 @@ upcall_clone_chan(rust_task *task, rust_task_id tid,
     // FIXME: This should be removed.
     LOG_UPCALL_ENTRY(task);
     rust_task *target = task->kernel->get_task_by_id(tid);
-    return chan->clone(target);
+    rust_chan *c = chan->clone(target);
+    target->deref();
+    return c;
 }
 
 extern "C" CDECL rust_task *
@@ -197,6 +199,7 @@ upcall_kill(rust_task *task, rust_task_id tid) {
     LOG_UPCALL_ENTRY(task);
     rust_task *target = task->kernel->get_task_by_id(tid);
     target->kill();
+    target->deref();
 }
 
 /**
@@ -315,7 +318,9 @@ extern "C" CDECL rust_str *
 upcall_dup_str(rust_task *task, rust_task_id tid, rust_str *str) {
     LOG_UPCALL_ENTRY(task);
     rust_task *target = task->kernel->get_task_by_id(tid);
-    return make_str(target, (char const *)str->data, str->fill);
+    rust_str *s = make_str(target, (char const *)str->data, str->fill);
+    target->deref();
+    return s;
 }
 
 extern "C" CDECL rust_vec *
@@ -478,27 +483,25 @@ upcall_new_task(rust_task *spawner, rust_vec *name) {
     LOG_UPCALL_ENTRY(spawner);
     rust_task_id tid =
         spawner->kernel->create_task(spawner, (const char *)name->data);
-    rust_task *task = spawner->kernel->get_task_by_id(tid);
-    task->ref();
+    // let the kernel bump the ref count.
+    spawner->kernel->get_task_by_id(tid);
+    // no deref because we're letting the caller take ownership.
     return tid;
 }
 
 extern "C" CDECL void
 upcall_take_task(rust_task *task, rust_task_id tid) {
     LOG_UPCALL_ENTRY(task);
-    rust_task *target = task->kernel->get_task_by_id(tid);
-    if(target) {
-        target->ref();
-    }
+    // get_task_by_id increments the refcount.
+    task->kernel->get_task_by_id(tid);
 }
 
 extern "C" CDECL void
+drop_task(rust_task *task, rust_task_id tid);
+extern "C" CDECL void
 upcall_drop_task(rust_task *task, rust_task_id tid) {
     LOG_UPCALL_ENTRY(task);
-    rust_task *target = task->kernel->get_task_by_id(tid);
-    if(target) {
-        target->deref();
-    }
+    drop_task(task, tid);
 }
 
 extern "C" CDECL void
@@ -544,6 +547,7 @@ upcall_start_task(rust_task *spawner,
     memcpy((void*)task->rust_sp, (void*)args, args_sz);
 
     task->start(spawnee_fn, child_arg);
+    task->deref();
     return task;
 }
 
diff --git a/src/rt/rustrt.def.in b/src/rt/rustrt.def.in
index db31ef68067..ccc190077d8 100644
--- a/src/rt/rustrt.def.in
+++ b/src/rt/rustrt.def.in
@@ -26,6 +26,7 @@ debug_tydesc
 do_gc
 drop_chan
 drop_port
+drop_task
 get_port_id
 get_task_context
 get_task_id