Add 'do atomically { .. }' for exclusives

4 files changed, 38 insertions, 4 deletions

diff --git a/src/rt/rust_builtin.cpp b/src/rt/rust_builtin.cpp
index f86e9232625..784df64d944 100644
--- a/src/rt/rust_builtin.cpp
+++ b/src/rt/rust_builtin.cpp
@@ -854,18 +854,26 @@ rust_global_env_chan_ptr() {
 }
 
 extern "C" void
-rust_task_inhibit_kill() {
-    rust_task *task = rust_get_current_task();
+rust_task_inhibit_kill(rust_task *task) {
     task->inhibit_kill();
 }
 
 extern "C" void
-rust_task_allow_kill() {
-    rust_task *task = rust_get_current_task();
+rust_task_allow_kill(rust_task *task) {
     task->allow_kill();
 }
 
 extern "C" void
+rust_task_inhibit_yield(rust_task *task) {
+    task->inhibit_yield();
+}
+
+extern "C" void
+rust_task_allow_yield(rust_task *task) {
+    task->allow_yield();
+}
+
+extern "C" void
 rust_task_kill_other(rust_task *task) { /* Used for linked failure */
     task->kill();
 }
diff --git a/src/rt/rust_task.cpp b/src/rt/rust_task.cpp
index c28e3350bf2..e9879525ee7 100644
--- a/src/rt/rust_task.cpp
+++ b/src/rt/rust_task.cpp
@@ -39,6 +39,7 @@ rust_task::rust_task(rust_sched_loop *sched_loop, rust_task_state state,
     killed(false),
     reentered_rust_stack(false),
     disallow_kill(0),
+    disallow_yield(0),
     c_stack(NULL),
     next_c_sp(0),
     next_rust_sp(0)
@@ -234,9 +235,18 @@ rust_task::must_fail_from_being_killed_inner() {
     return killed && !reentered_rust_stack && disallow_kill == 0;
 }
 
+void rust_task_yield_fail(rust_task *task) {
+    LOG_ERR(task, task, "task %" PRIxPTR " yielded in an atomic section",
+            task);
+    task->fail();
+}
+
 // Only run this on the rust stack
 void
 rust_task::yield(bool *killed) {
+    if (disallow_yield > 0) {
+        call_on_c_stack(this, (void *)rust_task_yield_fail);
+    }
     // FIXME (#2875): clean this up
     if (must_fail_from_being_killed()) {
         {
@@ -672,6 +682,17 @@ rust_task::allow_kill() {
     disallow_kill--;
 }
 
+void rust_task::inhibit_yield() {
+    scoped_lock with(lifecycle_lock);
+    disallow_yield++;
+}
+
+void rust_task::allow_yield() {
+    scoped_lock with(lifecycle_lock);
+    assert(disallow_yield > 0 && "Illegal allow_yield(): already yieldable!");
+    disallow_yield--;
+}
+
 void *
 rust_task::wait_event(bool *killed) {
     scoped_lock with(lifecycle_lock);
diff --git a/src/rt/rust_task.h b/src/rt/rust_task.h
index 2db882d11e3..985e1ebe6e8 100644
--- a/src/rt/rust_task.h
+++ b/src/rt/rust_task.h
@@ -185,6 +185,7 @@ private:
     // Indicates that we've called back into Rust from C
     bool reentered_rust_stack;
     unsigned long disallow_kill;
+    unsigned long disallow_yield;
 
     // The stack used for running C code, borrowed from the scheduler thread
     stk_seg *c_stack;
@@ -318,6 +319,8 @@ public:
 
     void inhibit_kill();
     void allow_kill();
+    void inhibit_yield();
+    void allow_yield();
 };
 
 // FIXME (#2697): 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 2079db61466..62fb8482ad5 100644
--- a/src/rt/rustrt.def.in
+++ b/src/rt/rustrt.def.in
@@ -183,6 +183,8 @@ rust_port_drop
 rust_port_task
 rust_task_inhibit_kill
 rust_task_allow_kill
+rust_task_inhibit_yield
+rust_task_allow_yield
 rust_task_kill_other
 rust_task_kill_all
 rust_create_cond_lock