Split Disk and Executor from OS.

pixie-uefi/src/export_cov.rs

@@ -1,8 +1,7 @@
 use crate::os::disk::Disk;
-use crate::os::UefiOS;
 
-pub async fn export(os: UefiOS) {
-    let mut disk = Disk::open_with_size(os, 500 << 20);
+pub async fn export() {
+    let mut disk = Disk::open_with_size(500 << 20);
 
     let mut coverage = vec![];
     // SAFETY: we never create threads anyway.

pixie-uefi/src/flash.rs

@@ -16,8 +16,9 @@ use uefi::proto::console::text::Color;
 
 use crate::os::boot_options::BootOptions;
 use crate::os::error::{Error, Result};
+use crate::os::executor::Executor;
 use crate::os::net::{TcpStream, UdpSocket, ETH_PACKET_SIZE};
-use crate::os::{memory, UefiOS};
+use crate::os::{disk, memory, UefiOS};
 use crate::MIN_MEMORY;
 
 async fn fetch_image(stream: &TcpStream) -> Result<Image> {
@@ -135,7 +136,7 @@ pub async fn flash(os: UefiOS, server_addr: SocketAddrV4) -> Result<()> {
         );
     });
 
-    let mut disk = os.open_first_disk();
+    let mut disk = disk::Disk::largest();
 
     for (hash, (size, csize, pos)) in mem::take(&mut chunks_info) {
         let mut found = None;
@@ -179,7 +180,7 @@ pub async fn flash(os: UefiOS, server_addr: SocketAddrV4) -> Result<()> {
         );
         while !chunks_info.is_empty() {
             let recv = Box::pin(socket.recv_from(&mut buf));
-            let sleep = Box::pin(os.sleep_us(100_000));
+            let sleep = Box::pin(Executor::sleep_us(100_000));
             match select(recv, sleep).await {
                 Either::Left(((buf, _addr), _)) => {
                     stats.borrow_mut().pack_recv += 1;

pixie-uefi/src/main.rs

@@ -14,16 +14,16 @@ use uefi::{entry, Status};
 
 use crate::flash::flash;
 use crate::os::error::{Error, Result};
+use crate::os::executor::Executor;
 use crate::os::net::{TcpStream, UdpSocket, ETH_PACKET_SIZE};
 use crate::os::UefiOS;
-use crate::reboot_to_os::reboot_to_os;
 use crate::register::register;
 use crate::store::store;
 
 mod flash;
 mod os;
 mod parse_disk;
-mod reboot_to_os;
+mod power_control;
 mod register;
 mod store;
 
@@ -33,7 +33,7 @@ mod export_cov;
 // Memory to keep free for non-chunk storage.
 const MIN_MEMORY: u64 = 32 << 20;
 
-async fn server_discover(os: UefiOS) -> Result<SocketAddrV4> {
+async fn server_discover() -> Result<SocketAddrV4> {
     let socket = UdpSocket::bind(None).await?;
 
     let task1 = async {
@@ -43,7 +43,7 @@ async fn server_discover(os: UefiOS) -> Result<SocketAddrV4> {
             socket
                 .send_to(SocketAddrV4::new(Ipv4Addr::BROADCAST, ACTION_PORT), &msg)
                 .await?;
-            os.sleep_us(1_000_000).await;
+            Executor::sleep_us(1_000_000).await;
         })
     };
 
@@ -65,13 +65,13 @@ async fn server_discover(os: UefiOS) -> Result<SocketAddrV4> {
     Ok(server)
 }
 
-async fn shutdown(os: UefiOS) -> ! {
+async fn shutdown() -> ! {
     #[cfg(feature = "coverage")]
-    export_cov::export(os).await;
+    export_cov::export().await;
 
     log::info!("Shutting down...");
-    os.sleep_us(1_000_000).await;
-    os.shutdown()
+    Executor::sleep_us(1_000_000).await;
+    power_control::shutdown()
 }
 
 async fn get_action(stream: &TcpStream) -> Result<Action> {
@@ -97,18 +97,18 @@ async fn complete_action(stream: &TcpStream) -> Result<()> {
 }
 
 async fn run(os: UefiOS) -> Result<()> {
-    let server = server_discover(os).await?;
+    let server = server_discover().await?;
 
     let mut last_was_wait = false;
 
-    os.spawn("ping", async move {
+    Executor::spawn("ping", async move {
         let udp_socket = UdpSocket::bind(None).await.unwrap();
         loop {
             udp_socket
                 .send_to(SocketAddrV4::new(*server.ip(), PING_PORT), b"pixie")
                 .await
                 .unwrap();
-            os.sleep_us(10_000_000).await;
+            Executor::sleep_us(10_000_000).await;
         }
     });
 
@@ -136,17 +136,17 @@ async fn run(os: UefiOS) -> Result<()> {
                 last_was_wait = true;
                 const WAIT_10MSECS: u64 = 50;
                 for _ in 0..WAIT_10MSECS {
-                    os.deep_sleep_us(10_000);
-                    os.schedule().await;
+                    Executor::deep_sleep_us(10_000);
+                    Executor::sched_yield().await;
                 }
             } else {
                 last_was_wait = false;
                 log::info!("Command: {command:?}");
                 match command {
                     Action::Wait => unreachable!(),
-                    Action::Boot => reboot_to_os(os).await,
+                    Action::Boot => power_control::reboot_to_os().await,
                     Action::Restart => {}
-                    Action::Shutdown => shutdown(os).await,
+                    Action::Shutdown => shutdown().await,
                     Action::Register => register(os, server).await?,
                     Action::Store => store(os, server).await?,
                     Action::Flash => flash(os, server).await?,
@@ -158,7 +158,7 @@ async fn run(os: UefiOS) -> Result<()> {
                 tcp.force_close().await;
 
                 if command == Action::Restart {
-                    os.reset();
+                    power_control::reset()
                 }
             }
         }

pixie-uefi/src/os/disk.rs

@@ -7,8 +7,9 @@ use uefi::boot::{OpenProtocolParams, ScopedProtocol};
 use uefi::proto::media::block::BlockIO;
 use uefi::Handle;
 
+use crate::os::executor::Executor;
+
 use super::error::Result;
-use super::UefiOS;
 
 fn open_disk(handle: Handle) -> Result<ScopedProtocol<BlockIO>> {
     let image_handle = uefi::boot::image_handle();
@@ -35,13 +36,12 @@ pub struct DiskPartition {
 
 pub struct Disk {
     block: ScopedProtocol<BlockIO>,
-    os: UefiOS,
 }
 
 // TODO(veluca): consider making parts of this actually async, i.e. by using DiskIo2/BlockIO2 if
 // available; support having more than one disk.
 impl Disk {
-    pub fn new(os: UefiOS) -> Disk {
+    pub fn largest() -> Disk {
         let (_size, handle) = uefi::boot::find_handles::<BlockIO>()
             .unwrap()
             .into_iter()
@@ -60,11 +60,11 @@ impl Disk {
             .expect("Disk not found");
 
         let block = open_disk(handle).unwrap();
-        Disk { block, os }
+        Disk { block }
     }
 
     #[cfg(feature = "coverage")]
-    pub fn open_with_size(os: UefiOS, base_size: i64) -> Disk {
+    pub fn open_with_size(base_size: i64) -> Disk {
         let (_size, handle) = uefi::boot::find_handles::<BlockIO>()
             .unwrap()
             .into_iter()
@@ -83,7 +83,7 @@ impl Disk {
             .expect("Disk not found");
 
         let block = open_disk(handle).unwrap();
-        Disk { block, os }
+        Disk { block }
     }
 
     pub fn size(&self) -> u64 {
@@ -124,7 +124,7 @@ impl Disk {
     }
 
     pub async fn read(&self, offset: u64, buf: &mut [u8]) -> Result<()> {
-        self.os.schedule().await;
+        Executor::sched_yield().await;
         self.read_sync(offset, buf)
     }
 
@@ -158,7 +158,7 @@ impl Disk {
     }
 
     pub async fn write(&mut self, offset: u64, buf: &[u8]) -> Result<()> {
-        self.os.schedule().await;
+        Executor::sched_yield().await;
         self.write_sync(offset, buf)
     }
 

pixie-uefi/src/os/executor.rs

@@ -2,25 +2,26 @@ use alloc::boxed::Box;
 use alloc::collections::VecDeque;
 use alloc::sync::Arc;
 use alloc::task::Wake;
-use core::cell::RefCell;
-use core::future::Future;
+use alloc::vec::Vec;
+use core::future::{poll_fn, Future};
 use core::pin::Pin;
-use core::task::{Context, Waker};
+use core::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+use core::task::{Context, Poll, Waker};
+use spin::Mutex;
+use uefi::boot::{EventType, TimerTrigger, Tpl};
 
-use super::sync::SyncRefCell;
 use crate::os::timer::Timer;
 
-pub(super) type BoxFuture<T = ()> = Pin<Box<dyn Future<Output = T> + 'static>>;
+struct BoxFuture(Pin<Box<dyn Future<Output = ()> + 'static>>);
 
-struct TaskInner {
-    pub in_queue: bool,
-    pub future: Option<BoxFuture>,
-    pub micros: i64,
-}
+// SAFETY: there are no threads.
+unsafe impl Send for BoxFuture {}
 
-pub(super) struct Task {
-    pub name: &'static str,
-    inner: RefCell<TaskInner>,
+struct Task {
+    name: &'static str,
+    in_queue: AtomicBool,
+    future: Mutex<BoxFuture>,
+    micros: AtomicU64,
 }
 
 impl Task {
@@ -30,65 +31,111 @@ impl Task {
     {
         Arc::new(Task {
             name,
-            inner: RefCell::new(TaskInner {
-                in_queue: false,
-                future: Some(Box::pin(future)),
-                micros: 0,
-            }),
+            future: Mutex::new(BoxFuture(Box::pin(future))),
+            micros: AtomicU64::new(0),
+            in_queue: AtomicBool::new(false),
         })
     }
-
-    pub(super) fn micros(&self) -> i64 {
-        self.inner.borrow().micros
-    }
 }
 
-// SAFETY: we never create threads anyway.
-unsafe impl Send for Task {}
-unsafe impl Sync for Task {}
-
 impl Wake for Task {
     fn wake(self: Arc<Self>) {
-        if !self.inner.borrow().in_queue {
-            self.inner.borrow_mut().in_queue = true;
-            EXECUTOR.borrow_mut().ready_tasks.push_back(self);
+        if !self.in_queue.swap(true, Ordering::Relaxed) {
+            EXECUTOR.lock().ready_tasks.push_back(self);
         }
     }
 }
 
-static EXECUTOR: SyncRefCell<Executor> = SyncRefCell::new(Executor {
+static EXECUTOR: Mutex<Executor> = Mutex::new(Executor {
     ready_tasks: VecDeque::new(),
+    tasks: vec![],
 });
 
 pub struct Executor {
     // TODO(veluca): scheduling.
     ready_tasks: VecDeque<Arc<Task>>,
+    tasks: Vec<Arc<Task>>,
 }
 
 impl Executor {
     pub fn run() -> ! {
         loop {
             let task = EXECUTOR
-                .borrow_mut()
+                .lock()
                 .ready_tasks
                 .pop_front()
                 .expect("Executor should never run out of ready tasks");
 
-            task.inner.borrow_mut().in_queue = false;
+            task.in_queue.store(false, Ordering::Relaxed);
             let waker = Waker::from(task.clone());
             let mut context = Context::from_waker(&waker);
-            let mut fut = task.inner.borrow_mut().future.take().unwrap();
+            let mut fut = task.future.try_lock().unwrap();
             let begin = Timer::micros();
-            let status = fut.as_mut().poll(&mut context);
+            let _ = fut.0.as_mut().poll(&mut context);
             let end = Timer::micros();
-            task.inner.borrow_mut().micros += end - begin;
-            if status.is_pending() {
-                task.inner.borrow_mut().future = Some(fut);
-            }
+            task.micros
+                .fetch_add((end - begin) as u64, Ordering::Relaxed);
         }
     }
 
-    pub(super) fn spawn(task: Arc<Task>) {
-        EXECUTOR.borrow_mut().ready_tasks.push_back(task);
+    /// Interrupt task execution.
+    /// This is useful to yield the CPU to other tasks.
+    pub fn sched_yield() -> impl Future<Output = ()> {
+        let mut ready = false;
+        poll_fn(move |cx| {
+            if ready {
+                Poll::Ready(())
+            } else {
+                ready = true;
+                cx.waker().wake_by_ref();
+                Poll::Pending
+            }
+        })
+    }
+
+    pub fn sleep_us(us: u64) -> impl Future<Output = ()> {
+        let tgt = Timer::micros() as u64 + us;
+        poll_fn(move |cx| {
+            let now = Timer::micros() as u64;
+            if now >= tgt {
+                Poll::Ready(())
+            } else {
+                // TODO(veluca): actually suspend the task.
+                cx.waker().wake_by_ref();
+                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();
+    }
+
+    /// Spawn a new task.
+    pub fn spawn<Fut>(name: &'static str, f: Fut)
+    where
+        Fut: Future<Output = ()> + 'static,
+    {
+        let task = Task::new(name, f);
+        let mut executor = EXECUTOR.lock();
+        executor.tasks.push(task.clone());
+        executor.ready_tasks.push_back(task);
+    }
+
+    pub fn top_tasks(n: usize) -> Vec<(u64, &'static str)> {
+        let mut tasks: Vec<_> = EXECUTOR
+            .lock()
+            .tasks
+            .iter()
+            .map(|x| (x.micros.load(Ordering::Relaxed), x.name))
+            .collect();
+        tasks.sort_by_key(|t| -(t.0 as i64));
+        tasks.truncate(n);
+        tasks
     }
 }