Implement GlobalAlloc for ESP32?

[stevefan1999-personal]

Currently I put TLSF in as a global static mut with a static memory provider from EspHeap:

#[global_allocator]
static ALLOCATOR: EspHeap = EspHeap::empty();

unsafe fn init_heap() {
    const HEAP_SIZE: usize = 32 * 1024;

    extern "C" {
        static mut _heap_start: u32;
    }

    unsafe {
        let heap_start = &_heap_start as *const _ as usize;
        ALLOCATOR.init(heap_start as *mut u8, HEAP_SIZE);
    }

    let mut TLSF: FlexTlsf<GlobalAllocAsFlexSource<EspHeap, 1024>, u16, u16, 12, 16> = FlexTlsf::new(GlobalAllocAsFlexSource(ALLOCATOR));

    unsafe {
        let mut ptr1 = TLSF.allocate(Layout::new::<u64>()).unwrap().cast::<u64>();
        let mut ptr2 = TLSF.allocate(Layout::new::<u64>()).unwrap().cast::<u64>();
        *ptr1.as_mut() = 42;
        *ptr2.as_mut() = 56;
        assert_eq!(*ptr1.as_ref(), 42);
        assert_eq!(*ptr2.as_ref(), 56);
        TLSF.deallocate(ptr1.cast(), Layout::new::<u64>().align());
        TLSF.deallocate(ptr2.cast(), Layout::new::<u64>().align());
    }

}

It seems like EspHeap does not offer a default method either so I can't take the heap from system (so EspHeap is layered upon). I may need to directly implement a flex source. What are the pointers for me to implement a FlexSource and eventually a global allocator?

[stevefan1999-personal]

I copied some code and ended up like this:

    alloc::{GlobalAlloc, Layout},
    cell::{RefCell, UnsafeCell},
    marker::PhantomData,
    mem::MaybeUninit,
    ops,
    ptr::{
        self,
        NonNull,
    },
};
use const_default::ConstDefault;
use rlsf::Tlsf;
use spin::Mutex;

pub struct GlobalTlsf {
    #[cfg(not(doc))]
    inner: Mutex<RefCell<TheTlsf>>,
}

type TheTlsf = Tlsf<'static, usize, usize, { usize::BITS as usize }, { usize::BITS as usize }>;

unsafe impl Send for GlobalTlsf {}
unsafe impl Sync for GlobalTlsf {}

impl GlobalTlsf {
    #[inline]
    pub const fn new() -> Self {
        Self {
            inner: Mutex::new(RefCell::new(ConstDefault::DEFAULT)),
        }
    }

    pub fn init(&mut self, array: &'static mut [MaybeUninit<u8>]) {
        let mut inner = self.inner.lock();
        inner.get_mut().insert_free_block(array);
    }
}


unsafe impl GlobalAlloc for GlobalTlsf {
    #[inline]
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let mut inner = self.inner.lock();
        let mut inner = inner.get_mut();
        inner.get_mut()
            .allocate(layout)
            .map(NonNull::as_ptr)
            .unwrap_or(ptr::null_mut())
    }

    #[inline]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let mut inner = self.inner.lock();
        let mut inner = inner.get_mut();
        let ptr = NonNull::new_unchecked(ptr);
        inner.deallocate(ptr, layout.align());
    }

    #[inline]
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        let mut inner = self.inner.lock();
        let mut inner = inner.get_mut();
        let ptr = NonNull::new_unchecked(ptr);
        let new_layout = Layout::from_size_align_unchecked(new_size, layout.align());
        if let Some(new_ptr) = inner.allocate(new_layout) {
            ptr::copy_nonoverlapping(ptr.as_ptr(), new_ptr.as_ptr(), layout.size().min(new_size));
            inner.deallocate(ptr, layout.align());
            new_ptr.as_ptr()
        } else {
            ptr::null_mut()
        }
    }
}

Then initial by:

static mut HEAP: [MaybeUninit<u8>; 65535] = [MaybeUninit::uninit(); 65535];

#[global_allocator]
static mut ALLOCATOR: GlobalTlsf = GlobalTlsf::new();

unsafe fn init_heap() {
    ALLOCATOR.init(&mut HEAP);
}

Not sure if this is safe but it seems like it works

[yvt]

Yes, that's the right way to implement a default allocator on unsupported targets, except that you don't need RefCell because Mutex already provides interior mutability.

Just in case though: Locking by spin::Mutex may be liable to deadlock on an embedded target if it's not done carefully. E.g., on a single-core system that supports application-level interrupt handlers, locking it in an interrupt handler will deadlock if it's already locked; on an RTOS-based system with strict priority-based scheduling, preemption while holding a lock will cause deadlock if a higher-priority task later attempts to lock it.

[stevefan1999-personal]

@yvt I'm using embassy and is suspectable to this. Any recommendation to mitigate locks? Of course we can't redesign TLSF to go lock-free...

[yvt]

It appears Embassy is based on cooperative (non-preemptive) scheduling. In this case, if you don't need to use the allocator in interrupt handlers, spin::Mutex is fine. You might want to use try_lock instead of lock to fail fast on misuse.

If you need to use it in interrupt handlers, interrupts need to be disabled while holding a lock. The main disadvantage of this approach is an increased interrupt latency. You should remove the custom implementation of GlobalAlloc::realloc to keep the worst-case latency bounded. The specific way of implementing this approach varies based on the ecosystem crates you'd like to use. E.g.,:

• With critical-section ^1 and riscv ^0.10: Wrap Tlsf with critical_section::Mutex<RefCell<_>>. Enable riscv's critical-section-single-hart feature to get critical_section::with working.
• With riscv ^0.8: Wrap Tlsf with riscv::interrupt::Mutex<RefCell<_>>.
• Or, whatever ways specific to ESP32 or Embassy that I'm not aware of.

(These "Mutex"es need RefCell because they are actually more like reentrant mutexes.)

[yvt]

Just skimmed Embassy documentation, but this looks useful: embassy_sync::blocking_mutex::Mutex

[stevefan1999-personal]

@yvt Based on your suggestion I implemented the following instead:

use core::{
    alloc::{GlobalAlloc, Layout},
    mem::MaybeUninit,
    ptr::{self, NonNull},
};

use const_default::ConstDefault;
use embassy_sync::blocking_mutex::raw::NoopRawMutex;
use embassy_sync::mutex::Mutex;
use riscv::interrupt;
use rlsf::Tlsf;

pub struct GlobalTlsf<'a> {
    #[cfg(not(doc))]
    inner: Mutex<NoopRawMutex, TheTlsf<'a>>,
}

type TheTlsf<'a> = Tlsf<'a, usize, usize, { usize::BITS as usize }, { usize::BITS as usize }>;

unsafe impl Send for GlobalTlsf<'_> {}

unsafe impl Sync for GlobalTlsf<'_> {}

impl<'a> GlobalTlsf<'a> {
    #[inline]
    pub const fn new() -> Self {
        Self {
            inner: Mutex::new(ConstDefault::DEFAULT),
        }
    }

    pub fn init(&mut self, array: &'a mut [MaybeUninit<u8>]) {
        interrupt::free(|| {
            let mut inner = self.inner.try_lock().expect("deadlock detected");
            inner.insert_free_block(array);
        })
    }
}

unsafe impl GlobalAlloc for GlobalTlsf<'_> {
    #[inline]
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        interrupt::free(|| {
            let mut inner = self.inner.try_lock().expect("deadlock detected");
            inner
                .allocate(layout)
                .map(NonNull::as_ptr)
                .unwrap_or(ptr::null_mut())
        })
    }

    #[inline]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        interrupt::free(|| {
            let mut inner = self.inner.try_lock().expect("deadlock detected");
            let ptr = NonNull::new_unchecked(ptr);
            inner.deallocate(ptr, layout.align());
        })
    }
}

Not sure if it should be added to main crate or just kept as an example for other to implement on their own