From 952bb5c431cafc4412dae5c143457fd9d1c0107b Mon Sep 17 00:00:00 2001
From: Emil Fresk <emil.fresk@gmail.com>
Date: Sun, 12 Jun 2022 21:24:26 +0200
Subject: [PATCH] Restart executor on finish if there are retries

---
 macros/src/codegen/dispatchers.rs | 22 ++++++++++++++--------
 src/export.rs                     |  9 ++++++---
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/macros/src/codegen/dispatchers.rs b/macros/src/codegen/dispatchers.rs
index e6caa78196..1d1a4fb8af 100644
--- a/macros/src/codegen/dispatchers.rs
+++ b/macros/src/codegen/dispatchers.rs
@@ -164,12 +164,6 @@ pub fn codegen(app: &App, analysis: &Analysis, extra: &Extra) -> Vec<TokenStream
                     #(#arms)*
                 }
             }
-
-            while let Some((task, index)) = retry_queue.pop() {
-                rtic::export::interrupt::free(|_| {
-                    (&mut *#rq.get_mut()).enqueue_unchecked((task, index));
-                });
-            }
         ));
 
         for (name, _task) in app.software_tasks.iter().filter_map(|(name, task)| {
@@ -185,14 +179,26 @@ pub fn codegen(app: &App, analysis: &Analysis, extra: &Extra) -> Vec<TokenStream
             stmts.push(quote!(
                 if #executor_run_ident.load(core::sync::atomic::Ordering::Relaxed) {
                     #executor_run_ident.store(false, core::sync::atomic::Ordering::Relaxed);
-                    (&mut *#exec_name.get_mut()).poll(||  {
+                    if (&mut *#exec_name.get_mut()).poll(||  {
                         #executor_run_ident.store(true, core::sync::atomic::Ordering::Release);
                         rtic::pend(#device::#enum_::#interrupt);
-                    });
+                    }) && !retry_queue.is_empty() {
+                        // If the retry queue is not empty and the executor finished, restart this
+                        // dispatch to check if the executor should be restarted.
+                        rtic::pend(#device::#enum_::#interrupt);
+                    }
                 }
             ));
         }
 
+        stmts.push(quote!(
+            while let Some((task, index)) = retry_queue.pop() {
+                rtic::export::interrupt::free(|_| {
+                    (&mut *#rq.get_mut()).enqueue_unchecked((task, index));
+                });
+            }
+        ));
+
         let doc = format!("Interrupt handler to dispatch tasks at priority {}", level);
         let attribute = &interrupts[&level].1.attrs;
         items.push(quote!(
diff --git a/src/export.rs b/src/export.rs
index 703c0e04d8..9ef721f986 100644
--- a/src/export.rs
+++ b/src/export.rs
@@ -65,7 +65,7 @@ pub mod executor {
             self.task = Some(future);
         }
 
-        pub fn poll(&mut self, wake: fn()) {
+        pub fn poll(&mut self, wake: fn()) -> bool {
             if let Some(future) = &mut self.task {
                 unsafe {
                     let waker = Waker::from_raw(RawWaker::new(wake as *const (), &WAKER_VTABLE));
@@ -75,10 +75,13 @@ pub mod executor {
                     match future.poll(&mut cx) {
                         Poll::Ready(_) => {
                             self.task = None;
+                            true // Only true if we finished now
                         }
-                        Poll::Pending => {}
-                    };
+                        Poll::Pending => false,
+                    }
                 }
+            } else {
+                false
             }
         }
     }