Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 31 additions & 5 deletions pixie-shared/src/util.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
}
}
}
79 changes: 38 additions & 41 deletions pixie-uefi/src/flash.rs
Original file line number Diff line number Diff line change
@@ -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;

Expand All @@ -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<Image> {
Expand All @@ -32,6 +32,7 @@ async fn fetch_image(stream: &TcpStream) -> Result<Image> {
Ok(postcard::from_bytes(&buf)?)
}

#[derive(Clone, PartialEq, Eq)]
struct Stats {
chunks: usize,
unique: usize,
Expand Down Expand Up @@ -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;
Expand All @@ -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) {
Expand All @@ -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);
Expand Down Expand Up @@ -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");

Expand Down
16 changes: 7 additions & 9 deletions pixie-uefi/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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;
Expand All @@ -113,9 +113,7 @@ async fn run(os: UefiOS) -> Result<()> {
});

loop {
// Clear any UI drawer.
os.set_ui_drawer(|_| {});
Comment thread
veluca93 marked this conversation as resolved.

update_content(|d| d.clear());
if !last_was_wait {
log::debug!("Sending request for command");
}
Expand Down Expand Up @@ -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?;
Expand All @@ -167,5 +165,5 @@ async fn run(os: UefiOS) -> Result<()> {

#[entry]
fn main() -> Status {
UefiOS::start(run)
os::start(run)
}
3 changes: 1 addition & 2 deletions pixie-uefi/src/os/disk.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
Loading