[pixie-uefi] Compute total memory available and optimize buffer size.

Author: Virv12

pixie-uefi/src/flash.rs

@@ -1,13 +1,18 @@
-use crate::os::{
-    error::{Error, Result},
-    mpsc, TcpStream, UefiOS, PACKET_SIZE,
+use crate::{
+    os::{
+        error::{Error, Result},
+        mpsc, TcpStream, UefiOS, PACKET_SIZE,
+    },
+    MIN_MEMORY,
 };
 use alloc::{boxed::Box, collections::BTreeMap, rc::Rc, string::ToString, vec::Vec};
 use core::{cell::RefCell, mem, net::SocketAddrV4};
 use futures::future::{select, Either};
 use log::info;
 use lz4_flex::decompress;
-use pixie_shared::{ChunkHash, Image, TcpRequest, UdpRequest, BODY_LEN, CHUNKS_PORT, HEADER_LEN};
+use pixie_shared::{
+    ChunkHash, Image, TcpRequest, UdpRequest, BODY_LEN, CHUNKS_PORT, HEADER_LEN, MAX_CHUNK_SIZE,
+};
 use uefi::proto::console::text::Color;
 
 struct PartialChunk {
@@ -232,6 +237,9 @@ pub async fn flash(os: UefiOS, server_addr: SocketAddrV4) -> Result<()> {
     let task1 = async {
         let mut tx = tx;
         let mut last_seen = Vec::new();
+        let total_mem = os.get_total_mem();
+        let max_chunks = (total_mem.saturating_sub(MIN_MEMORY) as usize / MAX_CHUNK_SIZE).max(128);
+        log::debug!("Total memory: {total_mem}. Max chunks in memory: {max_chunks}");
         while !chunks_info.is_empty() {
             let recv = Box::pin(socket.recv(&mut buf));
             let sleep = Box::pin(os.sleep_us(100_000));
@@ -247,7 +255,7 @@ pub async fn flash(os: UefiOS, server_addr: SocketAddrV4) -> Result<()> {
                     }
 
                     assert_eq!(last_seen.len(), received.len());
-                    if last_seen.len() > 128 {
+                    if last_seen.len() > max_chunks {
                         let hash = last_seen.remove(0);
                         received
                             .remove(&hash)

pixie-uefi/src/main.rs

@@ -28,6 +28,8 @@ mod reboot_to_os;
 mod register;
 mod store;
 
+const MIN_MEMORY: u64 = 500 << 20;
+
 async fn server_discover(os: UefiOS) -> Result<SocketAddrV4> {
     let socket = os.udp_bind(None).await?;
 

pixie-uefi/src/os/mod.rs

@@ -25,7 +25,8 @@ use core::{
     task::{Context, Poll},
 };
 use uefi::{
-    boot::{EventType, ScopedProtocol, TimerTrigger, Tpl},
+    boot::{EventType, MemoryType, ScopedProtocol, TimerTrigger, Tpl},
+    mem::memory_map::MemoryMap,
     proto::{
         console::{
             serial::Serial,
@@ -571,6 +572,14 @@ impl UefiOS {
     pub fn shutdown(&self) -> ! {
         uefi::runtime::reset(uefi::runtime::ResetType::SHUTDOWN, Status::SUCCESS, None)
     }
+
+    pub fn get_total_mem(&self) -> u64 {
+        uefi::boot::memory_map(MemoryType::LOADER_DATA)
+            .expect("Failed to get memory map")
+            .entries()
+            .map(|entry| entry.page_count * 4096)
+            .sum()
+    }
 }
 
 impl log::Log for UefiOS {

pixie-uefi/src/store.rs

@@ -3,7 +3,7 @@ use crate::{
         error::{Error, Result},
         mpsc, TcpStream, UefiOS,
     },
-    parse_disk,
+    parse_disk, MIN_MEMORY,
 };
 use alloc::{rc::Rc, vec::Vec};
 use core::{cell::RefCell, net::SocketAddrV4};
@@ -102,10 +102,15 @@ pub async fn store(os: UefiOS, server_address: SocketAddrV4) -> Result<()> {
     let mut total_size = 0;
     let mut total_csize = 0;
 
-    let (tx1, mut rx1) = mpsc::channel(32);
-    let (tx2, mut rx2) = mpsc::channel(32);
-    let (tx3, mut rx3) = mpsc::channel(32);
-    let (tx4, mut rx4) = mpsc::channel(32);
+    let total_mem = os.get_total_mem();
+    let channel_size =
+        (total_mem.saturating_sub(MIN_MEMORY) as usize / (4 * MAX_CHUNK_SIZE)).max(32);
+    log::debug!("Total memory: {total_mem}. Channel size: {channel_size}");
+
+    let (tx1, mut rx1) = mpsc::channel(channel_size);
+    let (tx2, mut rx2) = mpsc::channel(channel_size);
+    let (tx3, mut rx3) = mpsc::channel(channel_size);
+    let (tx4, mut rx4) = mpsc::channel(channel_size);
 
     let task1 = async {
         let mut tx1 = tx1;