Add heap usage stats with defmt (#2619)

* feat(esp-alloc): Add heap usage stats and provide `esp_alloc::get_info!()` macro

* refactor(esp-alloc): Feature gate internal memory usage that requires extra computation.

- Introduce the `internal-heap-stats` feature for `esp-alloc`.
- Add `esp_alloc::get_info!()` to `psram_quad` example to show usage and
ensure coverage of the feature in tests.

* refactor(esp-alloc): Remove `get_info!()` macro in favour of documenting `HEAP.stats()`

* Implement defmt::Format for HeapStats and RegionStats

* rustfmt

* show usage percent + move bar drawing logic to separate functions

* update doc comments

* Fixed a typo in qa-test/src/bin/psram_quad.rs

Co-authored-by: Scott Mabin <scott@mabez.dev>

* minor improvements to write bar functions

* Aligned the indentation in Cargo.toml

Co-authored-by: Kirill Mikhailov <62840029+playfulFence@users.noreply.github.com>

* Fixed a typo in docs

Co-authored-by: Kirill Mikhailov <62840029+playfulFence@users.noreply.github.com>

* Nitpicking x2

Co-authored-by: Kirill Mikhailov <62840029+playfulFence@users.noreply.github.com>

* Surround a function call with backticks

Co-authored-by: Kirill Mikhailov <62840029+playfulFence@users.noreply.github.com>

* rustfmt

---------

Co-authored-by: Anthony Grondin <104731965+AnthonyGrondin@users.noreply.github.com>
Co-authored-by: Scott Mabin <scott@mabez.dev>
Co-authored-by: Kirill Mikhailov <62840029+playfulFence@users.noreply.github.com>

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
 
 ### Fixed

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 get stats about the heap usage at anytime with:
+//! ```rust
+//! let stats: HeapStats = esp_alloc::HEAP.stats();
+//! // HeapStats implements the Display and defmt::Format traits, so you can pretty-print the heap stats.
+//! println!("{}", stats);
+//! ```
+//!
+//! ```txt
+//! HEAP INFO
+//! Size: 131068
+//! Current usage: 46148
+//! Max usage: 46148
+//! Total freed: 0
+//! Total allocated: 46148
+//! Memory Layout:
+//! Internal | ████████████░░░░░░░░░░░░░░░░░░░░░░░ | Used: 35% (Used 46148 of 131068, free: 84920)
+//! 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,22 @@ pub static HEAP: EspHeap = EspHeap::empty();
 
 const NON_REGION: Option<HeapRegion> = None;
 
-#[derive(EnumSetType)]
+const BAR_WIDTH: usize = 35;
+
+fn write_bar(f: &mut core::fmt::Formatter<'_>, usage_percent: usize) -> core::fmt::Result {
+    let used_blocks = BAR_WIDTH * usage_percent / 100;
+    (0..used_blocks).try_for_each(|_| write!(f, "█"))?;
+    (used_blocks..BAR_WIDTH).try_for_each(|_| write!(f, "░"))
+}
+
+#[cfg(feature = "defmt")]
+fn write_bar_defmt(fmt: defmt::Formatter, usage_percent: usize) {
+    let used_blocks = BAR_WIDTH * usage_percent / 100;
+    (0..used_blocks).for_each(|_| defmt::write!(fmt, "█"));
+    (used_blocks..BAR_WIDTH).for_each(|_| defmt::write!(fmt, "░"));
+}
+
+#[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 +123,75 @@ 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 usage_percent = self.used * 100 / self.size;
+
+        // 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, " | ")?;
+
+        write_bar(f, usage_percent)?;
+
+        write!(
+            f,
+            " | Used: {}% (Used {} of {}, free: {})",
+            usage_percent, self.used, self.size, self.free
+        )
+    }
+}
+
+#[cfg(feature = "defmt")]
+impl defmt::Format for RegionStats {
+    fn format(&self, fmt: defmt::Formatter) {
+        let usage_percent = self.used * 100 / self.size;
+
+        if self.capabilities.contains(MemoryCapability::Internal) {
+            defmt::write!(fmt, "Internal");
+        } else if self.capabilities.contains(MemoryCapability::External) {
+            defmt::write!(fmt, "External");
+        } else {
+            defmt::write!(fmt, "Unknown");
+        }
+
+        defmt::write!(fmt, " | ");
+
+        write_bar_defmt(fmt, usage_percent);
+
+        defmt::write!(
+            fmt,
+            " | Used: {}% (Used {} of {}, free: {})",
+            usage_percent,
+            self.used,
+            self.size,
+            self.free
+        );
+    }
+}
+
 /// A memory region to be used as heap memory
 pub struct HeapRegion {
     heap: Heap,
@@ -112,6 +218,104 @@ 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   | ")?;
+                write_bar(f, 0)?;
+                writeln!(f, " |")?;
+            }
+        }
+        Ok(())
+    }
+}
+
+#[cfg(feature = "defmt")]
+impl defmt::Format for HeapStats {
+    fn format(&self, fmt: defmt::Formatter) {
+        defmt::write!(fmt, "HEAP INFO\n");
+        defmt::write!(fmt, "Size: {}\n", self.size);
+        defmt::write!(fmt, "Current usage: {}\n", self.current_usage);
+        #[cfg(feature = "internal-heap-stats")]
+        {
+            defmt::write!(fmt, "Max usage: {}\n", self.max_usage);
+            defmt::write!(fmt, "Total freed: {}\n", self.total_freed);
+            defmt::write!(fmt, "Total allocated: {}\n", self.total_allocated);
+        }
+        defmt::write!(fmt, "Memory Layout:\n");
+        for region in self.region_stats.iter() {
+            if let Some(region) = region.as_ref() {
+                defmt::write!(fmt, "{}\n", region);
+            } else {
+                defmt::write!(fmt, "Unused   | ");
+                write_bar_defmt(fmt, 0);
+                defmt::write!(fmt, " |\n");
+            }
+        }
+    }
+}
+
+/// 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,6 +324,8 @@ 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 {
@@ -127,6 +333,12 @@ impl EspHeap {
     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 +401,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 +489,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 +515,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 +546,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 +556,16 @@ 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();
+            }
         })
     }
 }

qa-test/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: esp-hal/quad-psram
+//% FEATURES: esp-hal/quad-psram esp-alloc/internal-heap-stats
 
 #![no_std]
 #![no_main]
@@ -52,6 +52,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 {}