Add Event

Author: Virv12

pixie-uefi/src/os/executor/event.rs

@@ -0,0 +1,57 @@
+use alloc::sync::{Arc, Weak};
+use core::future::Future;
+use core::pin::Pin;
+use core::sync::atomic::{AtomicBool, Ordering};
+use core::task::{Context, Poll};
+
+use futures::task::AtomicWaker;
+
+#[derive(Debug, Default)]
+struct Inner {
+    triggered: AtomicBool,
+    waker: AtomicWaker,
+}
+
+pub struct Event {
+    inner: Arc<Inner>,
+}
+
+impl Event {
+    pub fn new() -> Self {
+        Self {
+            inner: Default::default(),
+        }
+    }
+
+    pub fn trigger(&self) -> EventTrigger {
+        EventTrigger {
+            inner: Arc::downgrade(&self.inner),
+        }
+    }
+}
+
+impl Future for Event {
+    type Output = ();
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
+        if self.inner.triggered.load(Ordering::Relaxed) {
+            Poll::Ready(())
+        } else {
+            self.inner.waker.register(cx.waker());
+            Poll::Pending
+        }
+    }
+}
+
+pub struct EventTrigger {
+    inner: Weak<Inner>,
+}
+
+impl EventTrigger {
+    pub fn trigger(&self) {
+        if let Some(inner) = self.inner.upgrade() {
+            inner.triggered.store(true, Ordering::Relaxed);
+            inner.waker.wake();
+        }
+    }
+}

pixie-uefi/src/os/executor.rs pixie-uefi/src/os/executor/mod.rspixie-uefi/src/os/executor.rs renamed to pixie-uefi/src/os/executor/mod.rs

@@ -12,14 +12,16 @@ use core::task::{Context, Poll, Waker};
 use core::time::Duration;
 
 use futures::channel::oneshot;
-use futures::task::AtomicWaker;
 use spin::Mutex;
 use uefi::proto::console::text::Color;
 
+use crate::os::executor::event::{Event, EventTrigger};
 use crate::os::send_wrapper::SendWrapper;
 use crate::os::timer::Timer;
 use crate::os::ui::DrawArea;
 
+pub mod event;
+
 type BoxFuture = SendWrapper<Pin<Box<dyn Future<Output = ()> + 'static>>>;
 
 struct Task {
@@ -55,11 +57,9 @@ impl Wake for Task {
     }
 }
 
-pub(super) type WrappedWaker = Arc<AtomicWaker>;
-
 struct TimedWait {
     wake_at: i64,
-    waker: WrappedWaker,
+    event: EventTrigger,
 }
 
 impl PartialEq for TimedWait {
@@ -99,7 +99,7 @@ static EXECUTOR: Mutex<Executor> = Mutex::new(Executor {
 });
 
 pub struct Executor {
-    wake_on_interrupt: Vec<Waker>,
+    wake_on_interrupt: Vec<EventTrigger>,
     timed_wait: BinaryHeap<TimedWait>,
     ready_tasks: VecDeque<Arc<Task>>,
     tasks: Vec<Arc<Task>>,
@@ -209,7 +209,7 @@ impl Executor {
         let mut do_wake = |force_interrupt_wake| {
             // Wake timed-waiting tasks.
             loop {
-                let waker = {
+                let event = {
                     let mut ex = EXECUTOR.lock();
                     let Some(w) = ex.timed_wait.peek() else {
                         break;
@@ -218,52 +218,53 @@ impl Executor {
                         break;
                     }
                     let w = ex.timed_wait.pop().unwrap();
-                    w.waker
+                    w.event
                 };
-                waker.wake();
+                event.trigger();
             }
             // Since we don't notice interrupts that happened while we are not hlt-ing,
             // make sure that we wake up all the interrupt-based waiting tasks every at
             // most INTERRUPT_MICROS micros to make it unlikely to miss interrupts.
             if last_interrupt_wakeup + INTERRUPT_MICROS <= Timer::micros() || force_interrupt_wake {
                 last_interrupt_wakeup = Timer::micros();
                 let to_wake = core::mem::take(&mut EXECUTOR.lock().wake_on_interrupt);
-                for w in to_wake {
-                    w.wake();
+                for e in to_wake {
+                    e.trigger();
                 }
             }
         };
 
         loop {
             do_wake(false);
             let task = EXECUTOR.lock().ready_tasks.pop_front();
-            if let Some(task) = task {
-                // It is possible for a done task to end up in the queue (if it wakes
-                // itself during execution). If that happens, we just remove it from
-                // the queue here.
-                if task.done.load(Ordering::Relaxed) {
-                    continue;
-                }
-                task.in_queue.store(false, Ordering::Relaxed);
-                let waker = Waker::from(task.clone());
-                let mut context = Context::from_waker(&waker);
-                let mut fut = task.future.try_lock().unwrap();
-                let begin = Timer::micros();
-                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);
-                }
-            } else {
+            let Some(task) = task else {
                 // If we don't have anything ready, sleep until the next interrupt.
                 // SAFETY: hlt is available on all reasonable x86 processors and has no safety
                 // requirements.
                 unsafe {
                     core::arch::asm!("hlt");
                 }
                 do_wake(true);
+                continue;
+            };
+
+            // It is possible for a done task to end up in the queue (if it wakes
+            // itself during execution). If that happens, we just remove it from
+            // the queue here.
+            if task.done.load(Ordering::Relaxed) {
+                continue;
+            }
+            task.in_queue.store(false, Ordering::Relaxed);
+            let waker = Waker::from(task.clone());
+            let mut context = Context::from_waker(&waker);
+            let mut fut = task.future.try_lock().unwrap();
+            let begin = Timer::micros();
+            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);
             }
         }
     }
@@ -283,41 +284,23 @@ impl Executor {
         })
     }
 
+    // Wakes a task as soon as *any* interrupt is received.
+    pub fn wait_for_interrupt() -> impl Future<Output = ()> {
+        let event = Event::new();
+        EXECUTOR.lock().wake_on_interrupt.push(event.trigger());
+        event
+    }
+
     // Note: there are no guarantees on whether the amount of time we will sleep for
     // will be exceeded.
     pub fn sleep(time: Duration) -> impl Future<Output = ()> {
         let tgt = Timer::micros() + time.as_micros() as i64;
-        let mut ww = None;
-        poll_fn(move |cx| {
-            let now = Timer::micros();
-            if now >= tgt {
-                Poll::Ready(())
-            } else {
-                Self::wake_at_micros(tgt, cx.waker(), &mut ww);
-                Poll::Pending
-            }
-        })
-    }
-
-    // Wakes a task as soon as *any* interrupt is received.
-    pub(super) fn wake_on_interrupt(waker: &Waker) {
-        EXECUTOR.lock().wake_on_interrupt.push(waker.clone());
-    }
-
-    pub(super) fn wake_at_micros(
-        micros: i64,
-        waker: &Waker,
-        previous_waker: &mut Option<WrappedWaker>,
-    ) {
-        if !previous_waker.is_some() {
-            let w = Arc::new(AtomicWaker::new());
-            EXECUTOR.lock().timed_wait.push(TimedWait {
-                wake_at: micros,
-                waker: w.clone(),
-            });
-            *previous_waker = Some(w);
-        }
-        previous_waker.as_ref().unwrap().register(waker);
+        let event = Event::new();
+        EXECUTOR.lock().timed_wait.push(TimedWait {
+            wake_at: tgt,
+            event: event.trigger(),
+        });
+        event
     }
 
     /// Spawn a new task.

pixie-uefi/src/os/input.rs

@@ -1,6 +1,3 @@
-use core::future::{poll_fn, Future};
-use core::task::Poll;
-
 use spin::lazy::Lazy;
 use spin::Mutex;
 use uefi::boot::ScopedProtocol;
@@ -16,16 +13,11 @@ static INPUT: Lazy<Mutex<SendWrapper<ScopedProtocol<Input>>>> = Lazy::new(|| {
     Mutex::new(SendWrapper(input))
 });
 
-pub fn read_key() -> impl Future<Output = Result<Key>> {
-    poll_fn(move |cx| {
-        let key = INPUT.lock().read_key();
-        match key {
-            Err(e) => Poll::Ready(Err(e.into())),
-            Ok(Some(key)) => Poll::Ready(Ok(key)),
-            Ok(None) => {
-                Executor::wake_on_interrupt(cx.waker());
-                Poll::Pending
-            }
+pub async fn read_key() -> Result<Key> {
+    loop {
+        if let Some(key) = INPUT.lock().read_key()? {
+            break Ok(key);
         }
-    })
+        Executor::wait_for_interrupt().await;
+    }
 }

pixie-uefi/src/os/net/mod.rs

@@ -1,14 +1,10 @@
 use alloc::string::{String, ToString};
-use alloc::sync::Arc;
 use alloc::vec::Vec;
 use core::fmt::Write;
-use core::future::{poll_fn, Future};
 use core::net::Ipv4Addr;
 use core::sync::atomic::{AtomicU64, Ordering};
-use core::task::Poll;
 use core::time::Duration;
 
-use futures::task::AtomicWaker;
 use smoltcp::iface::{
     Config, Interface, PollIngressSingleResult, PollResult, SocketHandle, SocketSet,
 };
@@ -25,7 +21,8 @@ use uefi::Handle;
 
 use super::timer::Timer;
 use crate::os::boot_options::BootOptions;
-use crate::os::executor::{Executor, WrappedWaker};
+use crate::os::executor::event::{Event as ExecutorEvent, EventTrigger};
+use crate::os::executor::Executor;
 use crate::os::net::interface::SnpDevice;
 pub use crate::os::net::tcp::TcpStream;
 pub use crate::os::net::udp::UdpSocket;
@@ -51,7 +48,7 @@ struct NetworkData {
 
 static NETWORK_DATA: Mutex<Option<NetworkData>> = Mutex::new(None);
 
-static WAITING_FOR_IP: Mutex<Vec<WrappedWaker>> = Mutex::new(vec![]);
+static WAITING_FOR_IP: Mutex<Vec<EventTrigger>> = Mutex::new(vec![]);
 
 fn with_net<T, F: FnOnce(&mut NetworkData) -> T>(f: F) -> T {
     let mut mg = NETWORK_DATA.try_lock().expect("Network is locked");
@@ -127,27 +124,28 @@ pub(super) fn init() {
         dhcp_socket_handle,
     });
 
-    Executor::spawn("[net_poll]", {
-        poll_fn(move |cx| {
+    Executor::spawn("[net_poll]", async {
+        loop {
             let wait = poll();
             match wait {
                 None => {
-                    Executor::wake_on_interrupt(cx.waker());
+                    Executor::wait_for_interrupt().await;
                 }
-                Some(x) if x < 200 => {
+                Some(wait) if wait < 200 => {
                     // Immediately wake if we want call poll() again in a very short time.
-                    cx.waker().wake_by_ref();
+                    Executor::sched_yield().await;
                 }
                 Some(wait) => {
-                    // Halve the waiting time, to try to ensure that we don't exceed the suggested
-                    // waiting time.
-                    let deadline = Timer::micros() + wait as i64 / 2;
-                    Executor::wake_on_interrupt(cx.waker());
-                    Executor::wake_at_micros(deadline, cx.waker(), &mut None);
+                    futures::future::select(
+                        Executor::wait_for_interrupt(),
+                        // Halve the waiting time, to try to ensure that we don't exceed the suggested
+                        // waiting time.
+                        Executor::sleep(Duration::from_micros(wait / 2)),
+                    )
+                    .await;
                 }
             }
-            Poll::<()>::Pending
-        })
+        }
     });
 
     Executor::spawn("[show_ip]", async {
@@ -169,21 +167,13 @@ pub(super) fn init() {
     speed::spawn_network_speed_task();
 }
 
-pub fn wait_for_ip() -> impl Future<Output = ()> {
-    let mut last_waker = None;
-    poll_fn(move |cx| {
-        if ip().is_some() {
-            Poll::Ready(())
-        } else {
-            if last_waker.is_none() {
-                let waker = Arc::new(AtomicWaker::new());
-                WAITING_FOR_IP.lock().push(waker.clone());
-                last_waker = Some(waker);
-            }
-            last_waker.as_ref().unwrap().register(cx.waker());
-            Poll::Pending
-        }
-    })
+pub async fn wait_for_ip() {
+    if ip().is_some() {
+        return;
+    }
+    let event = ExecutorEvent::new();
+    WAITING_FOR_IP.lock().push(event.trigger());
+    event.await;
 }
 
 fn ip() -> Option<Ipv4Addr> {
@@ -232,8 +222,8 @@ fn poll() -> Option<u64> {
                     .unwrap();
             }
             let to_wake = core::mem::take(&mut *WAITING_FOR_IP.lock());
-            for w in to_wake {
-                w.wake();
+            for e in to_wake {
+                e.trigger();
             }
         } else {
             interface.update_ip_addrs(|a| {