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/main.rs

@@ -132,11 +132,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_us(100_000).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,9 @@ 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 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 +53,31 @@ impl Wake for Task {
     }
 }
 
-static EXECUTOR: Mutex<Executor> = Mutex::new(Executor {
-    ready_tasks: VecDeque::new(),
-    tasks: vec![],
-});
+struct TimedWait {
+    wake_at: i64,
+    waker: Waker,
+}
+
+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 +87,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<Waker>,
+    timed_wait: BinaryHeap<TimedWait>,
     ready_tasks: VecDeque<Arc<Task>>,
     tasks: Vec<Arc<Task>>,
 }
@@ -167,24 +196,70 @@ impl Executor {
 
     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 +279,31 @@ impl Executor {
         })
     }
 
+    // Note: there are no guarantees on whether the amount of time we will sleep for
+    // will be exceeded.
     pub fn sleep_us(us: u64) -> impl Future<Output = ()> {
-        let tgt = Timer::micros() as u64 + us;
+        let tgt = Timer::micros() + us as i64;
         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());
                 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) {
+        EXECUTOR.lock().wake_on_interrupt.push(waker);
+    }
+
+    pub(super) fn wake_at_micros(micros: i64, waker: Waker) {
+        EXECUTOR.lock().timed_wait.push(TimedWait {
+            wake_at: micros,
+            waker,
+        });
     }
 
     /// 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;
 use crate::os::send_wrapper::SendWrapper;
 
 static INPUT: Lazy<Mutex<SendWrapper<ScopedProtocol<Input>>>> = Lazy::new(|| {
@@ -22,7 +23,7 @@ pub fn read_key() -> impl Future<Output = Result<Key>> {
             Err(e) => Poll::Ready(Err(e.into())),
             Ok(Some(key)) => Poll::Ready(Ok(key)),
             Ok(None) => {
-                cx.waker().wake_by_ref();
+                Executor::wake_on_interrupt(cx.waker().clone());
                 Poll::Pending
             }
         }

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.

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

@@ -1,11 +1,14 @@
 use alloc::string::{String, ToString};
+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::task::{Poll, Waker};
 
-use smoltcp::iface::{Config, Interface, PollResult, SocketHandle, SocketSet};
+use smoltcp::iface::{
+    Config, Interface, PollIngressSingleResult, PollResult, SocketHandle, SocketSet,
+};
 use smoltcp::socket::dhcpv4::{Event, Socket as Dhcpv4Socket};
 use smoltcp::wire::{DhcpOption, HardwareAddress, IpCidr};
 use spin::Mutex;
@@ -23,6 +26,7 @@ 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;
+use crate::os::send_wrapper::SendWrapper;
 use crate::os::timer::rdtsc;
 use crate::os::ui;
 
@@ -44,6 +48,8 @@ struct NetworkData {
 
 static NETWORK_DATA: Mutex<Option<NetworkData>> = Mutex::new(None);
 
+static WAITING_FOR_IP: Mutex<Vec<Waker>> = 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");
     f(mg.as_mut().expect("Network is not initialized"))
@@ -97,7 +103,7 @@ pub(super) fn init() {
         &snp.mode().current_address.0[..6],
     ));
 
-    let mut device = SnpDevice::new(snp);
+    let mut device = SnpDevice::new(SendWrapper(snp));
 
     let mut interface_config = Config::new(hw_addr);
     interface_config.random_seed = rdtsc() as u64;
@@ -121,9 +127,16 @@ pub(super) fn init() {
     Executor::spawn(
         "[net_poll]",
         poll_fn(move |cx| {
-            poll();
-            // TODO(veluca): figure out whether we can suspend the task.
-            cx.waker().wake_by_ref();
+            let wait = poll();
+            if wait == 0 {
+                cx.waker().wake_by_ref();
+            } else {
+                // 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().clone());
+                Executor::wake_at_micros(deadline, cx.waker().clone());
+            }
             Poll::<()>::Pending
         }),
     );
@@ -152,7 +165,7 @@ pub fn wait_for_ip() -> impl Future<Output = ()> {
         if ip().is_some() {
             Poll::Ready(())
         } else {
-            cx.waker().wake_by_ref();
+            WAITING_FOR_IP.lock().push(cx.waker().clone());
             Poll::Pending
         }
     })
@@ -167,25 +180,27 @@ fn get_ephemeral_port() -> u16 {
     ((ans % (60999 - 49152)) + 49152) as u16
 }
 
-fn poll() {
+// Returns # of microseconds to wait until we should call poll() again (possibly 0).
+fn poll() -> u64 {
     let now = Timer::instant();
 
     let mut data = NETWORK_DATA.lock();
 
-    let Some(NetworkData {
+    let NetworkData {
         interface,
         device,
         socket_set,
         dhcp_socket_handle,
-    }) = data.as_mut()
-    else {
-        return;
-    };
+    } = data.as_mut().unwrap();
 
-    let status = interface.poll(now, device, socket_set);
+    let status_out = interface.poll_egress(now, device, socket_set);
+    let status_in = interface.poll_ingress_single(now, device, socket_set);
 
-    if status == PollResult::None {
-        return;
+    if status_in == PollIngressSingleResult::None && status_out == PollResult::None {
+        return interface
+            .poll_delay(now, socket_set)
+            .map(|x| x.micros())
+            .unwrap_or(0);
     }
 
     let dhcp_status = socket_set
@@ -203,11 +218,16 @@ fn poll() {
                     .add_default_ipv4_route(router)
                     .unwrap();
             }
+            let to_wake = core::mem::take(&mut *WAITING_FOR_IP.lock());
+            for w in to_wake {
+                w.wake();
+            }
         } else {
             interface.update_ip_addrs(|a| {
                 a.clear();
             });
             interface.routes_mut().remove_default_ipv4_route();
         }
     }
+    0
 }

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

@@ -65,7 +65,12 @@ impl TcpStream {
             let state = with_net(|n| n.socket_set.get_mut::<TcpSocket>(self.handle).state());
             let res = f(state);
             if matches!(res, Poll::Pending) {
-                cx.waker().wake_by_ref();
+                with_net(|n| {
+                    let socket = n.socket_set.get_mut::<TcpSocket>(self.handle);
+                    // Every meaningful state change causes either send or recv to become possible (possibly erroring).
+                    socket.register_send_waker(cx.waker());
+                    socket.register_recv_waker(cx.waker());
+                });
             }
             res
         })