From b8ba7ea92948d1013d38798f18584b788afc5573 Mon Sep 17 00:00:00 2001
From: mwojcik <mw@m-labs.hk>
Date: Mon, 27 Mar 2023 11:17:55 +0800
Subject: [PATCH] dma: fix nested block_on panic

---
 src/runtime/src/comms.rs    | 2 +-
 src/runtime/src/rtio_dma.rs | 9 +++++----
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/src/runtime/src/comms.rs b/src/runtime/src/comms.rs
index 89eab8f..814c693 100644
--- a/src/runtime/src/comms.rs
+++ b/src/runtime/src/comms.rs
@@ -329,7 +329,7 @@ async fn handle_run_kernel(
             }
             kernel::Message::DmaEraseRequest(name) => {
                 // prevent possible OOM when we have large DMA record replacement.
-                rtio_dma::erase(name, aux_mutex, routing_table, timer);
+                rtio_dma::erase(name, aux_mutex, routing_table, timer).await;
             }
             kernel::Message::DmaGetRequest(name) => {
                 let result = rtio_dma::retrieve(name);
diff --git a/src/runtime/src/rtio_dma.rs b/src/runtime/src/rtio_dma.rs
index a3a0da6..01f3c2d 100644
--- a/src/runtime/src/rtio_dma.rs
+++ b/src/runtime/src/rtio_dma.rs
@@ -235,14 +235,14 @@ pub mod remote_dma {
         trace_set.await_done(timeout, timer).await
     }
 
-    pub fn erase(
+    pub async fn erase(
         aux_mutex: &Rc<Mutex<bool>>, 
         routing_table: &RoutingTable, 
         timer: GlobalTimer,
         id: u32
     ) {
         let trace_set = unsafe { TRACES.get_mut(&id).unwrap() };
-        task::block_on(trace_set.erase(aux_mutex, routing_table, timer));
+        trace_set.erase(aux_mutex, routing_table, timer).await;
         unsafe { TRACES.remove(&id); }
     }
 
@@ -304,6 +304,7 @@ pub fn put_record(mut recorder: DmaRecorder) -> u32 {
         // sends whole chunks, to limit comms/kernel CPU communication,
         // and as only comms core has access to varios DMA buffers.
         let mut ptr = 0;
+        recorder.buffer.push(0);
         while recorder.buffer[ptr] != 0 {
             // ptr + 3 = tgt >> 24 (destination)
             let len = recorder.buffer[ptr] as usize;
@@ -347,13 +348,13 @@ pub fn put_record(mut recorder: DmaRecorder) -> u32 {
     ptr
 }
 
-pub fn erase(name: String, _aux_mutex: &Rc<Mutex<bool>>, 
+pub async fn erase(name: String, _aux_mutex: &Rc<Mutex<bool>>, 
     _routing_table: &RoutingTable, _timer: GlobalTimer
 ) {
     let _entry = DMA_RECORD_STORE.lock().remove(&name);
     #[cfg(has_drtio)]
     if let Some((id, _v, _d)) = _entry {
-        remote_dma::erase(_aux_mutex, _routing_table, _timer, id);
+        remote_dma::erase(_aux_mutex, _routing_table, _timer, id).await;
     }
 }