Optimize async future memory usage with bump allocator for embedded systems

.gitignore

@@ -1,3 +1,4 @@
 /.vscode
 /target
 /Cargo.lock
+*.rlib

examples/light_eth.rs

@@ -9,7 +9,9 @@
 //!
 //! The example implements a fictitious Light device (an On-Off Matter cluster).
 #![recursion_limit = "256"]
+#![feature(maybe_uninit_uninit_array_transpose)]
 
+use core::mem::MaybeUninit;
 use core::pin::pin;
 
 use embassy_futures::select::select;
@@ -18,6 +20,7 @@ use embassy_time::{Duration, Timer};
 use env_logger::Target;
 use log::info;
 
+use rs_matter_stack::bump_alloc::BumpAlloc;
 use rs_matter_stack::eth::EthMatterStack;
 use rs_matter_stack::matter::dm::clusters::desc;
 use rs_matter_stack::matter::dm::clusters::desc::ClusterHandler as _;
@@ -77,11 +80,17 @@ fn main() -> Result<(), Error> {
             Async(desc::DescHandler::new(Dataver::new_rand(stack.matter().rand())).adapt()),
         );
 
+    static ALLOC_MEM: StaticCell<[u8; 90000]> = StaticCell::new();
+    let alloc_mem = ALLOC_MEM.uninit();
+    let ptr = alloc_mem.as_mut_ptr() as *mut MaybeUninit<u8>;
+    let p: &'static mut [MaybeUninit<u8>] = unsafe { core::slice::from_raw_parts_mut(ptr as _, 90000) };
+    let bump_alloc = BumpAlloc::new(p);
+
     // Run the Matter stack with our handler
     // Using `pin!` is completely optional, but saves some memory due to `rustc`
     // not being very intelligent w.r.t. stack usage in async functions
     let store = stack.create_shared_store(DirKvBlobStore::new_default());
-    let mut matter = pin!(stack.run_preex(
+    let mut matter = pin!(stack.run_preex2(
         // The Matter stack needs UDP sockets to communicate with other Matter devices
         edge_nal_std::Stack::new(),
         // Will try to find a default network interface
@@ -94,6 +103,7 @@ fn main() -> Result<(), Error> {
         (NODE, handler),
         // No user task future to run
         (),
+        &bump_alloc,
     ));
 
     // Just for demoing purposes:

src/bump_alloc.rs

@@ -0,0 +1,133 @@
+//! A simple bump allocator for fixed-size memory chunks
+//!
+//! This allocator is designed for allocating objects that have a known maximum lifetime.
+//! All allocations are freed when the allocator is dropped.
+
+use core::alloc::Layout;
+use core::cell::RefCell;
+use core::mem::MaybeUninit;
+use core::pin::Pin;
+use core::ptr::NonNull;
+
+struct Inner<'a> {
+    memory: &'a mut [MaybeUninit<u8>],
+    offset: usize,
+}
+
+/// A bump allocator that uses a provided memory chunk
+pub struct BumpAlloc<'a>(RefCell<Inner<'a>>);
+
+impl<'a> BumpAlloc<'a> {
+    /// Create a new bump allocator with the provided memory chunk
+    pub const fn new(memory: &'a mut [MaybeUninit<u8>]) -> Self {
+        Self(RefCell::new(Inner { memory, offset: 0 }))
+    }
+
+    /// Allocate memory for an object and pin it
+    pub fn box_pin<T>(&self, object: T) -> Result<Pin<BumpBox<'_, T>>, AllocError>
+    where
+        T: Sized,
+    {
+        let layout = Layout::new::<T>();
+        let ptr = self.alloc_raw(layout)?;
+
+        // Safety: We just allocated the memory and it's properly aligned
+        unsafe {
+            let ptr = ptr.as_ptr() as *mut T;
+            ptr.write(object);
+            Ok(Pin::new_unchecked(BumpBox {
+                ptr: NonNull::new_unchecked(ptr),
+                _allocator: core::marker::PhantomData,
+            }))
+        }
+    }
+
+    fn alloc_raw(&self, layout: Layout) -> Result<NonNull<u8>, AllocError> {
+        let size = layout.size();
+        let align = layout.align();
+
+        let mut inner = self.0.borrow_mut();
+
+        let current_offset = inner.offset;
+
+        // Align the offset
+        let aligned_offset = (current_offset + align - 1) & !(align - 1);
+        let new_offset = aligned_offset + size;
+
+        if new_offset > inner.memory.len() {
+            return Err(AllocError);
+        }
+
+        inner.offset = new_offset;
+
+        // Safety: We checked bounds and alignment
+        let ptr = unsafe { inner.memory.as_mut_ptr().add(aligned_offset) as *mut u8 };
+
+        Ok(unsafe { NonNull::new_unchecked(ptr) })
+    }
+
+    /// Get the current memory usage
+    pub fn used(&self) -> usize {
+        self.0.borrow().offset
+    }
+
+    /// Get the total capacity
+    pub fn capacity(&self) -> usize {
+        self.0.borrow().memory.len()
+    }
+}
+
+/// A box-like container that uses bump allocation
+pub struct BumpBox<'a, T> {
+    ptr: NonNull<T>,
+    _allocator: core::marker::PhantomData<&'a ()>,
+}
+
+impl<T> core::ops::Deref for BumpBox<'_, T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        unsafe { self.ptr.as_ref() }
+    }
+}
+
+impl<T> core::ops::DerefMut for BumpBox<'_, T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        unsafe { self.ptr.as_mut() }
+    }
+}
+
+impl<T: core::future::Future> core::future::Future for BumpBox<'_, T> {
+    type Output = T::Output;
+
+    fn poll(
+        self: Pin<&mut Self>,
+        cx: &mut core::task::Context<'_>,
+    ) -> core::task::Poll<Self::Output> {
+        // Safety: We maintain the pin invariant
+        let future = unsafe { self.map_unchecked_mut(|s| &mut **s) };
+        future.poll(cx)
+    }
+}
+
+impl<T> Drop for BumpBox<'_, T> {
+    fn drop(&mut self) {
+        // Safety: The pointer is valid and we own the data
+        unsafe {
+            self.ptr.as_ptr().drop_in_place();
+        }
+    }
+}
+
+/// Error type for allocation failures
+#[derive(Debug)]
+pub struct AllocError;
+
+impl core::fmt::Display for AllocError {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        write!(f, "Bump allocator out of memory")
+    }
+}
+
+#[cfg(feature = "std")]
+impl std::error::Error for AllocError {}

src/eth.rs

@@ -14,6 +14,7 @@ use rs_matter::transport::network::NoNetwork;
 use rs_matter::utils::init::{init, Init};
 use rs_matter::utils::select::Coalesce;
 
+use crate::bump_alloc::BumpAlloc;
 use crate::mdns::Mdns;
 use crate::nal::NetStack;
 use crate::network::{Embedding, Network};
@@ -207,6 +208,34 @@ where
         .await
     }
 
+    pub async fn run_preex2<U, N, M, S, H, X>(
+        &self,
+        net_stack: U,
+        netif: N,
+        mdns: M,
+        store: &SharedKvBlobStore<'_, S>,
+        handler: H,
+        user: X,
+        alloc: &BumpAlloc<'_>,
+    ) -> Result<(), Error>
+    where
+        U: NetStack,
+        N: NetifDiag + NetChangeNotif,
+        M: Mdns,
+        S: KvBlobStore,
+        H: AsyncHandler + AsyncMetadata,
+        X: UserTask,
+    {
+        self.run_with_memory(
+            PreexistingEthernet::new(net_stack, netif, mdns),
+            store,
+            handler,
+            user,
+            alloc,
+        )
+        .await
+    }
+
     /// Run the Matter stack for an Ethernet network.
     ///
     /// Parameters:
@@ -247,6 +276,50 @@ where
         select(&mut net_task, &mut persist_task).coalesce().await
     }
 
+    /// Run the Matter stack for an Ethernet network with bump allocation.
+    /// Uses the provided memory buffer for allocations instead of heap.
+    ///
+    /// Parameters:
+    /// - `ethernet` - a user-provided `Ethernet` implementation
+    /// - `persist` - a user-provided `Persist` implementation
+    /// - `handler` - a user-provided DM handler implementation
+    /// - `user` - a user-provided future that will be polled only when the netif interface is up
+    /// - `memory` - a memory buffer for internal allocations (recommended size: 16KB+)
+    pub async fn run_with_memory<N, S, H, X>(
+        &self,
+        ethernet: N,
+        store: &SharedKvBlobStore<'_, S>,
+        handler: H,
+        user: X,
+        alloc: &BumpAlloc<'_>,
+    ) -> Result<(), Error>
+    where
+        N: Ethernet,
+        S: KvBlobStore,
+        H: AsyncHandler + AsyncMetadata,
+        X: UserTask,
+    {
+        info!("Matter Stack memory: {}B", core::mem::size_of_val(self));
+
+        let persist = self.create_persist(store);
+
+        persist.load().await?;
+
+        self.matter().reset_transport()?;
+
+        if !self.is_commissioned().await? {
+            self.matter()
+                .enable_basic_commissioning(DiscoveryCapabilities::IP, 0)
+                .await?; // TODO
+        }
+
+        let mut net_task =
+            unwrap!(alloc.box_pin(self.run_ethernet_with_memory(ethernet, handler, user, alloc)));
+        let mut persist_task = unwrap!(alloc.box_pin(self.run_psm(&persist)));
+
+        select(&mut net_task, &mut persist_task).coalesce().await
+    }
+
     async fn run_ethernet<N, H, X>(&self, mut ethernet: N, handler: H, user: X) -> Result<(), Error>
     where
         N: Ethernet,
@@ -255,6 +328,25 @@ where
     {
         Ethernet::run(&mut ethernet, MatterStackEthernetTask(self, handler, user)).await
     }
+
+    async fn run_ethernet_with_memory<N, H, X>(
+        &self,
+        mut ethernet: N,
+        handler: H,
+        user: X,
+        alloc: &BumpAlloc<'_>,
+    ) -> Result<(), Error>
+    where
+        N: Ethernet,
+        H: AsyncHandler + AsyncMetadata,
+        X: UserTask,
+    {
+        Ethernet::run(
+            &mut ethernet,
+            MatterStackEthernetTaskWithMemory(self, handler, user, alloc),
+        )
+        .await
+    }
 }
 
 struct MatterStackEthernetTask<'a, E, H, X>(&'a MatterStack<'a, Eth<E>>, H, X)
@@ -263,6 +355,17 @@ where
     H: AsyncMetadata + AsyncHandler,
     X: UserTask;
 
+struct MatterStackEthernetTaskWithMemory<'a, 'b, E, H, X>(
+    &'a MatterStack<'a, Eth<E>>,
+    H,
+    X,
+    &'a BumpAlloc<'b>,
+)
+where
+    E: Embedding + 'static,
+    H: AsyncMetadata + AsyncHandler,
+    X: UserTask;
+
 impl<E, H, X> EthernetTask for MatterStackEthernetTask<'_, E, H, X>
 where
     E: Embedding + 'static,
@@ -295,3 +398,49 @@ where
             .await
     }
 }
+
+impl<E, H, X> EthernetTask for MatterStackEthernetTaskWithMemory<'_, '_, E, H, X>
+where
+    E: Embedding + 'static,
+    H: AsyncMetadata + AsyncHandler,
+    X: UserTask,
+{
+    async fn run<S, C, M>(&mut self, net_stack: S, netif: C, mut mdns: M) -> Result<(), Error>
+    where
+        S: NetStack,
+        C: NetifDiag + NetChangeNotif,
+        M: Mdns,
+    {
+        info!("Ethernet driver started with bump allocator");
+
+        // Use bump allocator instead of stack allocation for largest futures
+        let net_task = self
+            .3
+            .box_pin(self.0.run_oper_net(
+                &net_stack,
+                &netif,
+                &mut mdns,
+                core::future::pending(),
+                Option::<(NoNetwork, NoNetwork)>::None,
+            ))
+            .map_err(|_| rs_matter::error::Error::new(rs_matter::error::ErrorCode::NoMemory))?;
+
+        let handler = self.0.root_handler(&(), &true, &netif, &self.1);
+        let handler_task = self
+            .3
+            .box_pin(self.0.run_handler((&self.1, handler)))
+            .map_err(|_| rs_matter::error::Error::new(rs_matter::error::ErrorCode::NoMemory))?;
+
+        let mut user_task = pin!(self.2.run(&net_stack, &netif));
+
+        info!(
+            "Bump allocator usage: {}/{} bytes",
+            self.3.used(),
+            self.3.capacity()
+        );
+
+        select3(net_task, handler_task, &mut user_task)
+            .coalesce()
+            .await
+    }
+}

src/lib.rs

@@ -9,6 +9,8 @@
 #![warn(clippy::large_stack_frames)]
 #![warn(clippy::large_types_passed_by_value)]
 
+pub mod bump_alloc;
+
 use core::future::Future;
 use core::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV6};
 use core::pin::pin;