From 144411ce25afd95b193d02b2d54c1c55158e1afb Mon Sep 17 00:00:00 2001
From: Luca Versari <veluca93@gmail.com>
Date: Sun, 7 Dec 2025 23:02:04 +0100
Subject: [PATCH] Give names to store/flash/register tasks.

This is done by implementing a JoinHandle for Executor::spawn.
---
 pixie-uefi/src/main.rs        |  8 +++++---
 pixie-uefi/src/os/executor.rs | 32 +++++++++++++++++++++++++++-----
 pixie-uefi/src/os/net/mod.rs  |  2 +-
 3 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/pixie-uefi/src/main.rs b/pixie-uefi/src/main.rs
index 6019ba27..8b58c964 100644
--- a/pixie-uefi/src/main.rs
+++ b/pixie-uefi/src/main.rs
@@ -145,9 +145,11 @@ async fn run() -> Result<()> {
                     Action::Boot => power_control::reboot_to_os().await,
                     Action::Restart => {}
                     Action::Shutdown => shutdown().await,
-                    Action::Register => register(server).await?,
-                    Action::Store => store(server).await?,
-                    Action::Flash => flash(server).await?,
+                    Action::Register => {
+                        Executor::spawn("register", register(server)).join().await?
+                    }
+                    Action::Store => Executor::spawn("store", store(server)).join().await?,
+                    Action::Flash => Executor::spawn("flash", flash(server)).join().await?,
                 }
 
                 let tcp = TcpStream::connect(server).await?;
diff --git a/pixie-uefi/src/os/executor.rs b/pixie-uefi/src/os/executor.rs
index 1a3689b4..76eb25c8 100644
--- a/pixie-uefi/src/os/executor.rs
+++ b/pixie-uefi/src/os/executor.rs
@@ -8,6 +8,7 @@ use core::future::{poll_fn, Future};
 use core::pin::Pin;
 use core::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 use core::task::{Context, Poll, Waker};
+use futures::channel::oneshot;
 
 use spin::Mutex;
 use uefi::boot::{EventType, TimerTrigger, Tpl};
@@ -25,6 +26,7 @@ struct Task {
     future: Mutex<BoxFuture>,
     micros: AtomicU64,
     last_micros: AtomicU64,
+    done: AtomicBool,
 }
 
 impl Task {
@@ -38,13 +40,14 @@ impl Task {
             micros: AtomicU64::new(0),
             last_micros: AtomicU64::new(0),
             in_queue: AtomicBool::new(false),
+            done: AtomicBool::new(false),
         })
     }
 }
 
 impl Wake for Task {
     fn wake(self: Arc<Self>) {
-        if !self.in_queue.swap(true, Ordering::Relaxed) {
+        if !self.in_queue.swap(true, Ordering::Relaxed) && !self.done.load(Ordering::Relaxed) {
             EXECUTOR.lock().ready_tasks.push_back(self);
         }
     }
@@ -55,6 +58,14 @@ static EXECUTOR: Mutex<Executor> = Mutex::new(Executor {
     tasks: vec![],
 });
 
+pub struct JoinHandle<T>(oneshot::Receiver<T>);
+
+impl<T> JoinHandle<T> {
+    pub async fn join(self) -> T {
+        self.0.await.expect("tasks should never be cancelled")
+    }
+}
+
 pub struct Executor {
     // TODO(veluca): scheduling.
     ready_tasks: VecDeque<Arc<Task>>,
@@ -145,6 +156,9 @@ impl Executor {
                     t.last_micros
                         .store(t.micros.load(Ordering::Relaxed), Ordering::Relaxed);
                 }
+
+                // Clear completed tasks.
+                tasks.retain(|t| !t.done.load(Ordering::Relaxed));
             }
             last = Timer::micros() as u64;
             Self::sleep_us(1_000_000).await;
@@ -165,10 +179,13 @@ impl Executor {
             let mut context = Context::from_waker(&waker);
             let mut fut = task.future.try_lock().unwrap();
             let begin = Timer::micros();
-            let _ = fut.0.as_mut().poll(&mut context);
+            let done = fut.0.as_mut().poll(&mut context);
             let end = Timer::micros();
             task.micros
                 .fetch_add((end - begin) as u64, Ordering::Relaxed);
+            if done.is_ready() {
+                task.done.swap(true, Ordering::Relaxed);
+            }
         }
     }
 
@@ -211,13 +228,18 @@ impl Executor {
     }
 
     /// Spawn a new task.
-    pub fn spawn<Fut>(name: &'static str, f: Fut)
+    pub fn spawn<Fut, T: 'static>(name: &'static str, f: Fut) -> JoinHandle<T>
     where
-        Fut: Future<Output = ()> + 'static,
+        Fut: Future<Output = T> + 'static,
     {
-        let task = Task::new(name, f);
+        let (send, recv) = oneshot::channel();
+        let task = Task::new(name, async move {
+            let t = f.await;
+            let _ = send.send(t);
+        });
         let mut executor = EXECUTOR.lock();
         executor.tasks.push(task.clone());
         executor.ready_tasks.push_back(task);
+        JoinHandle(recv)
     }
 }
diff --git a/pixie-uefi/src/os/net/mod.rs b/pixie-uefi/src/os/net/mod.rs
index d4edd68d..04be355d 100644
--- a/pixie-uefi/src/os/net/mod.rs
+++ b/pixie-uefi/src/os/net/mod.rs
@@ -124,7 +124,7 @@ pub(super) fn init() {
             poll();
             // TODO(veluca): figure out whether we can suspend the task.
             cx.waker().wake_by_ref();
-            Poll::Pending
+            Poll::<()>::Pending
         }),
     );