core: Add new weak task API

4 files changed, 34 insertions, 8 deletions

diff --git a/src/rt/rust_builtin.cpp b/src/rt/rust_builtin.cpp
index 221afb89b23..a5e1260d4a5 100644
--- a/src/rt/rust_builtin.cpp
+++ b/src/rt/rust_builtin.cpp
@@ -1038,6 +1038,18 @@ rust_get_global_data_ptr() {
     return &task->kernel->global_data;
 }
 
+extern "C" void
+rust_inc_weak_task_count() {
+    rust_task *task = rust_get_current_task();
+    task->kernel->inc_weak_task_count();
+}
+
+extern "C" void
+rust_dec_weak_task_count() {
+    rust_task *task = rust_get_current_task();
+    task->kernel->dec_weak_task_count();
+}
+
 //
 // Local Variables:
 // mode: C++
diff --git a/src/rt/rust_kernel.cpp b/src/rt/rust_kernel.cpp
index 9c6ba9dcda3..d270ac07633 100644
--- a/src/rt/rust_kernel.cpp
+++ b/src/rt/rust_kernel.cpp
@@ -377,17 +377,12 @@ rust_kernel::weaken_task(rust_port_id chan) {
         KLOG_("Weakening task with channel %" PRIdPTR, chan);
         weak_task_chans.push_back(chan);
     }
-    uintptr_t new_non_weak_tasks = sync::decrement(non_weak_tasks);
-    KLOG_("New non-weak tasks %" PRIdPTR, new_non_weak_tasks);
-    if (new_non_weak_tasks == 0) {
-        begin_shutdown();
-    }
+    inc_weak_task_count();
 }
 
 void
 rust_kernel::unweaken_task(rust_port_id chan) {
-    uintptr_t new_non_weak_tasks = sync::increment(non_weak_tasks);
-    KLOG_("New non-weak tasks %" PRIdPTR, new_non_weak_tasks);
+    dec_weak_task_count();
     {
         scoped_lock with(weak_task_lock);
         KLOG_("Unweakening task with channel %" PRIdPTR, chan);
@@ -400,6 +395,21 @@ rust_kernel::unweaken_task(rust_port_id chan) {
 }
 
 void
+rust_kernel::inc_weak_task_count() {
+    uintptr_t new_non_weak_tasks = sync::decrement(non_weak_tasks);
+    KLOG_("New non-weak tasks %" PRIdPTR, new_non_weak_tasks);
+    if (new_non_weak_tasks == 0) {
+        begin_shutdown();
+    }
+}
+
+void
+rust_kernel::dec_weak_task_count() {
+    uintptr_t new_non_weak_tasks = sync::increment(non_weak_tasks);
+    KLOG_("New non-weak tasks %" PRIdPTR, new_non_weak_tasks);
+}
+
+void
 rust_kernel::end_weak_tasks() {
     std::vector<rust_port_id> chancopies;
     {
diff --git a/src/rt/rust_kernel.h b/src/rt/rust_kernel.h
index 99b230f7872..f90ecf01a7b 100644
--- a/src/rt/rust_kernel.h
+++ b/src/rt/rust_kernel.h
@@ -187,6 +187,8 @@ public:
     void unregister_task();
     void weaken_task(rust_port_id chan);
     void unweaken_task(rust_port_id chan);
+    void inc_weak_task_count();
+    void dec_weak_task_count();
 
     bool send_to_port(rust_port_id chan, void *sptr);
 
diff --git a/src/rt/rustrt.def.in b/src/rt/rustrt.def.in
index 8c26832f349..5be823d8fde 100644
--- a/src/rt/rustrt.def.in
+++ b/src/rt/rustrt.def.in
@@ -211,4 +211,6 @@ linenoiseHistoryLoad
 rust_raw_thread_start
 rust_raw_thread_join_delete
 rust_register_exit_function
-rust_get_global_data_ptr
\ No newline at end of file
+rust_get_global_data_ptr
+rust_inc_weak_task_count
+rust_dec_weak_task_count
\ No newline at end of file