Split off UI from OS and redesign it.

This gets rid of the last parts of the OS struct.

Author: veluca93

pixie-uefi/src/flash.rs

@@ -12,13 +12,12 @@ use lz4_flex::decompress;
 use pixie_shared::chunk_codec::Decoder;
 use pixie_shared::util::BytesFmt;
 use pixie_shared::{ChunkHash, Image, TcpRequest, UdpRequest, CHUNKS_PORT, MAX_CHUNK_SIZE};
-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::{disk, memory, UefiOS};
+use crate::os::{disk, memory};
 use crate::MIN_MEMORY;
 
 async fn fetch_image(stream: &TcpStream) -> Result<Image> {
@@ -75,7 +74,7 @@ fn handle_packet(
     Ok(Some((pos, data)))
 }
 
-pub async fn flash(os: UefiOS, server_addr: SocketAddrV4) -> Result<()> {
+pub async fn flash(server_addr: SocketAddrV4) -> Result<()> {
     let stream = TcpStream::connect(server_addr).await?;
     let image = fetch_image(&stream).await?;
     stream.shutdown().await;
@@ -103,37 +102,14 @@ pub async fn flash(os: UefiOS, server_addr: SocketAddrV4) -> Result<()> {
     }));
 
     let stats2 = stats.clone();
-    os.set_ui_drawer(move |os| {
-        os.write_with_color(
-            &format!("{} total chunks\n", stats2.borrow().chunks),
-            Color::White,
-            Color::Black,
-        );
-        os.write_with_color(
-            &format!("{} unique chunks\n", stats2.borrow().unique),
-            Color::White,
-            Color::Black,
-        );
-        os.write_with_color(
-            &format!("{} chunks to fetch\n", stats2.borrow().fetch),
-            Color::White,
-            Color::Black,
-        );
-        os.write_with_color(
-            &format!("{} chunks received\n", stats2.borrow().recv),
-            Color::White,
-            Color::Black,
-        );
-        os.write_with_color(
-            &format!("{} packets received\n", stats2.borrow().pack_recv),
-            Color::White,
-            Color::Black,
-        );
-        os.write_with_color(
-            &format!("{} chunks requested\n", stats2.borrow().requested),
-            Color::White,
-            Color::Black,
-        );
+    crate::os::ui::set_content_drawer(move |draw_area| {
+        draw_area.clear();
+        draw_area.writeln(&format!("{} total chunks", stats2.borrow().chunks));
+        draw_area.writeln(&format!("{} unique chunks", stats2.borrow().unique));
+        draw_area.writeln(&format!("{} chunks to fetch", stats2.borrow().fetch));
+        draw_area.writeln(&format!("{} chunks received", stats2.borrow().recv));
+        draw_area.writeln(&format!("{} packets received", stats2.borrow().pack_recv));
+        draw_area.writeln(&format!("{} chunks requested", stats2.borrow().requested));
     });
 
     let mut disk = disk::Disk::largest();

pixie-uefi/src/main.rs

@@ -16,7 +16,6 @@ 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::register::register;
 use crate::store::store;
 
@@ -96,7 +95,7 @@ async fn complete_action(stream: &TcpStream) -> Result<()> {
     Ok(())
 }
 
-async fn run(os: UefiOS) -> Result<()> {
+async fn run() -> Result<()> {
     let server = server_discover().await?;
 
     let mut last_was_wait = false;
@@ -114,7 +113,7 @@ async fn run(os: UefiOS) -> Result<()> {
 
     loop {
         // Clear any UI drawer.
-        os.set_ui_drawer(|_| {});
+        os::ui::set_content_drawer(|_| {});
 
         if !last_was_wait {
             log::debug!("Sending request for command");
@@ -147,9 +146,9 @@ async fn run(os: UefiOS) -> Result<()> {
                     Action::Boot => power_control::reboot_to_os().await,
                     Action::Restart => {}
                     Action::Shutdown => shutdown().await,
-                    Action::Register => register(os, server).await?,
-                    Action::Store => store(os, server).await?,
-                    Action::Flash => flash(os, server).await?,
+                    Action::Register => register(server).await?,
+                    Action::Store => store(server).await?,
+                    Action::Flash => flash(server).await?,
                 }
 
                 let tcp = TcpStream::connect(server).await?;
@@ -167,5 +166,5 @@ async fn run(os: UefiOS) -> Result<()> {
 
 #[entry]
 fn main() -> Status {
-    UefiOS::start(run)
+    os::start(run)
 }

pixie-uefi/src/os/disk.rs

@@ -7,9 +7,8 @@ use uefi::boot::{OpenProtocolParams, ScopedProtocol};
 use uefi::proto::media::block::BlockIO;
 use uefi::Handle;
 
-use crate::os::executor::Executor;
-
 use super::error::Result;
+use crate::os::executor::Executor;
 
 fn open_disk(handle: Handle) -> Result<ScopedProtocol<BlockIO>> {
     let image_handle = uefi::boot::image_handle();

pixie-uefi/src/os/executor.rs

@@ -7,10 +7,13 @@ use core::future::{poll_fn, Future};
 use core::pin::Pin;
 use core::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 use core::task::{Context, Poll, Waker};
+
 use spin::Mutex;
 use uefi::boot::{EventType, TimerTrigger, Tpl};
+use uefi::proto::console::text::Color;
 
 use crate::os::timer::Timer;
+use crate::os::ui::DrawArea;
 
 struct BoxFuture(Pin<Box<dyn Future<Output = ()> + 'static>>);
 
@@ -22,6 +25,7 @@ struct Task {
     in_queue: AtomicBool,
     future: Mutex<BoxFuture>,
     micros: AtomicU64,
+    last_micros: AtomicU64,
 }
 
 impl Task {
@@ -33,6 +37,7 @@ impl Task {
             name,
             future: Mutex::new(BoxFuture(Box::pin(future))),
             micros: AtomicU64::new(0),
+            last_micros: AtomicU64::new(0),
             in_queue: AtomicBool::new(false),
         })
     }
@@ -58,7 +63,99 @@ pub struct Executor {
 }
 
 impl Executor {
+    async fn draw_tasks() {
+        let mut draw_area = DrawArea::tasks();
+        let (w, h) = draw_area.size();
+        const TASK_LEN: usize = 34;
+        let num_w = (w - 2) / TASK_LEN;
+        let total_w = num_w * TASK_LEN + 2;
+        let pad = (w - total_w) / 2;
+        let mut last = Timer::micros() as u64;
+        Self::sleep_us(100_000).await;
+        loop {
+            draw_area.clear();
+            let cur = Timer::micros() as u64;
+            let elapsed = cur.saturating_sub(last).max(1) as f64;
+            {
+                let mut executor = EXECUTOR.lock();
+                let tasks = &mut executor.tasks;
+                // Sort by *descending* time used since last draw.
+                tasks.sort_unstable_by_key(|f| {
+                    f.last_micros.load(Ordering::Relaxed) as i64
+                        - f.micros.load(Ordering::Relaxed) as i64
+                });
+
+                draw_area.advance(pad);
+                draw_area.write("\u{250C}");
+                for x in 0..num_w {
+                    for _ in 0..TASK_LEN {
+                        draw_area.write("\u{2500}");
+                    }
+                    draw_area.write(if x + 1 == num_w {
+                        "\u{2510}"
+                    } else {
+                        "\u{252C}"
+                    });
+                }
+                draw_area.newline();
+                for y in 0..(h - 2) {
+                    draw_area.advance(pad);
+                    draw_area.write("\u{2502}");
+                    for x in 0..num_w {
+                        let idx = x * h + y;
+                        if idx >= tasks.len() {
+                            draw_area.advance(TASK_LEN);
+                        } else {
+                            let task = &tasks[idx];
+                            let total_cpu = task.micros.load(Ordering::Relaxed);
+                            let last_cpu = task.last_micros.load(Ordering::Relaxed);
+                            let frac = ((total_cpu - last_cpu) as f64 / elapsed).min(1.0);
+                            draw_area.write_with_color(
+                                &format!(
+                                    " {:15}{:5.1}%{:10.3}s ",
+                                    &task.name[..task.name.len().min(15)],
+                                    frac * 100.0,
+                                    total_cpu as f64 * 0.000_001,
+                                ),
+                                if frac >= 0.5 {
+                                    Color::Red
+                                } else if frac >= 0.1 {
+                                    Color::Yellow
+                                } else {
+                                    Color::White
+                                },
+                                Color::Black,
+                            );
+                        }
+                        draw_area.write("\u{2502}");
+                    }
+                    draw_area.newline();
+                }
+                draw_area.advance(pad);
+                draw_area.write("\u{2514}");
+                for x in 0..num_w {
+                    for _ in 0..TASK_LEN {
+                        draw_area.write("\u{2500}");
+                    }
+                    draw_area.write(if x + 1 == num_w {
+                        "\u{2518}"
+                    } else {
+                        "\u{2534}"
+                    });
+                }
+
+                for t in tasks.iter() {
+                    t.last_micros
+                        .store(t.micros.load(Ordering::Relaxed), Ordering::Relaxed);
+                }
+            }
+            last = Timer::micros() as u64;
+            Self::sleep_us(1_000_000).await;
+        }
+    }
+
     pub fn run() -> ! {
+        Self::spawn("[show_tasks]", Self::draw_tasks());
         loop {
             let task = EXECUTOR
                 .lock()
@@ -126,16 +223,4 @@ impl Executor {
         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
-    }
 }

pixie-uefi/src/os/input.rs

@@ -0,0 +1,31 @@
+use core::future::{poll_fn, Future};
+use core::task::Poll;
+
+use spin::lazy::Lazy;
+use spin::Mutex;
+use uefi::boot::ScopedProtocol;
+use uefi::proto::console::text::{Input, Key};
+
+use crate::os::error::Result;
+use crate::os::send_wrapper::SendWrapper;
+
+static INPUT: Lazy<Mutex<SendWrapper<ScopedProtocol<Input>>>> = Lazy::new(|| {
+    let input_handles = uefi::boot::find_handles::<Input>().unwrap();
+    let input = uefi::boot::open_protocol_exclusive::<Input>(input_handles[0]).unwrap();
+    Mutex::new(SendWrapper(input))
+});
+
+pub fn read_key() -> impl Future<Output = Result<Key>> {
+    poll_fn(move |cx| {
+        let key = INPUT.lock().read_key();
+        if let Err(e) = key {
+            return Poll::Ready(Err(e.into()));
+        }
+        let key = key.unwrap();
+        if let Some(key) = key {
+            return Poll::Ready(Ok(key));
+        }
+        cx.waker().wake_by_ref();
+        Poll::Pending
+    })
+}