Add task tracking option for heap usage monitoring

Add back a feature that was available in the old heap implementation
in release/v2.1 and earlier: keep track of which task allocates each
block from the heap. The task handle is conditionally added as
another word in the heap poisoning header under this configuration
option CONFIG_HEAP_TASK_TRACKING.

To allow custom monitoring and debugging code to be added, add helper
functions in multi_heap.c and multi_heap_poisoning.c to provide access
to information in the block headers.

components/heap/Kconfig

@@ -37,4 +37,13 @@ config HEAP_TRACING_STACK_DEPTH
         More stack frames uses more memory in the heap trace buffer (and slows down allocation), but
         can provide useful information.
 
+config HEAP_TASK_TRACKING
+    bool "Enable heap task tracking"
+    depends on !HEAP_POISONING_DISABLED
+    help
+        Enables tracking the task responsible for each heap allocation.
+
+        This function depends on heap poisoning being enabled and adds four more bytes of overhead for each block
+        allocated.
+
 endmenu

components/heap/component.mk

@@ -6,6 +6,10 @@ COMPONENT_OBJS := heap_caps_init.o heap_caps.o multi_heap.o heap_trace.o
 
 ifndef CONFIG_HEAP_POISONING_DISABLED
 COMPONENT_OBJS += multi_heap_poisoning.o
+
+ifdef CONFIG_HEAP_TASK_TRACKING
+COMPONENT_OBJS += esp_heap_debug.o
+endif
 endif
 
 ifdef CONFIG_HEAP_TRACING

components/heap/multi_heap.c

@@ -54,6 +54,14 @@ size_t multi_heap_free_size(multi_heap_handle_t heap)
 size_t multi_heap_minimum_free_size(multi_heap_handle_t heap)
     __attribute__((alias("multi_heap_minimum_free_size_impl")));
 
+void* multi_heap_get_block_address(multi_heap_block_handle_t block)
+    __attribute__((alias("multi_heap_get_block_address_impl")));
+
+void* multi_heap_get_block_owner(multi_heap_block_handle_t block)
+{
+    return NULL;
+}
+
 #endif
 
 #define ALIGN(X) ((X) & ~(sizeof(void *)-1))
@@ -279,6 +287,11 @@ static void split_if_necessary(heap_t *heap, heap_block_t *block, size_t size, h
     heap->free_bytes += block_data_size(new_block);
 }
 
+void* multi_heap_get_block_address_impl(multi_heap_block_handle_t block)
+{
+    return ((char *)block + offsetof(heap_block_t, data));
+}
+
 size_t multi_heap_get_allocated_size_impl(multi_heap_handle_t heap, void *p)
 {
     heap_block_t *pb = get_block(p);
@@ -339,6 +352,26 @@ void inline multi_heap_internal_unlock(multi_heap_handle_t heap)
     MULTI_HEAP_UNLOCK(heap->lock);
 }
 
+multi_heap_block_handle_t multi_heap_get_first_block(multi_heap_handle_t heap)
+{
+    return &heap->first_block;
+}
+
+multi_heap_block_handle_t multi_heap_get_next_block(multi_heap_handle_t heap, multi_heap_block_handle_t block)
+{
+    heap_block_t *next = get_next_block(block);
+    /* check for valid free last block to avoid assert in assert_valid_block */
+    if (next == heap->last_block && is_last_block(next) && is_free(next))
+	return NULL;
+    assert_valid_block(heap, next);
+    return next;
+}
+
+bool multi_heap_is_free(multi_heap_block_handle_t block)
+{
+    return is_free(block);
+}
+
 void *multi_heap_malloc_impl(multi_heap_handle_t heap, size_t size)
 {
     heap_block_t *best_block = NULL;

components/heap/multi_heap_internal.h

@@ -13,6 +13,9 @@
 // limitations under the License.
 #pragma once
 
+/* Opaque handle to a heap block */
+typedef const struct heap_block *multi_heap_block_handle_t;
+
 /* Internal definitions for the "implementation" of the multi_heap API,
    as defined in multi_heap.c.
 
@@ -28,6 +31,7 @@ void multi_heap_get_info_impl(multi_heap_handle_t heap, multi_heap_info_t *info)
 size_t multi_heap_free_size_impl(multi_heap_handle_t heap);
 size_t multi_heap_minimum_free_size_impl(multi_heap_handle_t heap);
 size_t multi_heap_get_allocated_size_impl(multi_heap_handle_t heap, void *p);
+void* multi_heap_get_block_address_impl(multi_heap_block_handle_t block);
 
 /* Some internal functions for heap poisoning use */
 
@@ -45,3 +49,20 @@ void multi_heap_internal_poison_fill_region(void *start, size_t size, bool is_fr
 void multi_heap_internal_lock(multi_heap_handle_t heap);
 
 void multi_heap_internal_unlock(multi_heap_handle_t heap);
+
+/* Some internal functions for heap debugging code to use */
+
+/* Get the handle to the first (fixed free) block in a heap */
+multi_heap_block_handle_t multi_heap_get_first_block(multi_heap_handle_t heap);
+
+/* Get the handle to the next block in a heap, with validation */
+multi_heap_block_handle_t multi_heap_get_next_block(multi_heap_handle_t heap, multi_heap_block_handle_t block);
+
+/* Test if a heap block is free */
+bool multi_heap_is_free(const multi_heap_block_handle_t block);
+
+/* Get the data address of a heap block */
+void* multi_heap_get_block_address(multi_heap_block_handle_t block);
+
+/* Get the owner identification for a heap block */
+void* multi_heap_get_block_owner(multi_heap_block_handle_t block);

components/heap/multi_heap_platform.h

@@ -63,6 +63,16 @@ inline static void multi_heap_assert(bool condition, const char *format, int lin
     multi_heap_assert((CONDITION), "CORRUPT HEAP: multi_heap.c:%d detected at 0x%08x\n", \
                       __LINE__, (intptr_t)(ADDRESS))
 
+#ifdef CONFIG_HEAP_TASK_TRACKING
+#define MULTI_HEAP_BLOCK_OWNER TaskHandle_t task;
+#define MULTI_HEAP_SET_BLOCK_OWNER(HEAD) HEAD->task = xTaskGetCurrentTaskHandle()
+#define MULTI_HEAP_GET_BLOCK_OWNER(HEAD) (HEAD->task)
+#else
+#define MULTI_HEAP_BLOCK_OWNER
+#define MULTI_HEAP_SET_BLOCK_OWNER(HEAD)
+#define MULTI_HEAP_GET_BLOCK_OWNER(HEAD)
+#endif
+
 #else // ESP_PLATFORM
 
 #include <assert.h>
@@ -73,4 +83,8 @@ inline static void multi_heap_assert(bool condition, const char *format, int lin
 #define MULTI_HEAP_UNLOCK(PLOCK)
 
 #define MULTI_HEAP_ASSERT(CONDITION, ADDRESS) assert((CONDITION) && "Heap corrupt")
+
+#define MULTI_HEAP_BLOCK_OWNER
+#define MULTI_HEAP_SET_BLOCK_OWNER(HEAD)
+#define MULTI_HEAP_GET_BLOCK_OWNER(HEAD)
 #endif

components/heap/multi_heap_poisoning.c

@@ -47,6 +47,7 @@
 
 typedef struct {
     uint32_t head_canary;
+    MULTI_HEAP_BLOCK_OWNER
     size_t alloc_size;
 } poison_head_t;
 
@@ -67,6 +68,7 @@ static uint8_t *poison_allocated_region(poison_head_t *head, size_t alloc_size)
     poison_tail_t *tail = (poison_tail_t *)(data + alloc_size);
     head->alloc_size = alloc_size;
     head->head_canary = HEAD_CANARY_PATTERN;
+    MULTI_HEAP_SET_BLOCK_OWNER(head);
 
     uint32_t tail_canary = TAIL_CANARY_PATTERN;
     if ((intptr_t)tail % sizeof(void *) == 0) {
@@ -258,6 +260,12 @@ void *multi_heap_realloc(multi_heap_handle_t heap, void *p, size_t size)
     return result;
 }
 
+void* multi_heap_get_block_address(multi_heap_block_handle_t block)
+{
+    char *head = multi_heap_get_block_address_impl(block);
+    return head + sizeof(poison_head_t);
+}
+
 size_t multi_heap_get_allocated_size(multi_heap_handle_t heap, void *p)
 {
     poison_head_t *head = verify_allocated_region(p, true);
@@ -269,6 +277,12 @@ size_t multi_heap_get_allocated_size(multi_heap_handle_t heap, void *p)
     return 0;
 }
 
+void* multi_heap_get_block_owner(multi_heap_block_handle_t block)
+{
+    poison_head_t *head = (poison_head_t*)multi_heap_get_block_address_impl(block);
+    return MULTI_HEAP_GET_BLOCK_OWNER(head);
+}
+
 multi_heap_handle_t multi_heap_register(void *start, size_t size)
 {
     if (start != NULL) {