2 files changed, 6 insertions, 13 deletions

diff --git a/embassy-executor/Cargo.toml b/embassy-executor/Cargo.toml
index 910d6fa37..45b0955bf 100644
--- a/embassy-executor/Cargo.toml
+++ b/embassy-executor/Cargo.toml
@@ -31,6 +31,7 @@ flavors = [
 default = []
 std = ["embassy-macros/std", "critical-section/std"]
 wasm = ["dep:wasm-bindgen", "dep:js-sys", "embassy-macros/wasm"]
+riscv = ["embassy-macros/riscv"]
 
 # Enable nightly-only features
 nightly = []
diff --git a/embassy-executor/src/arch/riscv32.rs b/embassy-executor/src/arch/riscv32.rs
index 2a4b006da..76eb8b114 100644
--- a/embassy-executor/src/arch/riscv32.rs
+++ b/embassy-executor/src/arch/riscv32.rs
@@ -55,19 +55,11 @@ impl Executor {
             unsafe {
                 self.inner.poll();
                 // we do not care about race conditions between the load and store operations, interrupts
-                //will only set this value to true.
-                critical_section::with(|_| {
-                    // if there is work to do, loop back to polling
-                    // TODO can we relax this?
-                    if SIGNAL_WORK_THREAD_MODE.load(Ordering::SeqCst) {
-                        SIGNAL_WORK_THREAD_MODE.store(false, Ordering::SeqCst);
-                    }
-                    // if not, wait for interrupt
-                    else {
-                        core::arch::asm!("wfi");
-                    }
-                });
-                // if an interrupt occurred while waiting, it will be serviced here
+                // will only set this value to true.
+                // if there is work to do, loop back to polling
+                if !SIGNAL_WORK_THREAD_MODE.fetch_and(false, Ordering::SeqCst) {
+                    core::arch::asm!("wfi");
+                }
             }
         }
     }