feat(esp-alloc): Add heap usage stats

esp-alloc/CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- `esp_alloc::HEAP.stats()` can now be used to get heap usage informations (#2137)
+
 ### Changed
 
 - a global allocator is created in esp-alloc, now you need to add individual memory regions (up to 3) to the allocator (#2099)

esp-alloc/Cargo.toml

@@ -23,10 +23,24 @@ default-target = "riscv32imc-unknown-none-elf"
 features       = ["nightly"]
 
 [dependencies]
+defmt                 = { version = "0.3.8", optional = true }
+cfg-if                = "1.0.0"
 critical-section      = "1.1.3"
 enumset               = "1.1.5"
 linked_list_allocator = { version = "0.10.5", default-features = false, features = ["const_mut_refs"] }
+document-features = "0.2.10"
 
 [features]
 default = []
 nightly = []
+
+## Implement `defmt::Format` on certain types.
+defmt = ["dep:defmt"]
+
+## Enable this feature if you want to keep stats about the internal heap usage such as:
+## - Max memory usage since initialization of the heap
+## - Total allocated memory since initialization of the heap
+## - Total freed memory since initialization of the heap
+##
+## ⚠️ Note: Enabling this feature will require extra computation every time alloc/dealloc is called.
+internal-heap-stats = []

esp-alloc/src/lib.rs

@@ -51,7 +51,28 @@
 //! ```rust
 //! let large_buffer: Vec<u8, _> = Vec::with_capacity_in(1048576, &PSRAM_ALLOCATOR);
 //! ```
-
+//!
+//! You can also gets stats about the heap usage at anytime with:
+//! ```rust
+//! let stats: HeapStats = esp_alloc::HEAP.stats();
+//! // HeapStats implements the Display trait, so you can pretty print the heap stats.
+//! println!("{}", stats);
+//! ```
+//!
+//! ```txt
+//! HEAP INFO
+//! Size: 2097152
+//! Current usage: 512028
+//! Max usage: 512028
+//! Total freed: 0
+//! Total allocated: 512028
+//! Memory Layout:
+//! External | ████████░░░░░░░░░░░░░░░░░░░░░░░░░░░ | Used: 512028 / Total: 2097152 (Free: 1585124)
+//! Unused   | ░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░ |
+//! Unused   | ░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░ |
+//! ```
+//! ## Feature Flags
+#![doc = document_features::document_features!()]
 #![no_std]
 #![cfg_attr(feature = "nightly", feature(allocator_api))]
 #![doc(html_logo_url = "https://avatars.githubusercontent.com/u/46717278")]
@@ -63,6 +84,7 @@ use core::alloc::{AllocError, Allocator};
 use core::{
     alloc::{GlobalAlloc, Layout},
     cell::RefCell,
+    fmt::Display,
     ptr::{self, NonNull},
 };
 
@@ -76,7 +98,9 @@ pub static HEAP: EspHeap = EspHeap::empty();
 
 const NON_REGION: Option<HeapRegion> = None;
 
-#[derive(EnumSetType)]
+const BAR_WIDTH: usize = 35;
+
+#[derive(EnumSetType, Debug)]
 /// Describes the properties of a memory region
 pub enum MemoryCapability {
     /// Memory must be internal; specifically it should not disappear when
@@ -86,6 +110,61 @@ pub enum MemoryCapability {
     External,
 }
 
+/// Stats for a heap region
+#[derive(Debug)]
+pub struct RegionStats {
+    /// Total usable size of the heap region in bytes.
+    size: usize,
+
+    /// Currently used size of the heap region in bytes.
+    used: usize,
+
+    /// Free size of the heap region in bytes.
+    free: usize,
+
+    /// Capabilities of the memory region.
+    capabilities: EnumSet<MemoryCapability>,
+}
+
+impl Display for RegionStats {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        let used_blocks = BAR_WIDTH * self.used / self.size;
+        let free_blocks = BAR_WIDTH - used_blocks;
+
+        // Display Memory type
+        if self.capabilities.contains(MemoryCapability::Internal) {
+            write!(f, "Internal")?;
+        } else if self.capabilities.contains(MemoryCapability::External) {
+            write!(f, "External")?;
+        } else {
+            write!(f, "Unknown")?;
+        }
+
+        write!(f, " | ")?;
+
+        // Display usage of the memory region using a bar graph.
+        for _ in 0..used_blocks {
+            write!(f, "█")?;
+        }
+        for _ in 0..free_blocks {
+            write!(f, "░")?;
+        }
+
+        write!(
+            f,
+            " | Used: {} / Total: {} (Free: {})",
+            self.used, self.size, self.free
+        )
+    }
+}
+
+#[cfg(feature = "defmt")]
+impl defmt::Format for RegionStats {
+    fn format(&self, fmt: defmt::Formatter) {
+        defmt::write!(fmt, "{}", defmt::Display2Format(self))
+    }
+}
+
 /// A memory region to be used as heap memory
 pub struct HeapRegion {
     heap: Heap,
@@ -112,6 +191,88 @@ impl HeapRegion {
 
         Self { heap, capabilities }
     }
+
+    /// Return stats for the current memory region
+    pub fn stats(&self) -> RegionStats {
+        RegionStats {
+            size: self.heap.size(),
+            used: self.heap.used(),
+            free: self.heap.free(),
+            capabilities: self.capabilities,
+        }
+    }
+}
+
+/// Stats for a heap allocator
+///
+/// Enable the "internal-heap-stats" feature if you want collect additional heap
+/// informations at the cost of extra cpu time during every alloc/dealloc.
+#[derive(Debug)]
+pub struct HeapStats {
+    /// Granular stats for all the configured memory regions.
+    region_stats: [Option<RegionStats>; 3],
+
+    /// Total size of all combined heap regions in bytes.
+    size: usize,
+
+    /// Current usage of the heap across all configured regions in bytes.
+    current_usage: usize,
+
+    /// Estimation of the max used heap in bytes.
+    #[cfg(feature = "internal-heap-stats")]
+    max_usage: usize,
+
+    /// Estimation of the total allocated bytes since initialization.
+    #[cfg(feature = "internal-heap-stats")]
+    total_allocated: usize,
+
+    /// Estimation of the total freed bytes since initialization.
+    #[cfg(feature = "internal-heap-stats")]
+    total_freed: usize,
+}
+
+impl Display for HeapStats {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        writeln!(f, "HEAP INFO")?;
+        writeln!(f, "Size: {}", self.size)?;
+        writeln!(f, "Current usage: {}", self.current_usage)?;
+        #[cfg(feature = "internal-heap-stats")]
+        {
+            writeln!(f, "Max usage: {}", self.max_usage)?;
+            writeln!(f, "Total freed: {}", self.total_freed)?;
+            writeln!(f, "Total allocated: {}", self.total_allocated)?;
+        }
+        writeln!(f, "Memory Layout: ")?;
+        for region in self.region_stats.iter() {
+            if let Some(region) = region.as_ref() {
+                region.fmt(f)?;
+                writeln!(f)?;
+            } else {
+                // Display unused memory regions
+                write!(f, "Unused   | ")?;
+                for _ in 0..BAR_WIDTH {
+                    write!(f, "░")?;
+                }
+                writeln!(f, " |")?;
+            }
+        }
+        Ok(())
+    }
+}
+
+#[cfg(feature = "defmt")]
+impl defmt::Format for HeapStats {
+    fn format(&self, fmt: defmt::Formatter) {
+        defmt::write!(fmt, "{}", defmt::Display2Format(self))
+    }
+}
+
+/// Internal stats to keep track across multiple regions.
+#[cfg(feature = "internal-heap-stats")]
+struct InternalHeapStats {
+    max_usage: usize,
+    total_allocated: usize,
+    total_freed: usize,
 }
 
 /// A memory allocator
@@ -120,13 +281,21 @@ impl HeapRegion {
 /// memory in regions satisfying specific needs.
 pub struct EspHeap {
     heap: Mutex<RefCell<[Option<HeapRegion>; 3]>>,
+    #[cfg(feature = "internal-heap-stats")]
+    internal_heap_stats: Mutex<RefCell<InternalHeapStats>>,
 }
 
 impl EspHeap {
     /// Crate a new UNINITIALIZED heap allocator
     pub const fn empty() -> Self {
         EspHeap {
             heap: Mutex::new(RefCell::new([NON_REGION; 3])),
+            #[cfg(feature = "internal-heap-stats")]
+            internal_heap_stats: Mutex::new(RefCell::new(InternalHeapStats {
+                max_usage: 0,
+                total_allocated: 0,
+                total_freed: 0,
+            })),
         }
     }
 
@@ -189,6 +358,51 @@ impl EspHeap {
         })
     }
 
+    /// Return usage stats for the [Heap].
+    ///
+    /// Note:
+    /// [HeapStats] directly implements [Display], so this function can be
+    /// called from within `println!()` to pretty-print the usage of the
+    /// heap.
+    pub fn stats(&self) -> HeapStats {
+        const EMPTY_REGION_STAT: Option<RegionStats> = None;
+        let mut region_stats: [Option<RegionStats>; 3] = [EMPTY_REGION_STAT; 3];
+
+        critical_section::with(|cs| {
+            let mut used = 0;
+            let mut free = 0;
+            let regions = self.heap.borrow_ref(cs);
+            for (id, region) in regions.iter().enumerate() {
+                if let Some(region) = region.as_ref() {
+                    let stats = region.stats();
+                    free += stats.free;
+                    used += stats.used;
+                    region_stats[id] = Some(region.stats());
+                }
+            }
+
+            cfg_if::cfg_if! {
+                if #[cfg(feature = "internal-heap-stats")] {
+                    let internal_heap_stats = self.internal_heap_stats.borrow_ref(cs);
+                    HeapStats {
+                        region_stats,
+                        size: free + used,
+                        current_usage: used,
+                        max_usage: internal_heap_stats.max_usage,
+                        total_allocated: internal_heap_stats.total_allocated,
+                        total_freed: internal_heap_stats.total_freed,
+                    }
+                } else {
+                    HeapStats {
+                        region_stats,
+                        size: free + used,
+                        current_usage: used,
+                    }
+                }
+            }
+        })
+    }
+
     /// Returns an estimate of the amount of bytes available.
     pub fn free(&self) -> usize {
         self.free_caps(EnumSet::empty())
@@ -232,6 +446,8 @@ impl EspHeap {
         layout: Layout,
     ) -> *mut u8 {
         critical_section::with(|cs| {
+            #[cfg(feature = "internal-heap-stats")]
+            let before = self.used();
             let mut regions = self.heap.borrow_ref_mut(cs);
             let mut iter = (*regions).iter_mut().filter(|region| {
                 if region.is_some() {
@@ -256,7 +472,22 @@ impl EspHeap {
                 }
             };
 
-            res.map_or(ptr::null_mut(), |allocation| allocation.as_ptr())
+            res.map_or(ptr::null_mut(), |allocation| {
+                #[cfg(feature = "internal-heap-stats")]
+                {
+                    let mut internal_heap_stats = self.internal_heap_stats.borrow_ref_mut(cs);
+                    drop(regions);
+                    // We need to call used because [linked_list_allocator::Heap] does internal size
+                    // alignment so we cannot use the size provided by the layout.
+                    let used = self.used();
+
+                    internal_heap_stats.total_allocated += used - before;
+                    internal_heap_stats.max_usage =
+                        core::cmp::max(internal_heap_stats.max_usage, used);
+                }
+
+                allocation.as_ptr()
+            })
         })
     }
 }
@@ -272,6 +503,8 @@ unsafe impl GlobalAlloc for EspHeap {
         }
 
         critical_section::with(|cs| {
+            #[cfg(feature = "internal-heap-stats")]
+            let before = self.used();
             let mut regions = self.heap.borrow_ref_mut(cs);
             let mut iter = (*regions).iter_mut();
 
@@ -280,6 +513,15 @@ unsafe impl GlobalAlloc for EspHeap {
                     region.heap.deallocate(NonNull::new_unchecked(ptr), layout);
                 }
             }
+
+            #[cfg(feature = "internal-heap-stats")]
+            {
+                let mut internal_heap_stats = self.internal_heap_stats.borrow_ref_mut(cs);
+                drop(regions);
+                // We need to call used because [linked_list_allocator::Heap] does internal size
+                // alignment so we cannot use the size provided by the layout.
+                internal_heap_stats.total_freed += before - self.used();
+            }
         })
     }
 }

examples/src/bin/psram_quad.rs

@@ -3,7 +3,7 @@
 //! You need an ESP32, ESP32-S2 or ESP32-S3 with at least 2 MB of PSRAM memory.
 
 //% CHIPS: esp32 esp32s2 esp32s3
-//% FEATURES: psram-2m
+//% FEATURES: psram-2m esp-alloc/internal-heap-stats
 
 #![no_std]
 #![no_main]
@@ -51,6 +51,8 @@ fn main() -> ! {
     let string = String::from("A string allocated in PSRAM");
     println!("'{}' allocated at {:p}", &string, string.as_ptr());
 
+    println!("{}", esp_alloc::HEAP.stats());
+
     println!("done");
 
     loop {}