heap: Make weak declaration for the alloc and free callbacks

- Fix "test get allocated size"
- Add tests for the free / alloc hooks
- Call alloc function hook on malloc/realloc/calloc base functions
- Add caps parameter to the allocation hook function

components/heap/heap_caps.c

@@ -15,6 +15,12 @@
 #include "heap_private.h"
 #include "esp_system.h"
 
+#define CALL_HOOK(hook, ...) {      \
+    if (hook != NULL) {             \
+        hook(__VA_ARGS__);          \
+    }                               \
+}
+
 /* Forward declaration for base function, put in IRAM.
  * These functions don't check for errors after trying to allocate memory. */
 static void *heap_caps_realloc_base( void *ptr, size_t size, uint32_t caps );
@@ -31,8 +37,6 @@ possible. This should optimize the amount of RAM accessible to the code without
 */
 
 static esp_alloc_failed_hook_t alloc_failed_callback;
-static esp_heap_trace_alloc_hook_t trace_alloc_callback;
-static esp_heap_trace_free_hook_t trace_free_callback;
 
 #ifdef CONFIG_HEAP_ABORT_WHEN_ALLOCATION_FAILS
 IRAM_ATTR static void hex_to_str(char buf[8], uint32_t n)
@@ -100,28 +104,6 @@ esp_err_t heap_caps_register_failed_alloc_callback(esp_alloc_failed_hook_t callb
     return ESP_OK;
 }
 
-esp_err_t heap_caps_register_trace_alloc_callback(esp_heap_trace_alloc_hook_t callback)
-{
-    if (callback == NULL) {
-        return ESP_ERR_INVALID_ARG;
-    }
-
-    trace_alloc_callback = callback;
-
-    return ESP_OK;
-}
-
-esp_err_t heap_caps_register_trace_free_callback(esp_heap_trace_free_hook_t callback)
-{
-    if (callback == NULL) {
-        return ESP_ERR_INVALID_ARG;
-    }
-
-    trace_free_callback = callback;
-
-    return ESP_OK;
-}
-
 bool heap_caps_match(const heap_t *heap, uint32_t caps)
 {
     return heap->heap != NULL && ((get_all_caps(heap) & caps) == caps);
@@ -136,11 +118,7 @@ IRAM_ATTR static void *heap_caps_malloc_base( size_t size, uint32_t caps)
 {
     void *ret = NULL;
 
-    if (size == 0) {
-        return NULL;
-    }
-
-    if (size > HEAP_SIZE_MAX) {
+    if (size == 0 || size > HEAP_SIZE_MAX ) {
         // Avoids int overflow when adding small numbers to size, or
         // calculating 'end' from start+size, by limiting 'size' to the possible range
         return NULL;
@@ -185,12 +163,15 @@ IRAM_ATTR static void *heap_caps_malloc_base( size_t size, uint32_t caps)
                         ret = multi_heap_malloc(heap->heap, size + 4);  // int overflow checked above
 
                         if (ret != NULL) {
-                            return dram_alloc_to_iram_addr(ret, size + 4);  // int overflow checked above
+                            uint32_t *iptr = dram_alloc_to_iram_addr(ret, size + 4);  // int overflow checked above
+                            CALL_HOOK(esp_heap_trace_alloc_hook, iptr, size, caps);
+                            return iptr;
                         }
                     } else {
                         //Just try to alloc, nothing special.
                         ret = multi_heap_malloc(heap->heap, size);
                         if (ret != NULL) {
+                            CALL_HOOK(esp_heap_trace_alloc_hook, ret, size, caps);
                             return ret;
                         }
                     }
@@ -211,9 +192,6 @@ IRAM_ATTR void *heap_caps_malloc( size_t size, uint32_t caps){
 
     void* ptr = heap_caps_malloc_base(size, caps);
 
-    if (trace_alloc_callback) {
-        trace_alloc_callback(ptr, size);
-    }
 
     if (!ptr && size > 0){
         heap_caps_alloc_failed(size, caps, __func__);
@@ -255,11 +233,6 @@ IRAM_ATTR void *heap_caps_malloc_default( size_t size )
             r=heap_caps_malloc_base( size, MALLOC_CAP_DEFAULT );
         }
 
-        // trace allocation
-        if (trace_alloc_callback) {
-            trace_alloc_callback(r, size);
-        }
-
         // allocation failure?
         if (r==NULL && size > 0){
             heap_caps_alloc_failed(size, MALLOC_CAP_DEFAULT, __func__);
@@ -294,11 +267,6 @@ IRAM_ATTR void *heap_caps_realloc_default( void *ptr, size_t size )
             r=heap_caps_realloc_base( ptr, size, MALLOC_CAP_DEFAULT);
         }
 
-        // trace allocation
-        if (trace_alloc_callback) {
-            trace_alloc_callback(r, size);
-        }
-
         // allocation failure?
         if (r==NULL && size>0){
             heap_caps_alloc_failed(size, MALLOC_CAP_DEFAULT, __func__);
@@ -323,9 +291,7 @@ IRAM_ATTR void *heap_caps_malloc_prefer( size_t size, size_t num, ... )
             break;
         }
     }
-    if (trace_alloc_callback) {
-        trace_alloc_callback(r, size);
-    }
+
     if (r == NULL && size > 0){
         heap_caps_alloc_failed(size, caps, __func__);
     }
@@ -349,9 +315,7 @@ IRAM_ATTR void *heap_caps_realloc_prefer( void *ptr, size_t size, size_t num, ..
             break;
         }
     }
-    if (trace_alloc_callback) {
-        trace_alloc_callback(r, size);
-    }
+
     if (r == NULL && size > 0){
         heap_caps_alloc_failed(size, caps, __func__);
     }
@@ -375,9 +339,7 @@ IRAM_ATTR void *heap_caps_calloc_prefer( size_t n, size_t size, size_t num, ...
             break;
         }
     }
-    if (trace_alloc_callback) {
-        trace_alloc_callback(r, size);
-    }
+
     if (r == NULL && size > 0){
         heap_caps_alloc_failed(size, caps, __func__);
     }
@@ -421,9 +383,7 @@ IRAM_ATTR void heap_caps_free( void *ptr)
     assert(heap != NULL && "free() target pointer is outside heap areas");
     multi_heap_free(heap->heap, ptr);
 
-    if (trace_free_callback) {
-        trace_free_callback(ptr);
-    }
+    CALL_HOOK(esp_heap_trace_free_hook, ptr);
 }
 
 /*
@@ -477,6 +437,7 @@ IRAM_ATTR static void *heap_caps_realloc_base( void *ptr, size_t size, uint32_t
         // (which will resize the block if it can)
         void *r = multi_heap_realloc(heap->heap, ptr, size);
         if (r != NULL) {
+            CALL_HOOK(esp_heap_trace_alloc_hook, r, size, caps);
             return r;
         }
     }
@@ -508,9 +469,7 @@ IRAM_ATTR void *heap_caps_realloc( void *ptr, size_t size, uint32_t caps)
 {
     ptr = heap_caps_realloc_base(ptr, size, caps);
 
-    if (trace_alloc_callback) {
-        trace_alloc_callback(ptr, size);
-    }
+
     if (ptr == NULL && size > 0){
         heap_caps_alloc_failed(size, caps, __func__);
     }
@@ -542,9 +501,7 @@ IRAM_ATTR void *heap_caps_calloc( size_t n, size_t size, uint32_t caps)
 {
     void* ptr = heap_caps_calloc_base(n, size, caps);
 
-    if (trace_alloc_callback) {
-        trace_alloc_callback(ptr, size);
-    }
+
     if (!ptr && size > 0){
         heap_caps_alloc_failed(size, caps, __func__);
     }
@@ -696,7 +653,7 @@ size_t heap_caps_get_allocated_size( void *ptr )
 
 IRAM_ATTR void *heap_caps_aligned_alloc(size_t alignment, size_t size, uint32_t caps)
 {
-    void *ptr = NULL;
+    void *ret = NULL;
 
     if(!alignment) {
         return NULL;
@@ -731,24 +688,20 @@ IRAM_ATTR void *heap_caps_aligned_alloc(size_t alignment, size_t size, uint32_t
                 //doesn't cover, see if they're available in other prios.
                 if ((get_all_caps(heap) & caps) == caps) {
                     //Just try to alloc, nothing special.
-                    ptr = multi_heap_aligned_alloc(heap->heap, size, alignment);
-                    if (ptr != NULL) {
-                        break;
+                    ret = multi_heap_aligned_alloc(heap->heap, size, alignment);
+                    if (ret != NULL) {
+                        CALL_HOOK(esp_heap_trace_alloc_hook, ret, size, caps);
+                        return ret;
                     }
                 }
             }
         }
     }
 
-    if (trace_alloc_callback) {
-        trace_alloc_callback(ptr, size);
-    }
-
-    if (size > 0 && ptr != NULL) {
-        heap_caps_alloc_failed(size, caps, __func__);
-    }
+    heap_caps_alloc_failed(size, caps, __func__);
 
-    return ptr;
+    //Nothing usable found.
+    return NULL;
 }
 
 IRAM_ATTR void heap_caps_aligned_free(void *ptr)

components/heap/include/esp_heap_caps.h

@@ -11,6 +11,7 @@
 #include "multi_heap.h"
 #include <sdkconfig.h>
 #include "esp_err.h"
+#include "esp_attr.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -57,17 +58,11 @@ esp_err_t heap_caps_register_failed_alloc_callback(esp_alloc_failed_hook_t callb
  * @brief callback called after every allocation
  * @param ptr the allocated memory
  * @param size in bytes of the allocation
+ * @param caps Bitwise OR of MALLOC_CAP_* flags indicating the type of memory allocated.
  * @note this hook is called on the same thread as the allocation, which may be within a low level operation.
  * You should refrain from doing heavy work, logging, flash writes, or any locking.
  */
-typedef void (*esp_heap_trace_alloc_hook_t) (void* ptr, size_t size);
-
-/**
- * @brief registers a callback function to be invoked after every heap allocation
- * @param callback caller defined callback to be invoked
- * @return ESP_OK if callback was registered.
- */
-esp_err_t heap_caps_register_trace_alloc_callback(esp_heap_trace_alloc_hook_t callback);
+__attribute__((weak)) IRAM_ATTR void esp_heap_trace_alloc_hook(void* ptr, size_t size, uint32_t caps);
 
 /**
  * @brief callback called after every free
@@ -75,14 +70,7 @@ esp_err_t heap_caps_register_trace_alloc_callback(esp_heap_trace_alloc_hook_t ca
  * @note this hook is called on the same thread as the allocation, which may be within a low level operation.
  * You should refrain from doing heavy work, logging, flash writes, or any locking.
  */
-typedef void (*esp_heap_trace_free_hook_t) (void* ptr);
-
-/**
- * @brief registers a callback function to be invoked after every heap allocation
- * @param callback caller defined callback to be invoked
- * @return ESP_OK if callback was registered.
- */
-esp_err_t heap_caps_register_trace_free_callback(esp_heap_trace_free_hook_t callback);
+__attribute__((weak)) IRAM_ATTR void esp_heap_trace_free_hook(void* ptr);
 
 /**
  * @brief Allocate a chunk of memory which has the given capabilities

components/heap/multi_heap_poisoning.c

@@ -370,7 +370,7 @@ multi_heap_handle_t multi_heap_register(void *start, size_t size)
     return multi_heap_register_impl(start, size);
 }
 
-static inline void subtract_poison_overhead(size_t *arg) {
+static inline __attribute__((always_inline)) void subtract_poison_overhead(size_t *arg) {
     if (*arg > POISON_OVERHEAD) {
         *arg -= POISON_OVERHEAD;
     } else {

components/heap/test_apps/main/test_malloc.c

@@ -162,22 +162,88 @@ TEST_CASE("malloc/calloc(0) should not call failure callback", "[heap]")
 
 TEST_CASE("test get allocated size", "[heap]")
 {
-    const size_t iterations = 32;
+    // random values to test, some are 4 bytes aligned, some are not
+    const size_t alloc_sizes[] = { 1035, 1064, 1541 };
+    const size_t iterations = sizeof(alloc_sizes) / sizeof(size_t);
+    void *ptr_array[iterations];
 
     for (size_t i = 0; i < iterations; i++) {
-        // minimum block size is 12, so to avoid unecessary logic in the test,
-        // set the minimum requested size to 12.
-        const size_t alloc_size = rand() % 1024 + 12;
-
-        void *ptr = heap_caps_malloc(alloc_size, MALLOC_CAP_DEFAULT);
-        TEST_ASSERT_NOT_NULL(ptr);
+        ptr_array[i] = heap_caps_malloc(alloc_sizes[i], MALLOC_CAP_DEFAULT);
+        TEST_ASSERT_NOT_NULL(ptr_array[i]);
 
         // test that the heap_caps_get_allocated_size() returns the right number of bytes (aligned to 4 bytes
         // since the heap component aligns to 4 bytes)
-        const size_t aligned_size = (alloc_size + 3) & ~3;
-        printf("initial size: %d, requested size : %d, allocated size: %d\n", alloc_size, aligned_size, heap_caps_get_allocated_size(ptr));
-        TEST_ASSERT_EQUAL(aligned_size, heap_caps_get_allocated_size(ptr));
+        const size_t aligned_size = (alloc_sizes[i] + 3) & ~3;
+        const size_t real_size = heap_caps_get_allocated_size(ptr_array[i]);
+        printf("initial size: %d, requested size : %d, allocated size: %d\n", alloc_sizes[i], aligned_size, real_size);
+        TEST_ASSERT_EQUAL(aligned_size, real_size);
+
+        heap_caps_free(ptr_array[i]);
+    }
+}
 
-        heap_caps_free(ptr);
+// provide the definition of alloc and free hooks
+static const size_t alloc_size = 1234; // make this size atypical to be able to rely on it in the hook
+static const size_t expected_calls = 2; // one call for malloc/calloc and one call for realloc
+static uint32_t *alloc_ptr = NULL;
+static bool test_success = false;
+static size_t counter = 0;
+
+static void reset_static_variables(void) {
+    test_success = false;
+    alloc_ptr = NULL;
+    counter = 0;
+}
+
+void esp_heap_trace_alloc_hook(void* ptr, size_t size, uint32_t caps)
+{
+    if (size == alloc_size) {
+        counter++;
+        if (counter == expected_calls) {
+            alloc_ptr = ptr;
+        }
     }
 }
+
+void esp_heap_trace_free_hook(void* ptr)
+{
+    if (alloc_ptr == ptr && counter == expected_calls) {
+        test_success = true;
+    }
+}
+
+TEST_CASE("test allocation and free function hooks", "[heap]")
+{
+    // alloc, realloc and free memory, at the end of the test, test_success will be set
+    // to true if both function hooks are called.
+    uint32_t *ptr = heap_caps_malloc(alloc_size, MALLOC_CAP_DEFAULT);
+    TEST_ASSERT_NOT_NULL(ptr);
+    ptr = heap_caps_realloc(ptr, alloc_size, MALLOC_CAP_32BIT);
+    heap_caps_free(ptr);
+
+    TEST_ASSERT_TRUE(test_success);
+
+    // re-init the static variables
+    reset_static_variables();
+
+    // calloc, realloc and free memory, at the end of the test, test_success will be set
+    // to true if both function hooks are called.
+    ptr = heap_caps_calloc(1, alloc_size, MALLOC_CAP_DEFAULT);
+    TEST_ASSERT_NOT_NULL(ptr);
+    ptr = heap_caps_realloc(ptr, alloc_size, MALLOC_CAP_32BIT);
+    heap_caps_free(ptr);
+
+    TEST_ASSERT_TRUE(test_success);
+
+    // re-init the static variables
+    reset_static_variables();
+
+    // aligned alloc, realloc and aligned free memory, at the end of the test, test_success
+    // will be set to true if both function hooks are called.
+    ptr = heap_caps_aligned_alloc(0x200, alloc_size, MALLOC_CAP_DEFAULT);
+    TEST_ASSERT_NOT_NULL(ptr);
+    ptr = heap_caps_realloc(ptr, alloc_size, MALLOC_CAP_32BIT);
+    heap_caps_free(ptr);
+
+    TEST_ASSERT_TRUE(test_success);
+}