Optimize async future memory usage with bump allocator for embedded systems

.gitignore

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

src/bump_alloc.rs

@@ -0,0 +1,130 @@
+/// A simple bump allocator for fixed-size memory chunks
+/// 
+/// This allocator is designed for async futures that have a known maximum lifetime.
+/// All allocations are freed when the allocator is dropped.
+use core::alloc::Layout;
+use core::cell::Cell;
+use core::mem::MaybeUninit;
+use core::pin::Pin;
+use core::ptr::NonNull;
+
+/// A bump allocator that uses a provided memory chunk
+pub struct BumpAllocator<'a> {
+    memory: &'a mut [MaybeUninit<u8>],
+    offset: Cell<usize>,
+}
+
+impl<'a> BumpAllocator<'a> {
+    /// Create a new bump allocator with the provided memory chunk
+    pub fn new(memory: &'a mut [MaybeUninit<u8>]) -> Self {
+        Self {
+            memory,
+            offset: Cell::new(0),
+        }
+    }
+
+    /// Allocate memory for a future and pin it
+    pub fn alloc_pin<F>(&mut self, future: F) -> Result<Pin<BumpBox<'_, F>>, AllocError>
+    where
+        F: core::future::Future,
+    {
+        let layout = Layout::new::<F>();
+        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 F;
+            ptr.write(future);
+            Ok(Pin::new_unchecked(BumpBox {
+                ptr: NonNull::new_unchecked(ptr),
+                _allocator: core::marker::PhantomData,
+            }))
+        }
+    }
+
+    fn alloc_raw(&mut self, layout: Layout) -> Result<NonNull<u8>, AllocError> {
+        let size = layout.size();
+        let align = layout.align();
+
+        let current_offset = self.offset.get();
+
+        // Align the offset
+        let aligned_offset = (current_offset + align - 1) & !(align - 1);
+        let new_offset = aligned_offset + size;
+
+        if new_offset > self.memory.len() {
+            return Err(AllocError);
+        }
+
+        self.offset.set(new_offset);
+
+        // Safety: We checked bounds and alignment
+        let ptr = unsafe {
+            self.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.offset.get()
+    }
+
+    /// Get the total capacity
+    pub fn capacity(&self) -> usize {
+        self.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

@@ -1,7 +1,12 @@
 use core::pin::pin;
+use core::mem::MaybeUninit;
+extern crate alloc;
+use alloc::boxed::Box;
 
 use embassy_futures::select::{select, select3};
 
+use crate::bump_alloc::BumpAllocator;
+
 use rs_matter::dm::clusters::gen_comm::CommPolicy;
 use rs_matter::dm::clusters::gen_diag::{GenDiag, NetifDiag};
 use rs_matter::dm::clusters::net_comm::NetworkType;
@@ -247,6 +252,49 @@ 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,
+        memory: &mut [MaybeUninit<u8>],
+    ) -> 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 = pin!(self.run_ethernet_with_memory(ethernet, handler, user, memory));
+        let mut persist_task = 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 +303,21 @@ 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,
+        memory: &mut [MaybeUninit<u8>]
+    ) -> Result<(), Error>
+    where
+        N: Ethernet,
+        H: AsyncHandler + AsyncMetadata,
+        X: UserTask,
+    {
+        Ethernet::run(&mut ethernet, MatterStackEthernetTaskWithMemory(self, handler, user, memory)).await
+    }
 }
 
 struct MatterStackEthernetTask<'a, E, H, X>(&'a MatterStack<'a, Eth<E>>, H, X)
@@ -263,6 +326,12 @@ where
     H: AsyncMetadata + AsyncHandler,
     X: UserTask;
 
+struct MatterStackEthernetTaskWithMemory<'a, E, H, X>(&'a MatterStack<'a, Eth<E>>, H, X, &'a mut [MaybeUninit<u8>])
+where
+    E: Embedding + 'static,
+    H: AsyncMetadata + AsyncHandler,
+    X: UserTask;
+
 impl<E, H, X> EthernetTask for MatterStackEthernetTask<'_, E, H, X>
 where
     E: Embedding + 'static,
@@ -277,7 +346,8 @@ where
     {
         info!("Ethernet driver started");
 
-        let mut net_task = pin!(self.0.run_oper_net(
+        // Box the largest futures to reduce stack frame size
+        let net_task = Box::pin(self.0.run_oper_net(
             &net_stack,
             &netif,
             &mut mdns,
@@ -286,11 +356,51 @@ where
         ));
 
         let handler = self.0.root_handler(&(), &true, &netif, &self.1);
-        let mut handler_task = pin!(self.0.run_handler((&self.1, handler)));
+        let handler_task = Box::pin(self.0.run_handler((&self.1, handler)));
 
         let mut user_task = pin!(self.2.run(&net_stack, &netif));
 
-        select3(&mut net_task, &mut handler_task, &mut user_task)
+        select3(net_task, handler_task, &mut user_task)
+            .coalesce()
+            .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");
+
+        // Create bump allocator from provided memory
+        let mut allocator = BumpAllocator::new(self.3);
+
+        // Use bump allocator instead of Box::pin for largest futures
+        let net_task = allocator.alloc_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, false))?;
+
+        let handler = self.0.root_handler(&(), &true, &netif, &self.1);
+        let handler_task = allocator.alloc_pin(self.0.run_handler((&self.1, handler)))
+            .map_err(|_| rs_matter::error::Error::new(rs_matter::error::ErrorCode::NoMemory, false))?;
+
+        let mut user_task = pin!(self.2.run(&net_stack, &netif));
+
+        info!("Bump allocator usage: {}/{} bytes", allocator.used(), allocator.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;

src/wireless/thread.rs

@@ -1,4 +1,6 @@
 use core::pin::pin;
+extern crate alloc;
+use alloc::boxed::Box;
 
 use embassy_futures::select::{select, select3, select4};
 use embassy_sync::blocking_mutex::raw::RawMutex;
@@ -415,12 +417,13 @@ where
             NoopWirelessNetCtl::new(NetworkType::Thread),
         );
 
-        let mut btp_task = pin!(self.0.run_btp(peripheral));
+        // Box the largest futures to reduce stack frame size
+        let btp_task = Box::pin(self.0.run_btp(peripheral));
 
         let handler = self.0.root_handler(&(), &(), &net_ctl, &false, &self.1);
-        let mut handler_task = pin!(self.0.run_handler((&self.1, handler)));
+        let handler_task = Box::pin(self.0.run_handler((&self.1, handler)));
 
-        select(&mut btp_task, &mut handler_task).coalesce().await
+        select(btp_task, handler_task).coalesce().await
     }
 }
 
@@ -452,7 +455,8 @@ where
 
         let stack = &mut self.0;
 
-        let mut net_task = pin!(stack.run_oper_net(
+        // Box the largest futures to reduce stack frame size
+        let net_task = Box::pin(stack.run_oper_net(
             &net_stack,
             &netif,
             &mut mdns,
@@ -466,14 +470,14 @@ where
         let handler = self
             .0
             .root_handler(&(), &netif, &net_ctl_s, &false, &self.1);
-        let mut handler_task = pin!(self.0.run_handler((&self.1, handler)));
+        let handler_task = Box::pin(self.0.run_handler((&self.1, handler)));
 
         let mut user_task = pin!(self.2.run(&net_stack, &netif));
 
         select4(
-            &mut net_task,
+            net_task,
             &mut mgr_task,
-            &mut handler_task,
+            handler_task,
             &mut user_task,
         )
         .coalesce()
@@ -507,17 +511,17 @@ where
 
         let stack = &mut self.0;
 
-        let mut net_task =
-            pin!(stack.run_net_coex(&net_stack, &netif, &net_ctl, &mut mdns, &mut gatt));
+        // Box the largest futures to reduce stack frame size
+        let net_task = Box::pin(stack.run_net_coex(&net_stack, &netif, &net_ctl, &mut mdns, &mut gatt));
 
         let net_ctl_s = NetCtlWithStatusImpl::new(&self.0.network.net_state, &net_ctl);
 
         let handler = self.0.root_handler(&(), &netif, &net_ctl_s, &true, &self.1);
-        let mut handler_task = pin!(self.0.run_handler((&self.1, handler)));
+        let handler_task = Box::pin(self.0.run_handler((&self.1, handler)));
 
         let mut user_task = pin!(self.2.run(&net_stack, &netif));
 
-        select3(&mut net_task, &mut handler_task, &mut user_task)
+        select3(net_task, handler_task, &mut user_task)
             .coalesce()
             .await
     }