Use talc to handle allocations.

Author: veluca93

pixie-uefi/Cargo.lock

@@ -27,7 +27,9 @@ postcard = { version = "1.1.3", default-features = false, features = ["alloc"] }
 rand = { version = "0.8.5", default-features = false }
 rand_xoshiro = { version = "0.6.0", default-features = false }
 smoltcp = { version = "0.12.0", default-features = false, features = ["alloc", "proto-ipv4", "medium-ethernet", "socket-udp", "socket-tcp", "socket-dhcpv4", "async", "socket-tcp-cubic"] }
-uefi = { version = "0.33.0", features = ["alloc", "global_allocator", "panic_handler"] }
+spin = "0.10.0"
+talc = { version = "4.4.3", default-features = false, features = ["lock_api"] }
+uefi = { version = "0.33.0", features = ["panic_handler", "alloc"] }
 
 [dependencies.pixie-shared]
 path = "../pixie-shared"

pixie-uefi/Cargo.toml

@@ -0,0 +1,37 @@
+use talc::*;
+use uefi::boot::MemoryType;
+
+const MAX_BLOCK_SIZE: usize = 1 << 30;
+
+struct AllocOnOom {
+    next_block_size: usize,
+}
+
+impl OomHandler for AllocOnOom {
+    fn handle_oom(talc: &mut Talc<Self>, layout: core::alloc::Layout) -> Result<(), ()> {
+        let bs = talc.oom_handler.next_block_size;
+        talc.oom_handler.next_block_size =
+            (talc.oom_handler.next_block_size * 2).min(MAX_BLOCK_SIZE);
+        match uefi::boot::allocate_pool(MemoryType::LOADER_DATA, bs) {
+            Err(e) => {
+                uefi::println!(
+                    "Cannot allocate new block ({e}). Triggered by allocation {layout:?}"
+                );
+                return Err(());
+            }
+            Ok(ptr) => {
+                let span = talc::Span::from_base_size(ptr.as_ptr(), bs);
+                // SAFETY: the memory was just allocated, so we have exclusive access to it
+                // and we can transfer ownership.
+                unsafe { talc.claim(span) }?;
+                return Ok(());
+            }
+        }
+    }
+}
+
+#[global_allocator]
+static ALLOCATOR: Talck<spin::Mutex<()>, AllocOnOom> = Talc::new(AllocOnOom {
+    next_block_size: 1 << 20,
+})
+.lock();

pixie-uefi/src/os/allocator.rs

@@ -44,6 +44,7 @@ use uefi::{
     CStr16, Event, Handle, Status,
 };
 
+mod allocator;
 mod boot_options;
 pub mod disk;
 pub mod error;

pixie-uefi/src/os/mod.rs

@@ -19,7 +19,7 @@ RUSTFLAGS='-C instrument-coverage' LLVM_PROFILE_FILE=../prof-out/test-pixie-shar
 popd
 
 pushd pixie-uefi
-RUSTFLAGS='-Cinstrument-coverage -Zno-profiler-runtime' cargo +nightly build -F coverage
+RUSTFLAGS='-Cinstrument-coverage -Zno-profiler-runtime' cargo +nightly build -F coverage -Zbuild-std
 popd
 
 pushd pixie-web