Implement an actual scheduler.

As disk operations are still blocking, it doesn't quite work as well as
one would want, but it's a start.

Author: veluca93

pixie-uefi/src/flash.rs

@@ -5,6 +5,7 @@ use core::cell::{Cell, RefCell};
 use core::fmt::Write;
 use core::mem;
 use core::net::SocketAddrV4;
+use core::time::Duration;
 
 use futures::future::{select, Either};
 use log::info;
@@ -120,7 +121,7 @@ pub async fn flash(server_addr: SocketAddrV4) -> Result<()> {
     let draw_task = async {
         let mut last_stats = stats.borrow().clone();
         while !done.get() {
-            Executor::sleep_us(100_000).await;
+            Executor::sleep(Duration::from_millis(100)).await;
             let stats = stats.borrow().clone();
             if stats == last_stats {
                 continue;
@@ -176,7 +177,7 @@ pub async fn flash(server_addr: SocketAddrV4) -> Result<()> {
         );
         while !chunks_info.is_empty() {
             let recv = Box::pin(socket.recv_from(&mut buf));
-            let sleep = Box::pin(Executor::sleep_us(100_000));
+            let sleep = Box::pin(Executor::sleep(Duration::from_millis(100)));
             match select(recv, sleep).await {
                 Either::Left(((buf, _addr), _)) => {
                     stats.borrow_mut().pack_recv += 1;

pixie-uefi/src/main.rs

@@ -7,6 +7,7 @@ extern crate alloc;
 
 use alloc::boxed::Box;
 use core::net::{Ipv4Addr, SocketAddrV4};
+use core::time::Duration;
 
 use futures::future::{self, Either};
 use pixie_shared::{Action, TcpRequest, UdpRequest, ACTION_PORT, PING_PORT};
@@ -43,7 +44,7 @@ async fn server_discover() -> Result<SocketAddrV4> {
             socket
                 .send_to(SocketAddrV4::new(Ipv4Addr::BROADCAST, ACTION_PORT), &msg)
                 .await?;
-            Executor::sleep_us(1_000_000).await;
+            Executor::sleep(Duration::from_secs(1)).await;
         })
     };
 
@@ -70,7 +71,7 @@ async fn shutdown() -> ! {
     export_cov::export().await;
 
     log::info!("Shutting down...");
-    Executor::sleep_us(1_000_000).await;
+    Executor::sleep(Duration::from_secs(1)).await;
     power_control::shutdown()
 }
 
@@ -108,7 +109,7 @@ async fn run() -> Result<()> {
                 .send_to(SocketAddrV4::new(*server.ip(), PING_PORT), b"pixie")
                 .await
                 .unwrap();
-            Executor::sleep_us(10_000_000).await;
+            Executor::sleep(Duration::from_secs(10)).await;
         }
     });
 
@@ -132,11 +133,7 @@ async fn run() -> Result<()> {
                     log::warn!("Started waiting for another command...");
                 }
                 last_was_wait = true;
-                const WAIT_10MSECS: u64 = 50;
-                for _ in 0..WAIT_10MSECS {
-                    Executor::deep_sleep_us(10_000);
-                    Executor::sched_yield().await;
-                }
+                Executor::sleep(Duration::from_millis(500)).await;
             } else {
                 last_was_wait = false;
                 log::info!("Command: {command:?}");

pixie-uefi/src/os/executor.rs

@@ -1,4 +1,5 @@
 use alloc::boxed::Box;
+use alloc::collections::binary_heap::BinaryHeap;
 use alloc::collections::VecDeque;
 use alloc::sync::Arc;
 use alloc::task::Wake;
@@ -8,10 +9,10 @@ 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 core::time::Duration;
 
+use futures::channel::oneshot;
 use spin::Mutex;
-use uefi::boot::{EventType, TimerTrigger, Tpl};
 use uefi::proto::console::text::Color;
 
 use crate::os::send_wrapper::SendWrapper;
@@ -53,10 +54,58 @@ impl Wake for Task {
     }
 }
 
-static EXECUTOR: Mutex<Executor> = Mutex::new(Executor {
-    ready_tasks: VecDeque::new(),
-    tasks: vec![],
-});
+#[derive(Clone)]
+pub(super) struct WrappedWaker(Arc<Mutex<Option<Waker>>>);
+
+impl WrappedWaker {
+    pub(super) fn empty() -> Self {
+        Self(Arc::new(Mutex::new(None)))
+    }
+
+    pub(super) fn clear(&self) {
+        *self.0.lock() = None;
+    }
+
+    fn wake(self) {
+        let mut w = self.0.lock();
+        if let Some(w) = w.take() {
+            w.wake();
+        }
+    }
+
+    fn replace(&self, waker: Waker) -> bool {
+        let mut w = self.0.lock();
+        let was_present = w.is_some();
+        *w = Some(waker);
+        was_present
+    }
+}
+
+struct TimedWait {
+    wake_at: i64,
+    waker: WrappedWaker,
+}
+
+impl PartialEq for TimedWait {
+    fn eq(&self, other: &Self) -> bool {
+        self.wake_at == other.wake_at
+    }
+}
+
+impl Eq for TimedWait {}
+
+impl PartialOrd for TimedWait {
+    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for TimedWait {
+    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
+        // Reversed order: min-heap.
+        other.wake_at.cmp(&self.wake_at)
+    }
+}
 
 pub struct JoinHandle<T>(oneshot::Receiver<T>);
 
@@ -66,8 +115,16 @@ impl<T> JoinHandle<T> {
     }
 }
 
+static EXECUTOR: Mutex<Executor> = Mutex::new(Executor {
+    wake_on_interrupt: vec![],
+    timed_wait: BinaryHeap::new(),
+    ready_tasks: VecDeque::new(),
+    tasks: vec![],
+});
+
 pub struct Executor {
-    // TODO(veluca): scheduling.
+    wake_on_interrupt: Vec<WrappedWaker>,
+    timed_wait: BinaryHeap<TimedWait>,
     ready_tasks: VecDeque<Arc<Task>>,
     tasks: Vec<Arc<Task>>,
 }
@@ -81,7 +138,7 @@ impl Executor {
         assert!((w - 1).is_multiple_of(TASK_LEN + 1));
         let num_w = (w - 1) / (TASK_LEN + 1);
         let mut last = Timer::micros() as u64;
-        Self::sleep_us(100_000).await;
+        Self::sleep(Duration::from_millis(100)).await;
         loop {
             draw_area.clear();
             let cur = Timer::micros() as u64;
@@ -161,30 +218,76 @@ impl Executor {
                 tasks.retain(|t| !t.done.load(Ordering::Relaxed));
             }
             last = Timer::micros() as u64;
-            Self::sleep_us(1_000_000).await;
+            Self::sleep(Duration::from_secs(1)).await
         }
     }
 
     pub fn run() -> ! {
         Self::spawn("[show_tasks]", Self::draw_tasks());
+
+        // Maximum amount of microseconds between wakeups of interrupt-based wakers.
+        const INTERRUPT_MICROS: i64 = 500;
+
+        let mut last_interrupt_wakeup = Timer::micros();
+
+        let mut do_wake = |force_interrupt_wake| {
+            // Wake timed-waiting tasks.
+            loop {
+                let waker = {
+                    let mut ex = EXECUTOR.lock();
+                    let Some(w) = ex.timed_wait.peek() else {
+                        break;
+                    };
+                    if w.wake_at > Timer::micros() {
+                        break;
+                    }
+                    let w = ex.timed_wait.pop().unwrap();
+                    w.waker
+                };
+                waker.wake();
+            }
+            // 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();
+                }
+            }
+        };
+
         loop {
-            let task = EXECUTOR
-                .lock()
-                .ready_tasks
-                .pop_front()
-                .expect("Executor should never run out of ready tasks");
-
-            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);
+            let task = EXECUTOR.lock().ready_tasks.pop_front();
+            if let Some(task) = task {
+                {
+                    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);
+                    }
+                }
+                do_wake(false);
+            } else {
+                do_wake(false);
+                if !EXECUTOR.lock().ready_tasks.is_empty() {
+                    continue;
+                }
+                // If we still 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);
             }
         }
     }
@@ -204,27 +307,39 @@ impl Executor {
         })
     }
 
-    pub fn sleep_us(us: u64) -> impl Future<Output = ()> {
-        let tgt = Timer::micros() as u64 + us;
+    // 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 ww = WrappedWaker::empty();
         poll_fn(move |cx| {
-            let now = Timer::micros() as u64;
+            let now = Timer::micros();
             if now >= tgt {
                 Poll::Ready(())
             } else {
-                // TODO(veluca): actually suspend the task.
-                cx.waker().wake_by_ref();
+                Self::wake_at_micros(tgt, cx.waker().clone(), &ww);
                 Poll::Pending
             }
         })
     }
 
-    /// **WARNING**: this function halts all tasks
-    pub fn deep_sleep_us(us: u64) {
-        // SAFETY: we are not using a callback
-        let e =
-            unsafe { uefi::boot::create_event(EventType::TIMER, Tpl::NOTIFY, None, None).unwrap() };
-        uefi::boot::set_timer(&e, TimerTrigger::Relative(10 * us)).unwrap();
-        uefi::boot::wait_for_event(&mut [e]).unwrap();
+    // Wakes a task as soon as *any* interrupt is received.
+    pub(super) fn wake_on_interrupt(waker: Waker, previous_waker: &WrappedWaker) {
+        if !previous_waker.replace(waker) {
+            EXECUTOR
+                .lock()
+                .wake_on_interrupt
+                .push(previous_waker.clone());
+        }
+    }
+
+    pub(super) fn wake_at_micros(micros: i64, waker: Waker, previous_waker: &WrappedWaker) {
+        if !previous_waker.replace(waker) {
+            EXECUTOR.lock().timed_wait.push(TimedWait {
+                wake_at: micros,
+                waker: previous_waker.clone(),
+            });
+        }
     }
 
     /// Spawn a new task.

pixie-uefi/src/os/input.rs

@@ -7,6 +7,7 @@ use uefi::boot::ScopedProtocol;
 use uefi::proto::console::text::{Input, Key};
 
 use crate::os::error::Result;
+use crate::os::executor::{Executor, WrappedWaker};
 use crate::os::send_wrapper::SendWrapper;
 
 static INPUT: Lazy<Mutex<SendWrapper<ScopedProtocol<Input>>>> = Lazy::new(|| {
@@ -16,13 +17,20 @@ static INPUT: Lazy<Mutex<SendWrapper<ScopedProtocol<Input>>>> = Lazy::new(|| {
 });
 
 pub fn read_key() -> impl Future<Output = Result<Key>> {
+    let ww = WrappedWaker::empty();
     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)),
+            Err(e) => {
+                ww.clear();
+                Poll::Ready(Err(e.into()))
+            }
+            Ok(Some(key)) => {
+                ww.clear();
+                Poll::Ready(Ok(key))
+            }
             Ok(None) => {
-                cx.waker().wake_by_ref();
+                Executor::wake_on_interrupt(cx.waker().clone(), &ww);
                 Poll::Pending
             }
         }

pixie-uefi/src/os/mod.rs

@@ -2,6 +2,7 @@ use core::ffi::c_void;
 use core::future::Future;
 use core::ptr::NonNull;
 use core::sync::atomic::{AtomicBool, Ordering};
+use core::time::Duration;
 
 use uefi::boot::{EventType, Tpl};
 use uefi::{Event, Status};
@@ -72,7 +73,7 @@ where
                 break;
             }
 
-            Executor::sleep_us(30_000_000).await;
+            Executor::sleep(Duration::from_secs(30)).await;
         }
     });
 

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

@@ -5,8 +5,9 @@ use uefi::proto::network::snp::{ReceiveFlags, SimpleNetwork};
 use uefi::Status;
 
 use super::ETH_PACKET_SIZE;
+use crate::os::send_wrapper::SendWrapper;
 
-type Snp = ScopedProtocol<SimpleNetwork>;
+type Snp = SendWrapper<ScopedProtocol<SimpleNetwork>>;
 
 pub struct SnpDevice {
     snp: Snp,
@@ -15,9 +16,6 @@ pub struct SnpDevice {
     rx_buf: [u8; ETH_PACKET_SIZE + 4],
 }
 
-// SAFETY: we never create threads anyway.
-unsafe impl Send for SnpDevice {}
-
 impl SnpDevice {
     pub fn new(snp: Snp) -> SnpDevice {
         // Shut down the SNP protocol if needed.