diff --git a/pixie-shared/src/util.rs b/pixie-shared/src/util.rs index 95c56f7a..32a4e595 100644 --- a/pixie-shared/src/util.rs +++ b/pixie-shared/src/util.rs @@ -2,16 +2,42 @@ pub struct BytesFmt(pub u64); impl core::fmt::Display for BytesFmt { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + let w = f.width().unwrap_or(0); + let prc = f.precision().unwrap_or(2); if self.0 < 1 << 10 { - write!(f, "{} B", self.0) + write!(f, "{:1$} B", self.0, w.saturating_sub(2)) } else if self.0 < 1 << 20 { - write!(f, "{:.2} KiB", self.0 as f64 / (1i64 << 10) as f64) + write!( + f, + "{:1$.2$} KiB", + self.0 as f64 / (1i64 << 10) as f64, + w.saturating_sub(4), + prc + ) } else if self.0 < 1 << 30 { - write!(f, "{:.2} MiB", self.0 as f64 / (1i64 << 20) as f64) + write!( + f, + "{:1$.2$} MiB", + self.0 as f64 / (1i64 << 20) as f64, + w.saturating_sub(4), + prc + ) } else if self.0 < 1 << 40 { - write!(f, "{:.2} GiB", self.0 as f64 / (1i64 << 30) as f64) + write!( + f, + "{:1$.2$} GiB", + self.0 as f64 / (1i64 << 30) as f64, + w.saturating_sub(4), + prc + ) } else { - write!(f, "{:.2} TiB", self.0 as f64 / (1i64 << 40) as f64) + write!( + f, + "{:1$.2$} TiB", + self.0 as f64 / (1i64 << 40) as f64, + w.saturating_sub(4), + prc + ) } } } diff --git a/pixie-uefi/src/flash.rs b/pixie-uefi/src/flash.rs index fae377d4..6167ea90 100644 --- a/pixie-uefi/src/flash.rs +++ b/pixie-uefi/src/flash.rs @@ -1,8 +1,8 @@ use alloc::boxed::Box; use alloc::collections::BTreeMap; -use alloc::rc::Rc; use alloc::vec::Vec; -use core::cell::RefCell; +use core::cell::{Cell, RefCell}; +use core::fmt::Write; use core::mem; use core::net::SocketAddrV4; @@ -12,13 +12,13 @@ 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::ui::{update_content, DrawArea}; +use crate::os::{disk, memory}; use crate::MIN_MEMORY; async fn fetch_image(stream: &TcpStream) -> Result { @@ -32,6 +32,7 @@ async fn fetch_image(stream: &TcpStream) -> Result { Ok(postcard::from_bytes(&buf)?) } +#[derive(Clone, PartialEq, Eq)] struct Stats { chunks: usize, unique: usize, @@ -75,7 +76,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; @@ -93,49 +94,43 @@ pub async fn flash(os: UefiOS, server_addr: SocketAddrV4) -> Result<()> { info!("Obtained chunks; {} distinct chunks", chunks_info.len()); - let stats = Rc::new(RefCell::new(Stats { + let stats = RefCell::new(Stats { chunks: image.disk.len(), unique: chunks_info.len(), fetch: 0, recv: 0, pack_recv: 0, requested: 0, - })); - - 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, - ); }); + let draw = |draw_area: &mut DrawArea| { + draw_area.clear(); + writeln!(draw_area, "{} total chunks", stats.borrow().chunks).unwrap(); + writeln!(draw_area, "{} unique chunks", stats.borrow().unique).unwrap(); + writeln!(draw_area, "{} chunks to fetch", stats.borrow().fetch).unwrap(); + writeln!(draw_area, "{} chunks received", stats.borrow().recv).unwrap(); + writeln!(draw_area, "{} packets received", stats.borrow().pack_recv).unwrap(); + writeln!(draw_area, "{} chunks requested", stats.borrow().requested).unwrap(); + }; + + update_content(draw); + + let done = Cell::new(false); + + let draw_task = async { + let mut last_stats = stats.borrow().clone(); + while !done.get() { + Executor::sleep_us(100_000).await; + let stats = stats.borrow().clone(); + if stats == last_stats { + continue; + } + update_content(draw); + last_stats = stats; + } + Ok(()) + }; + let mut disk = disk::Disk::largest(); for (hash, (size, csize, pos)) in mem::take(&mut chunks_info) { @@ -158,6 +153,7 @@ pub async fn flash(os: UefiOS, server_addr: SocketAddrV4) -> Result<()> { chunks_info.insert(hash, (size, csize, pos)); stats.borrow_mut().fetch = chunks_info.len(); } + update_content(draw); } info!("Disk scanned; {} chunks to fetch", stats.borrow().fetch); @@ -226,10 +222,11 @@ pub async fn flash(os: UefiOS, server_addr: SocketAddrV4) -> Result<()> { .send_to(server_addr, &postcard::to_allocvec(&msg)?) .await?; } + done.set(true); Ok(()) }; - let ((), ()) = futures::try_join!(task1, task2)?; + let ((), (), ()) = futures::try_join!(task1, task2, draw_task)?; info!("Fetch complete, updating boot options"); diff --git a/pixie-uefi/src/main.rs b/pixie-uefi/src/main.rs index fe0e66a3..6019ba27 100644 --- a/pixie-uefi/src/main.rs +++ b/pixie-uefi/src/main.rs @@ -16,7 +16,7 @@ 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::os::ui::update_content; use crate::register::register; use crate::store::store; @@ -96,7 +96,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; @@ -113,9 +113,7 @@ async fn run(os: UefiOS) -> Result<()> { }); loop { - // Clear any UI drawer. - os.set_ui_drawer(|_| {}); - + update_content(|d| d.clear()); if !last_was_wait { log::debug!("Sending request for command"); } @@ -147,9 +145,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 +165,5 @@ async fn run(os: UefiOS) -> Result<()> { #[entry] fn main() -> Status { - UefiOS::start(run) + os::start(run) } diff --git a/pixie-uefi/src/os/disk.rs b/pixie-uefi/src/os/disk.rs index 5aad1c80..5a2fe531 100644 --- a/pixie-uefi/src/os/disk.rs +++ b/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> { let image_handle = uefi::boot::image_handle(); diff --git a/pixie-uefi/src/os/executor.rs b/pixie-uefi/src/os/executor.rs index e5fc0ef8..1a3689b4 100644 --- a/pixie-uefi/src/os/executor.rs +++ b/pixie-uefi/src/os/executor.rs @@ -3,25 +3,28 @@ use alloc::collections::VecDeque; use alloc::sync::Arc; use alloc::task::Wake; use alloc::vec::Vec; +use core::fmt::Write; 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::send_wrapper::SendWrapper; use crate::os::timer::Timer; +use crate::os::ui::DrawArea; -struct BoxFuture(Pin + 'static>>); - -// SAFETY: there are no threads. -unsafe impl Send for BoxFuture {} +type BoxFuture = SendWrapper + 'static>>>; struct Task { name: &'static str, in_queue: AtomicBool, future: Mutex, micros: AtomicU64, + last_micros: AtomicU64, } impl Task { @@ -31,8 +34,9 @@ impl Task { { Arc::new(Task { name, - future: Mutex::new(BoxFuture(Box::pin(future))), + future: Mutex::new(SendWrapper(Box::pin(future))), micros: AtomicU64::new(0), + last_micros: AtomicU64::new(0), in_queue: AtomicBool::new(false), }) } @@ -57,8 +61,98 @@ pub struct Executor { tasks: Vec>, } +pub(super) const TASK_LEN: usize = 34; + impl Executor { + async fn draw_tasks() { + let mut draw_area = DrawArea::tasks(); + let (w, h) = draw_area.size(); + 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; + 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 + }); + + write!(draw_area, "\u{250C}").unwrap(); + for x in 0..num_w { + for _ in 0..TASK_LEN { + write!(draw_area, "\u{2500}").unwrap(); + } + if x + 1 == num_w { + write!(draw_area, "\u{2510}").unwrap(); + } else { + write!(draw_area, "\u{252C}").unwrap(); + } + } + draw_area.newline(); + for y in 0..(h - 2) { + write!(draw_area, "\u{2502}").unwrap(); + 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, + ); + } + write!(draw_area, "\u{2502}").unwrap(); + } + draw_area.newline(); + } + write!(draw_area, "\u{2514}").unwrap(); + for x in 0..num_w { + for _ in 0..TASK_LEN { + write!(draw_area, "\u{2500}").unwrap(); + } + if x + 1 == num_w { + write!(draw_area, "\u{2518}").unwrap(); + } else { + write!(draw_area, "\u{2534}").unwrap(); + } + } + draw_area.newline(); + + 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 +220,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 - } } diff --git a/pixie-uefi/src/os/input.rs b/pixie-uefi/src/os/input.rs new file mode 100644 index 00000000..9545de07 --- /dev/null +++ b/pixie-uefi/src/os/input.rs @@ -0,0 +1,30 @@ +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>>> = Lazy::new(|| { + let input_handles = uefi::boot::find_handles::().unwrap(); + let input = uefi::boot::open_protocol_exclusive::(input_handles[0]).unwrap(); + Mutex::new(SendWrapper(input)) +}); + +pub fn read_key() -> impl Future> { + 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) => { + cx.waker().wake_by_ref(); + Poll::Pending + } + } + }) +} diff --git a/pixie-uefi/src/os/logger.rs b/pixie-uefi/src/os/logger.rs index 631102b5..8278b822 100644 --- a/pixie-uefi/src/os/logger.rs +++ b/pixie-uefi/src/os/logger.rs @@ -1,25 +1,18 @@ -use alloc::{collections::vec_deque::VecDeque, string::String}; +use alloc::string::String; use core::fmt::Write; + use log::Level; use spin::Mutex; -use uefi::{boot::ScopedProtocol, proto::console::serial::Serial}; - -use crate::os::{timer::Timer, UefiOS}; - -static MESSAGES: Mutex> = Mutex::new(VecDeque::new()); -static SERIAL: Mutex> = Mutex::new(None); - -pub(super) struct LogEntry { - pub time: f64, - pub level: Level, - pub target: String, - pub msg: String, -} +use uefi::boot::ScopedProtocol; +use uefi::proto::console::serial::Serial; +use uefi::proto::console::text::Color; -struct SerialWrapper(ScopedProtocol); +use crate::os::send_wrapper::SendWrapper; +use crate::os::timer::Timer; +use crate::os::ui::{self, DrawArea}; -// SAFETY: There are no threads. -unsafe impl Send for SerialWrapper {} +static SERIAL: Mutex>>> = Mutex::new(None); +static DRAW_AREA: Mutex = Mutex::new(DrawArea::invalid()); struct Logger {} @@ -28,10 +21,13 @@ pub(super) fn init() { .ok() .map(|handles| uefi::boot::open_protocol_exclusive::(handles[0]).unwrap()); - *SERIAL.lock() = serial.map(SerialWrapper); + *SERIAL.lock() = serial.map(SendWrapper); log::set_logger(&Logger {}).unwrap(); log::set_max_level(log::LevelFilter::Trace); + + *DRAW_AREA.lock() = DrawArea::logs(); + DRAW_AREA.lock().clear(); } fn append_message(time: f64, level: log::Level, target: &str, msg: String) { @@ -51,32 +47,24 @@ fn append_message(time: f64, level: log::Level, target: &str, msg: String) { } { - let mut logs = MESSAGES.try_lock().expect("messages are locked"); - logs.push_back(LogEntry { - time, - level, - target: target.into(), - msg, - }); - const MAX_MESSAGES: usize = 10; - if logs.len() > MAX_MESSAGES { - logs.pop_front(); - } + let col = match level { + Level::Trace => Color::Cyan, + Level::Debug => Color::Blue, + Level::Info => Color::Green, + Level::Warn => Color::Yellow, + Level::Error => Color::Red, + }; + let mut draw_area = DRAW_AREA.lock(); + write!(draw_area, "[{time:.1}s ").unwrap(); + draw_area.write_with_color(&format!("{level:5} "), col, Color::Black); + writeln!(draw_area, "{target}] {msg}").unwrap(); } if level <= Level::Warn { - UefiOS { cant_build: () }.force_ui_redraw() + ui::flush(); } } -pub(super) fn for_each_log(f: F) { - MESSAGES - .try_lock() - .expect("messages are locked") - .iter() - .for_each(f); -} - impl log::Log for Logger { fn enabled(&self, _metadata: &log::Metadata) -> bool { true diff --git a/pixie-uefi/src/os/mod.rs b/pixie-uefi/src/os/mod.rs index 784b391d..d10a154d 100644 --- a/pixie-uefi/src/os/mod.rs +++ b/pixie-uefi/src/os/mod.rs @@ -1,294 +1,80 @@ -use alloc::boxed::Box; -use alloc::string::{String, ToString}; -use alloc::vec::Vec; -use core::cell::RefMut; use core::ffi::c_void; -use core::fmt::Write; -use core::future::{poll_fn, Future}; +use core::future::Future; use core::ptr::NonNull; -use core::task::Poll; +use core::sync::atomic::{AtomicBool, Ordering}; -use pixie_shared::util::BytesFmt; -use uefi::boot::{EventType, ScopedProtocol, Tpl}; -use uefi::proto::console::text::{Color, Input, Key, Output}; +use uefi::boot::{EventType, Tpl}; use uefi::{Event, Status}; -use crate::os::logger::{for_each_log, LogEntry}; - use self::error::Result; use self::executor::Executor; -use self::sync::SyncRefCell; use self::timer::Timer; pub mod boot_options; pub mod disk; pub mod error; pub mod executor; +pub mod input; mod logger; pub mod memory; pub mod net; -pub mod sync; +mod send_wrapper; mod timer; +pub mod ui; -struct UefiOSImpl { - input: ScopedProtocol, - vga: ScopedProtocol, - ui_buf: Vec<(String, Color, Color)>, - ui_pos: usize, - ui_drawer: Option>, -} +static INITIALIZED: AtomicBool = AtomicBool::new(false); -impl UefiOSImpl { - fn cols(&mut self) -> usize { - let mode = self.vga.current_mode().unwrap().unwrap(); - mode.columns() - } +pub fn start(mut f: F) -> ! +where + F: FnMut() -> Fut + 'static, + Fut: Future>, +{ + assert!(!INITIALIZED.swap(true, Ordering::Relaxed)); - pub fn write_with_color(&mut self, msg: &str, fg: Color, bg: Color) { - let lines: Vec<_> = msg.split('\n').collect(); - for (n, line) in lines.iter().enumerate() { - self.ui_buf.push((line.to_string(), fg, bg)); - self.ui_pos += line.len(); - if n != lines.len() - 1 { - let cols = self.cols(); - let colp = self.ui_pos % cols; - let n = cols - colp; - self.ui_buf - .push((String::from_utf8(vec![0x20; n]).unwrap(), fg, bg)); - self.ui_pos += n; - } - } - } + uefi::helpers::init().unwrap(); - pub fn maybe_advance_to_col(&mut self, col: usize) { - let (fg, bg) = if let Some((_, f, b)) = self.ui_buf[..].last() { - (*f, *b) - } else { - (Color::White, Color::Black) - }; - let cols = self.cols(); - let colp = self.ui_pos % cols; - let n = col - colp; - if colp < col { - self.ui_buf - .push((String::from_utf8(vec![0x20; n]).unwrap(), fg, bg)); - self.ui_pos += n; - } + unsafe extern "efiapi" fn exit_boot_services(_e: Event, _ctx: Option>) { + panic!("You must never exit boot services"); } - pub fn flush_ui_buf(&mut self) { - self.vga.set_cursor_position(0, 0).unwrap(); - let mode = self.vga.current_mode().unwrap().unwrap(); - let (cols, rows) = (mode.columns(), mode.rows()); - for (msg, fg, bg) in self.ui_buf.drain(..) { - self.vga.set_color(fg, bg).unwrap(); - write!(self.vga, "{msg}").unwrap(); - } - self.vga.set_color(Color::White, Color::Black).unwrap(); - if self.ui_pos + 1 < cols * rows { - // Clear any remaining chars. - let n = cols * rows - self.ui_pos - 1; - write!(self.vga, "{}", String::from_utf8(vec![0x20; n]).unwrap()).unwrap(); - } - self.ui_pos = 0; + unsafe { + uefi::boot::create_event( + EventType::SIGNAL_EXIT_BOOT_SERVICES, + Tpl::NOTIFY, + Some(exit_boot_services), + None, + ) + .unwrap(); } -} - -static OS: SyncRefCell> = SyncRefCell::new(None); - -#[non_exhaustive] -#[derive(Clone, Copy)] -pub struct UefiOS { - #[allow(dead_code)] - cant_build: (), -} - -unsafe extern "efiapi" fn exit_boot_services(_e: Event, _ctx: Option>) { - panic!("You must never exit boot services"); -} - -impl UefiOS { - pub fn start(mut f: F) -> ! - where - F: FnMut(UefiOS) -> Fut + 'static, - Fut: Future>, - { - // Never call this function twice. - assert!(OS.borrow().is_none()); - - uefi::helpers::init().unwrap(); - - // Ensure we never exit boot services. - // SAFETY: the callback panics on exit from boot services, and thus handles exit from boot - // services correctly by definition. - unsafe { - uefi::boot::create_event( - EventType::SIGNAL_EXIT_BOOT_SERVICES, - Tpl::NOTIFY, - Some(exit_boot_services), - None, - ) - .unwrap(); - } - - Timer::ensure_init(); - - let input_handles = uefi::boot::find_handles::().unwrap(); - let input = uefi::boot::open_protocol_exclusive::(input_handles[0]).unwrap(); - let vga_handles = uefi::boot::find_handles::().unwrap(); - let mut vga = uefi::boot::open_protocol_exclusive::(vga_handles[0]).unwrap(); + Timer::ensure_init(); + ui::init(); + logger::init(); + net::init(); - vga.clear().unwrap(); - - *OS.borrow_mut() = Some(UefiOSImpl { - input, - vga, - ui_buf: vec![], - ui_pos: 0, - ui_drawer: None, - }); - - let os = UefiOS { cant_build: () }; - - logger::init(); - net::init(); - - Executor::spawn("init", async move { - loop { - if let Err(err) = f(os).await { - log::error!("Error: {err:?}"); - } + Executor::spawn("init", async move { + loop { + if let Err(err) = f().await { + log::error!("Error: {err:?}"); } - }); - - Executor::spawn("[watchdog]", async move { - loop { - let err = uefi::boot::set_watchdog_timer(300, 0x10000, None); + } + }); - if let Err(err) = err { - if err.status() != Status::UNSUPPORTED { - log::error!("Error disabling watchdog: {err:?}"); - } + Executor::spawn("[watchdog]", async move { + loop { + let err = uefi::boot::set_watchdog_timer(300, 0x10000, None); - break; + if let Err(err) = err { + if err.status() != Status::UNSUPPORTED { + log::error!("Error disabling watchdog: {err:?}"); } - Executor::sleep_us(30_000_000).await; + break; } - }); - Executor::spawn("[draw_ui]", async move { - loop { - os.draw_ui(); - Executor::sleep_us(1_000_000).await; - } - }); - - Executor::run() - } - - fn borrow_mut(&self) -> RefMut<'static, UefiOSImpl> { - RefMut::map(OS.borrow_mut(), |f| f.as_mut().unwrap()) - } - - pub fn read_key(&self) -> impl Future> + '_ { - poll_fn(move |cx| { - let key = self.borrow_mut().input.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 - }) - } - - pub fn write_with_color(&self, msg: &str, fg: Color, bg: Color) { - self.borrow_mut().write_with_color(msg, fg, bg); - } - - fn draw_ui(&self) { - // Write the header. - { - let time = Timer::micros() as f32 * 0.000_001; - let ip = net::ip(); - let mut os = self.borrow_mut(); - - let mode = os.vga.current_mode().unwrap().unwrap(); - let cols = mode.columns(); - - os.write_with_color(&format!("uptime: {time:10.1}s"), Color::White, Color::Black); - os.maybe_advance_to_col(cols / 3); - - if let Some(ip) = ip { - os.write_with_color(&format!("IP: {ip}"), Color::White, Color::Black); - } else { - os.write_with_color("DHCP...", Color::Yellow, Color::Black); - } - - os.maybe_advance_to_col(3 * cols / 5); - - let (vtx, vrx) = net::speed(); - os.write_with_color( - &format!("rx: {}/s tx: {}/s\n\n", BytesFmt(vrx), BytesFmt(vtx)), - Color::White, - Color::Black, - ); - - for (micros, name) in Executor::top_tasks(7) { - os.write_with_color(name, Color::White, Color::Black); - os.maybe_advance_to_col(cols / 4); - os.write_with_color( - &format!("{:7.3}s\n", micros as f64 * 0.000_001), - Color::White, - Color::Black, - ); - } - - os.maybe_advance_to_col(cols); - - // TODO(veluca): find a better solution. - for_each_log( - |LogEntry { - time, - level, - target, - msg, - }| { - let fg_color = match level { - log::Level::Trace => Color::Cyan, - log::Level::Debug => Color::Blue, - log::Level::Info => Color::Green, - log::Level::Warn => Color::Yellow, - log::Level::Error => Color::Red, - }; - os.write_with_color(&format!("[{time:.1}s "), Color::White, Color::Black); - os.write_with_color(&format!("{level:5}"), fg_color, Color::Black); - os.write_with_color(&format!(" {target}] {msg}\n"), Color::White, Color::Black); - }, - ); - os.write_with_color("\n", Color::Black, Color::Black); - } - { - let ui = self.borrow_mut().ui_drawer.take(); - if let Some(ui) = &ui { - ui(*self); - } - self.borrow_mut().ui_drawer = ui; + Executor::sleep_us(30_000_000).await; } - // Actually draw the changes. - self.borrow_mut().flush_ui_buf(); - } + }); - pub fn force_ui_redraw(&self) { - self.draw_ui() - } - - pub fn set_ui_drawer(&self, f: F) { - self.borrow_mut().ui_drawer = Some(Box::new(f)); - } + Executor::run() } diff --git a/pixie-uefi/src/os/net/mod.rs b/pixie-uefi/src/os/net/mod.rs index a790414e..d4edd68d 100644 --- a/pixie-uefi/src/os/net/mod.rs +++ b/pixie-uefi/src/os/net/mod.rs @@ -1,4 +1,5 @@ use alloc::string::{String, ToString}; +use core::fmt::Write; use core::future::{poll_fn, Future}; use core::net::Ipv4Addr; use core::sync::atomic::{AtomicU64, Ordering}; @@ -8,6 +9,7 @@ use smoltcp::iface::{Config, Interface, PollResult, SocketHandle, SocketSet}; use smoltcp::socket::dhcpv4::{Event, Socket as Dhcpv4Socket}; use smoltcp::wire::{DhcpOption, HardwareAddress, IpCidr}; use spin::Mutex; +use uefi::proto::console::text::Color; use uefi::proto::device_path::build::DevicePathBuilder; use uefi::proto::device_path::text::{AllowShortcuts, DevicePathToText, DisplayOnly}; use uefi::proto::device_path::DevicePath; @@ -19,10 +21,10 @@ use super::timer::Timer; use crate::os::boot_options::BootOptions; use crate::os::executor::Executor; use crate::os::net::interface::SnpDevice; -use crate::os::net::speed::{RX_SPEED, TX_SPEED}; pub use crate::os::net::tcp::TcpStream; pub use crate::os::net::udp::UdpSocket; use crate::os::timer::rdtsc; +use crate::os::ui; mod interface; mod speed; @@ -42,10 +44,6 @@ struct NetworkData { static NETWORK_DATA: Mutex> = Mutex::new(None); -pub fn speed() -> (u64, u64) { - (TX_SPEED.bytes_per_second(), RX_SPEED.bytes_per_second()) -} - fn with_net T>(f: F) -> T { let mut mg = NETWORK_DATA.try_lock().expect("Network is locked"); f(mg.as_mut().expect("Network is not initialized")) @@ -130,7 +128,23 @@ pub(super) fn init() { }), ); - speed::spawn_update_network_speed_task(); + Executor::spawn("[show_ip]", async { + let mut draw_area = ui::DrawArea::ip(); + loop { + draw_area.clear(); + let ip = ip(); + let w = draw_area.size().0; + if let Some(ip) = ip { + write!(draw_area, "IP: {ip:>0$}", w - 4).unwrap(); + Executor::sleep_us(10_000_000).await + } else { + draw_area.write_with_color("DHCP...", Color::Yellow, Color::Black); + Executor::sleep_us(100_000).await + } + } + }); + + speed::spawn_network_speed_task(); } pub fn wait_for_ip() -> impl Future { @@ -144,9 +158,8 @@ pub fn wait_for_ip() -> impl Future { }) } -pub fn ip() -> Option { - let mg = NETWORK_DATA.try_lock().unwrap(); - mg.as_ref().and_then(|n| n.interface.ipv4_addr()) +fn ip() -> Option { + with_net(|n| n.interface.ipv4_addr()) } fn get_ephemeral_port() -> u16 { diff --git a/pixie-uefi/src/os/net/speed.rs b/pixie-uefi/src/os/net/speed.rs index 30765ad2..89770cb7 100644 --- a/pixie-uefi/src/os/net/speed.rs +++ b/pixie-uefi/src/os/net/speed.rs @@ -1,7 +1,12 @@ +use core::fmt::Write; use core::sync::atomic::{AtomicU64, Ordering}; +use pixie_shared::util::BytesFmt; +use uefi::proto::console::text::Color; + use crate::os::executor::Executor; use crate::os::timer::Timer; +use crate::os::ui; pub struct NetSpeed { total: AtomicU64, @@ -44,11 +49,28 @@ impl NetSpeed { pub(super) static TX_SPEED: NetSpeed = NetSpeed::new(); pub(super) static RX_SPEED: NetSpeed = NetSpeed::new(); -pub(super) fn spawn_update_network_speed_task() { +pub(super) fn spawn_network_speed_task() { Executor::spawn("[net_speed]", async move { + let mut draw_area = ui::DrawArea::net_speed(); loop { TX_SPEED.update_speed(); RX_SPEED.update_speed(); + draw_area.clear(); + let w = draw_area.size().0; + let vtx = TX_SPEED.bytes_per_second(); + let vrx = RX_SPEED.bytes_per_second(); + draw_area.write_with_color( + &format!("Network \u{2193}:{0:1$}", "", w - 22), + Color::Green, + Color::Black, + ); + writeln!(draw_area, "{:10.1}/s", BytesFmt(vrx)).unwrap(); + draw_area.write_with_color( + &format!("Network \u{2191}:{0:1$}", "", w - 22), + Color::Cyan, + Color::Black, + ); + writeln!(draw_area, "{:10.1}/s", BytesFmt(vtx)).unwrap(); Executor::sleep_us(1_000_000).await; } }); diff --git a/pixie-uefi/src/os/send_wrapper.rs b/pixie-uefi/src/os/send_wrapper.rs new file mode 100644 index 00000000..2fb485e0 --- /dev/null +++ b/pixie-uefi/src/os/send_wrapper.rs @@ -0,0 +1,19 @@ +use core::ops::{Deref, DerefMut}; + +pub struct SendWrapper(pub T); + +// SAFETY: there are no threads. +unsafe impl Send for SendWrapper {} + +impl Deref for SendWrapper { + type Target = T; + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl DerefMut for SendWrapper { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 + } +} diff --git a/pixie-uefi/src/os/sync.rs b/pixie-uefi/src/os/sync.rs deleted file mode 100644 index 5b310e13..00000000 --- a/pixie-uefi/src/os/sync.rs +++ /dev/null @@ -1,20 +0,0 @@ -use core::cell::RefCell; -use core::ops::Deref; - -pub struct SyncRefCell(RefCell); - -impl SyncRefCell { - pub const fn new(value: T) -> Self { - Self(RefCell::new(value)) - } -} - -unsafe impl Sync for SyncRefCell {} - -impl Deref for SyncRefCell { - type Target = RefCell; - - fn deref(&self) -> &Self::Target { - &self.0 - } -} diff --git a/pixie-uefi/src/os/ui.rs b/pixie-uefi/src/os/ui.rs new file mode 100644 index 00000000..ea2f69a1 --- /dev/null +++ b/pixie-uefi/src/os/ui.rs @@ -0,0 +1,308 @@ +use alloc::vec::Vec; +use core::fmt::Write; +use core::sync::atomic::{AtomicUsize, Ordering}; + +use pixie_shared::util::BytesFmt; +use spin::lazy::Lazy; +use spin::Mutex; +use uefi::boot::ScopedProtocol; +use uefi::proto::console::text::{Color, Output}; +use uefi::{CStr16, Char16}; + +use super::executor::{Executor, TASK_LEN}; +use super::memory; +use super::send_wrapper::SendWrapper; +use super::timer::Timer; + +const TASK_HEIGHT: usize = 7; // includes box +const LOG_HEIGHT: usize = 10; + +#[derive(Clone, Copy)] +struct ScreenChar { + c: Char16, + fg: Color, + bg: Color, +} + +impl Default for ScreenChar { + fn default() -> Self { + Self { + c: Char16::try_from(0x20).unwrap(), + fg: Color::White, + bg: Color::Black, + } + } +} + +impl PartialEq for ScreenChar { + fn eq(&self, other: &Self) -> bool { + other.c == self.c && other.fg as u8 == self.fg as u8 && other.bg as u8 == self.bg as u8 + } +} + +pub struct Screen { + vga: SendWrapper>, + front_buffer: Vec, + back_buffer: Vec, +} + +static WIDTH: AtomicUsize = AtomicUsize::new(0); +static HEIGHT: AtomicUsize = AtomicUsize::new(0); + +fn w() -> usize { + WIDTH.load(Ordering::Relaxed) +} + +fn h() -> usize { + HEIGHT.load(Ordering::Relaxed) +} + +static SCREEN: Lazy> = Lazy::new(|| { + let vga_handles = uefi::boot::find_handles::().unwrap(); + let vga = uefi::boot::open_protocol_exclusive::(vga_handles[0]).unwrap(); + let mode = vga.current_mode().unwrap().unwrap(); + let (w, h) = (mode.columns(), mode.rows()); + WIDTH.store(w, Ordering::Relaxed); + HEIGHT.store(h, Ordering::Relaxed); + Mutex::new(Screen { + vga: SendWrapper(vga), + front_buffer: vec![ScreenChar::default(); w * h], + back_buffer: vec![ScreenChar::default(); w * h], + }) +}); + +static CONTENT_DRAW_AREA: Mutex = Mutex::new(DrawArea::invalid()); + +pub(super) fn init() { + let mut screen = SCREEN.lock(); + screen.vga.clear().unwrap(); + let _ = screen.vga.enable_cursor(false); + + Executor::spawn("[flush_ui]", async move { + loop { + SCREEN.lock().flush(); + Executor::sleep_us(100_000).await; + } + }); + + *CONTENT_DRAW_AREA.lock() = DrawArea::content(); + + Executor::spawn("[show_timer]", async move { + let mut draw_area = DrawArea::time(); + loop { + draw_area.clear(); + let time = Timer::micros() as f32 * 0.000_001; + let w = draw_area.size().0; + write!(draw_area, "uptime:{0:1$}{time:12.1}s", "", w - 20).unwrap(); + Executor::sleep_us(50_000).await; + } + }); + + Executor::spawn("[show_memory]", async { + let mut draw_area = DrawArea::memory(); + loop { + draw_area.clear(); + let memory = memory::stats(); + let w = draw_area.size().0; + write!( + draw_area, + "RAM:{2:3$}{:10.1} / {:10.1}", + BytesFmt(memory.used), + BytesFmt(memory.free + memory.used), + "", + w - 27, + ) + .unwrap(); + Executor::sleep_us(1_000_000).await; + } + }); + + let logs_area = DrawArea::logs(); + for y in 0..logs_area.size.1 { + screen.back_buffer[(logs_area.offset.1 + y) * w()].c = Char16::try_from(0x25ba).unwrap(); + } +} + +impl Screen { + fn flush(&mut self) { + let mut i = 0; + while i < w() * h() { + if self.back_buffer[i] == self.front_buffer[i] { + i += 1; + continue; + } + + let start = i; + let fg = self.back_buffer[i].fg; + let bg = self.back_buffer[i].bg; + + let mut run = Vec::new(); + while i < w() * h() + && self.back_buffer[i].fg as u8 == fg as u8 + && self.back_buffer[i].bg as u8 == bg as u8 + && self.back_buffer[i] != self.front_buffer[i] + { + if i > start && i % w() == 0 { + break; + } + run.push(self.back_buffer[i].c); + self.front_buffer[i] = self.back_buffer[i]; + i += 1; + } + + run.push(Char16::try_from(0).unwrap()); + + let x = start % w(); + let y = start / w(); + self.vga.set_cursor_position(x, y).unwrap(); + self.vga.set_color(fg, bg).unwrap(); + self.vga + .output_string(CStr16::from_char16_until_nul(&run).unwrap()) + .unwrap(); + } + } +} + +const STATUS_WIDTH: usize = 32; + +pub struct DrawArea { + offset: (usize, usize), + size: (usize, usize), + pos: (usize, usize), + scroll: bool, +} + +impl DrawArea { + pub(super) const fn invalid() -> Self { + Self::new((0, 0), (0, 0), false) + } + + const fn new(offset: (usize, usize), size: (usize, usize), scroll: bool) -> Self { + Self { + offset, + size, + pos: (0, 0), + scroll, + } + } + + fn task_columns() -> usize { + (w() - STATUS_WIDTH - 1) / (TASK_LEN + 1) + } + + fn task_width() -> usize { + Self::task_columns() * (TASK_LEN + 1) + 1 + } + + fn task_side(col: usize, num_cols: usize) -> DrawArea { + Self::new( + (Self::task_width() + 1, col), + (w() - Self::task_width() - 1, num_cols), + false, + ) + } + + pub(super) fn ip() -> DrawArea { + Self::task_side(1, 1) + } + + pub(super) fn net_speed() -> DrawArea { + Self::task_side(2, 2) + } + + fn time() -> DrawArea { + Self::task_side(4, 1) + } + + pub(super) fn memory() -> DrawArea { + Self::task_side(5, 1) + } + + pub(super) fn tasks() -> DrawArea { + Self::new((0, 0), (Self::task_width(), TASK_HEIGHT), false) + } + + pub(super) fn logs() -> DrawArea { + Self::new((2, h() - LOG_HEIGHT), (w() - 3, LOG_HEIGHT), true) + } + + pub(super) fn content() -> DrawArea { + let offset = 1 + TASK_HEIGHT; + Self::new((0, offset), (w(), h() - offset - LOG_HEIGHT), false) + } + + pub fn size(&self) -> (usize, usize) { + self.size + } + + fn idx(&self, p: (usize, usize)) -> usize { + (self.offset.1 + p.1) * w() + self.offset.0 + p.0 + } + + pub fn clear(&mut self) { + self.pos = (0, 0); + let mut screen = SCREEN.lock(); + for y in 0..self.size.1 { + for x in 0..self.size.0 { + screen.back_buffer[self.idx((x, y))] = ScreenChar::default(); + } + } + } + + pub fn write_with_color(&mut self, msg: &str, fg: Color, bg: Color) { + let mut screen = SCREEN.lock(); + for c in msg.chars() { + if c == '\n' { + self.newline(); + continue; + } + if self.pos.0 == self.size.0 { + self.newline(); + } + while self.scroll && self.pos.1 >= self.size.1 { + self.pos.1 -= 1; + for y in 0..self.size.1 - 1 { + for x in 0..self.size.0 { + screen.back_buffer[self.idx((x, y))] = + screen.back_buffer[self.idx((x, y + 1))]; + } + } + for x in 0..self.size.0 { + screen.back_buffer[self.idx((x, self.size.1 - 1))] = ScreenChar::default(); + } + } + if self.pos.1 >= self.size.1 { + continue; + } + if let Ok(c16) = Char16::try_from(c) { + screen.back_buffer[self.idx(self.pos)] = ScreenChar { c: c16, fg, bg }; + } + self.pos.0 += 1; + } + } + + pub fn newline(&mut self) { + self.pos.0 = 0; + self.pos.1 += 1; + } + pub fn advance(&mut self, n: usize) { + self.pos.0 += n; + } +} + +impl Write for DrawArea { + fn write_str(&mut self, s: &str) -> core::fmt::Result { + self.write_with_color(s, Color::White, Color::Black); + Ok(()) + } +} + +pub fn flush() { + SCREEN.lock().flush(); +} + +pub fn update_content(f: F) { + f(&mut CONTENT_DRAW_AREA + .try_lock() + .expect("content draw area is locked")); +} diff --git a/pixie-uefi/src/register.rs b/pixie-uefi/src/register.rs index 5b0dee76..77377b69 100644 --- a/pixie-uefi/src/register.rs +++ b/pixie-uefi/src/register.rs @@ -10,7 +10,8 @@ use uefi::proto::console::text::{Color, Key, ScanCode}; use crate::os::error::{Error, Result}; use crate::os::net::{TcpStream, UdpSocket, ETH_PACKET_SIZE}; -use crate::os::UefiOS; +use crate::os::ui::DrawArea; +use crate::os::{input, ui}; #[derive(Debug, Default)] struct Data { @@ -18,49 +19,41 @@ struct Data { selected: usize, } -pub async fn register(os: UefiOS, server_addr: SocketAddrV4) -> Result<()> { +pub async fn register(server_addr: SocketAddrV4) -> Result<()> { let data = Rc::new(RefCell::new(Data::default())); let data2 = data.clone(); - os.set_ui_drawer(move |os| { + let draw_content = move |draw_area: &mut DrawArea| { + draw_area.clear(); let data2 = data2.borrow(); - os.write_with_color( - &format!("Group: {}\n", data2.station.group), - if data2.selected == 0 { + let fg = |i| { + if data2.selected == i { Color::Yellow } else { Color::White - }, + } + }; + draw_area.write_with_color( + &format!("Group: {}\n", data2.station.group), + fg(0), Color::Black, ); - os.write_with_color( + draw_area.write_with_color( &format!("Row: {}\n", data2.station.row), - if data2.selected == 1 { - Color::Yellow - } else { - Color::White - }, + fg(1), Color::Black, ); - os.write_with_color( + draw_area.write_with_color( &format!("Column: {}\n", data2.station.col), - if data2.selected == 2 { - Color::Yellow - } else { - Color::White - }, + fg(2), Color::Black, ); - os.write_with_color( + draw_area.write_with_color( &format!("Image: {}\n", data2.station.image), - if data2.selected == 3 { - Color::Yellow - } else { - Color::White - }, + fg(3), Color::Black, ); - }); + }; let udp = UdpSocket::bind(Some(HINT_PORT)).await?; let mut buf = [0; ETH_PACKET_SIZE]; @@ -69,18 +62,22 @@ pub async fn register(os: UefiOS, server_addr: SocketAddrV4) -> Result<()> { let mut images = Vec::new(); let mut groups = Vec::new(); + ui::update_content(&draw_content); + ui::flush(); + loop { let key = if hint { loop { let recv = Box::pin(udp.recv_from(&mut buf)); - let key = Box::pin(os.read_key()); + let key = Box::pin(input::read_key()); match select(recv, key).await { Either::Left(((buf, _), _)) => { let hint: HintPacket = postcard::from_bytes(buf)?; data.borrow_mut().station = hint.station; images = hint.images; groups = hint.groups.into_iter().map(|(k, _)| k).collect(); - os.force_ui_redraw(); + ui::update_content(&draw_content); + ui::flush(); } Either::Right((key, _)) => { hint = false; @@ -89,7 +86,7 @@ pub async fn register(os: UefiOS, server_addr: SocketAddrV4) -> Result<()> { } } } else { - os.read_key().await? + input::read_key().await? }; if key == Key::Special(ScanCode::DOWN) { @@ -157,7 +154,8 @@ pub async fn register(os: UefiOS, server_addr: SocketAddrV4) -> Result<()> { if key == Key::Printable('\r'.try_into().unwrap()) { break; } - os.force_ui_redraw(); + ui::update_content(&draw_content); + ui::flush(); } let msg = TcpRequest::Register(data.borrow().station.clone()); diff --git a/pixie-uefi/src/store.rs b/pixie-uefi/src/store.rs index 2d17e8f5..39880eb9 100644 --- a/pixie-uefi/src/store.rs +++ b/pixie-uefi/src/store.rs @@ -1,18 +1,18 @@ -use alloc::rc::Rc; use alloc::vec::Vec; use core::cell::RefCell; +use core::fmt::Write; use core::net::SocketAddrV4; use log::info; use lz4_flex::compress; use pixie_shared::util::BytesFmt; use pixie_shared::{Chunk, Image, Offset, TcpRequest, UdpRequest, MAX_CHUNK_SIZE}; -use uefi::proto::console::text::Color; use crate::os::boot_options::BootOptions; use crate::os::error::{Error, Result}; use crate::os::net::{TcpStream, UdpSocket}; -use crate::os::{disk, memory, UefiOS}; +use crate::os::ui::DrawArea; +use crate::os::{disk, memory, ui}; use crate::{parse_disk, MIN_MEMORY}; #[derive(Debug)] @@ -41,31 +41,28 @@ enum State { }, } -pub async fn store(os: UefiOS, server_address: SocketAddrV4) -> Result<()> { - let stats = Rc::new(RefCell::new(State::ReadingPartitions)); - let stats2 = stats.clone(); - os.set_ui_drawer(move |os| match &*stats2.borrow() { - State::ReadingPartitions => { - os.write_with_color("Reading partitions...", Color::White, Color::Black) - } - State::PushingChunks { - cur, - total, - tsize, - tcsize, - } => { - os.write_with_color( - &format!("Pushed {cur} out of {total} chunks\n"), - Color::White, - Color::Black, - ); - os.write_with_color( - &format!("total size {tsize}, compressed {tcsize}\n"), - Color::White, - Color::Black, - ); +pub async fn store(server_address: SocketAddrV4) -> Result<()> { + let stats = RefCell::new(State::ReadingPartitions); + + let draw = |draw_area: &mut DrawArea| { + draw_area.clear(); + match &*stats.borrow() { + State::ReadingPartitions => { + writeln!(draw_area, "Reading partitions...").unwrap(); + } + State::PushingChunks { + cur, + total, + tsize, + tcsize, + } => { + writeln!(draw_area, "Pushed {cur} out of {total} chunks").unwrap(); + writeln!(draw_area, "total size {tsize}, compressed {tcsize}").unwrap(); + } } - }); + }; + + ui::update_content(draw); let boid = BootOptions::reboot_target().expect("Could not find reboot target"); let bo_command = BootOptions::get(boid); @@ -189,6 +186,7 @@ pub async fn store(os: UefiOS, server_address: SocketAddrV4) -> Result<()> { &postcard::to_allocvec(&UdpRequest::ActionProgress(chunks.len(), total))?, ) .await?; + ui::update_content(draw); } Ok(chunks) };