1 files changed, 13 insertions, 5 deletions

diff --git a/embassy-executor/src/arch/riscv32.rs b/embassy-executor/src/arch/riscv32.rs
index 76eb8b114..2a4b006da 100644
--- a/embassy-executor/src/arch/riscv32.rs
+++ b/embassy-executor/src/arch/riscv32.rs
@@ -55,11 +55,19 @@ 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.
-                // if there is work to do, loop back to polling
-                if !SIGNAL_WORK_THREAD_MODE.fetch_and(false, Ordering::SeqCst) {
-                    core::arch::asm!("wfi");
-                }
+                //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
             }
         }
     }