Use memory tracking annotations in base alloc, disjoint_pool and jemalloc_pool

Author: igchor

src/base_alloc/base_alloc.c

@@ -11,6 +11,7 @@
 #include "base_alloc_internal.h"
 #include "utils_common.h"
 #include "utils_concurrency.h"
+#include "utils_sanitizers.h"
 
 // minimum size of a single pool of the base allocator
 #define MINIMUM_POOL_SIZE (ba_os_get_page_size())
@@ -80,7 +81,10 @@ static void ba_debug_checks(umf_ba_pool_t *pool) {
     umf_ba_chunk_t *next_chunk = pool->metadata.free_list;
     while (next_chunk) {
         n_free_chunks++;
+        utils_annotate_memory_defined(next_chunk, sizeof(umf_ba_chunk_t));
+        umf_ba_chunk_t *tmp = next_chunk;
         next_chunk = next_chunk->next;
+        utils_annotate_memory_inaccessible(tmp, sizeof(umf_ba_chunk_t));
     }
     assert(n_free_chunks == pool->metadata.n_chunks - pool->metadata.n_allocs);
 }
@@ -89,6 +93,9 @@ static void ba_debug_checks(umf_ba_pool_t *pool) {
 // ba_divide_memory_into_chunks - divide given memory into chunks of chunk_size and add them to the free_list
 static void ba_divide_memory_into_chunks(umf_ba_pool_t *pool, void *ptr,
                                          size_t size) {
+    // mark the memory temporarily accessible to perform the division
+    utils_annotate_memory_undefined(ptr, size);
+
     assert(pool->metadata.free_list == NULL);
     assert(size > pool->metadata.chunk_size);
 
@@ -112,6 +119,17 @@ static void ba_divide_memory_into_chunks(umf_ba_pool_t *pool, void *ptr,
 
     current_chunk->next = NULL;
     pool->metadata.free_list = ptr; // address of the first chunk
+
+    // mark the memory as unaccessible again
+    utils_annotate_memory_inaccessible(ptr, size);
+}
+
+static void *ba_os_alloc_annotated(size_t pool_size) {
+    void *ptr = ba_os_alloc(pool_size);
+    if (ptr) {
+        utils_annotate_memory_inaccessible(ptr, pool_size);
+    }
+    return ptr;
 }
 
 umf_ba_pool_t *umf_ba_create(size_t size) {
@@ -127,11 +145,14 @@ umf_ba_pool_t *umf_ba_create(size_t size) {
 
     pool_size = ALIGN_UP(pool_size, ba_os_get_page_size());
 
-    umf_ba_pool_t *pool = (umf_ba_pool_t *)ba_os_alloc(pool_size);
+    umf_ba_pool_t *pool = (umf_ba_pool_t *)ba_os_alloc_annotated(pool_size);
     if (!pool) {
         return NULL;
     }
 
+    // annotate metadata region as accessible
+    utils_annotate_memory_undefined(pool, offsetof(umf_ba_pool_t, data));
+
     pool->metadata.pool_size = pool_size;
     pool->metadata.chunk_size = chunk_size;
     pool->next_pool = NULL; // this is the only pool now
@@ -141,6 +162,8 @@ umf_ba_pool_t *umf_ba_create(size_t size) {
     pool->metadata.n_chunks = 0;
 #endif /* NDEBUG */
 
+    utils_annotate_memory_defined(pool, offsetof(umf_ba_pool_t, data));
+
     char *data_ptr = (char *)&pool->data;
     size_t size_left = pool_size - offsetof(umf_ba_pool_t, data);
 
@@ -163,12 +186,16 @@ void *umf_ba_alloc(umf_ba_pool_t *pool) {
     util_mutex_lock(&pool->metadata.free_lock);
     if (pool->metadata.free_list == NULL) {
         umf_ba_next_pool_t *new_pool =
-            (umf_ba_next_pool_t *)ba_os_alloc(pool->metadata.pool_size);
+            (umf_ba_next_pool_t *)ba_os_alloc_annotated(
+                pool->metadata.pool_size);
         if (!new_pool) {
             util_mutex_unlock(&pool->metadata.free_lock);
             return NULL;
         }
 
+        // annotate metadata region as accessible
+        utils_annotate_memory_undefined(new_pool, sizeof(umf_ba_next_pool_t));
+
         // add the new pool to the list of pools
         new_pool->next_pool = pool->next_pool;
         pool->next_pool = new_pool;
@@ -186,11 +213,20 @@ void *umf_ba_alloc(umf_ba_pool_t *pool) {
     }
 
     umf_ba_chunk_t *chunk = pool->metadata.free_list;
+
+    // mark the memory defined to read the next ptr, after this is done
+    // we'll mark the memory as undefined
+    utils_annotate_memory_defined(chunk, sizeof(chunk));
+
     pool->metadata.free_list = pool->metadata.free_list->next;
     pool->metadata.n_allocs++;
 #ifndef NDEBUG
     ba_debug_checks(pool);
 #endif /* NDEBUG */
+
+    VALGRIND_DO_MALLOCLIKE_BLOCK(chunk, pool->metadata.chunk_size, 0, 0);
+    utils_annotate_memory_undefined(chunk, pool->metadata.chunk_size);
+
     util_mutex_unlock(&pool->metadata.free_lock);
 
     return chunk;
@@ -234,6 +270,10 @@ void umf_ba_free(umf_ba_pool_t *pool, void *ptr) {
 #ifndef NDEBUG
     ba_debug_checks(pool);
 #endif /* NDEBUG */
+
+    VALGRIND_DO_FREELIKE_BLOCK(chunk, 0);
+    utils_annotate_memory_inaccessible(chunk, pool->metadata.chunk_size);
+
     util_mutex_unlock(&pool->metadata.free_lock);
 }
 

src/pool/pool_disjoint.cpp

@@ -26,6 +26,7 @@
 #include "pool_disjoint.h"
 #include "umf.h"
 #include "utils_math.h"
+#include "utils_sanitizers.h"
 
 typedef struct umf_disjoint_pool_shared_limits_t {
     size_t MaxSize;
@@ -329,6 +330,8 @@ class DisjointPool::AllocImpl {
               umf_disjoint_pool_params_t *params)
         : MemHandle{hProvider}, params(*params) {
 
+        VALGRIND_DO_CREATE_MEMPOOL(this, 0, 0);
+
         // Generate buckets sized such as: 64, 96, 128, 192, ..., CutOff.
         // Powers of 2 and the value halfway between the powers of 2.
         auto Size1 = this->params.MinBucketSize;
@@ -352,6 +355,8 @@ class DisjointPool::AllocImpl {
         }
     }
 
+    ~AllocImpl() { VALGRIND_DO_DESTROY_MEMPOOL(this); }
+
     void *allocate(size_t Size, size_t Alignment, bool &FromPool);
     void *allocate(size_t Size, bool &FromPool);
     void deallocate(void *Ptr, bool &ToPool);
@@ -392,6 +397,7 @@ static void *memoryProviderAlloc(umf_memory_provider_handle_t hProvider,
     if (ret != UMF_RESULT_SUCCESS) {
         throw MemoryProviderError{ret};
     }
+    utils_annotate_memory_inaccessible(ptr, size);
     return ptr;
 }
 
@@ -798,7 +804,9 @@ void *DisjointPool::AllocImpl::allocate(size_t Size, bool &FromPool) try {
 
     FromPool = false;
     if (Size > getParams().MaxPoolableSize) {
-        return memoryProviderAlloc(getMemHandle(), Size);
+        Ptr = memoryProviderAlloc(getMemHandle(), Size);
+        utils_annotate_memory_undefined(Ptr, Size);
+        return Ptr;
     }
 
     auto &Bucket = findBucket(Size);
@@ -813,6 +821,9 @@ void *DisjointPool::AllocImpl::allocate(size_t Size, bool &FromPool) try {
         Bucket.countAlloc(FromPool);
     }
 
+    VALGRIND_DO_MEMPOOL_ALLOC(this, Ptr, Size);
+    utils_annotate_memory_undefined(Ptr, Bucket.getSize());
+
     return Ptr;
 } catch (MemoryProviderError &e) {
     umf::getPoolLastStatusRef<DisjointPool>() = e.code;
@@ -848,7 +859,9 @@ void *DisjointPool::AllocImpl::allocate(size_t Size, size_t Alignment,
     // If not, just request aligned pointer from the system.
     FromPool = false;
     if (AlignedSize > getParams().MaxPoolableSize) {
-        return memoryProviderAlloc(getMemHandle(), Size, Alignment);
+        Ptr = memoryProviderAlloc(getMemHandle(), Size, Alignment);
+        utils_annotate_memory_undefined(Ptr, Size);
+        return Ptr;
     }
 
     auto &Bucket = findBucket(AlignedSize);
@@ -863,6 +876,9 @@ void *DisjointPool::AllocImpl::allocate(size_t Size, size_t Alignment,
         Bucket.countAlloc(FromPool);
     }
 
+    VALGRIND_DO_MEMPOOL_ALLOC(this, AlignPtrUp(Ptr, Alignment), Size);
+    utils_annotate_memory_undefined(AlignPtrUp(Ptr, Alignment), Size);
+
     return AlignPtrUp(Ptr, Alignment);
 } catch (MemoryProviderError &e) {
     umf::getPoolLastStatusRef<DisjointPool>() = e.code;
@@ -929,6 +945,9 @@ void DisjointPool::AllocImpl::deallocate(void *Ptr, bool &ToPool) {
                 Bucket.countFree();
             }
 
+            VALGRIND_DO_MEMPOOL_FREE(this, Ptr);
+            utils_annotate_memory_inaccessible(Ptr, Bucket.getSize());
+
             if (Bucket.getSize() <= Bucket.ChunkCutOff()) {
                 Bucket.freeChunk(Ptr, Slab, ToPool);
             } else {

src/pool/pool_jemalloc.c

@@ -16,6 +16,8 @@
 #include "base_alloc_global.h"
 #include "utils_common.h"
 #include "utils_concurrency.h"
+#include "utils_sanitizers.h"
+
 #include <umf/memory_pool.h>
 #include <umf/memory_pool_ops.h>
 #include <umf/pools/pool_jemalloc.h>
@@ -72,7 +74,14 @@ static void *arena_extent_alloc(extent_hooks_t *extent_hooks, void *new_addr,
         return NULL;
     }
 
+#ifndef __SANITIZE_ADDRESS__
+    // jemalloc might write to new extents in realloc, so we cannot
+    // mark them as unaccessible under asan
+    utils_annotate_memory_inaccessible(ptr, size);
+#endif
+
     if (*zero) {
+        utils_annotate_memory_defined(ptr, size);
         memset(ptr, 0, size); // TODO: device memory is not accessible by host
     }
 
@@ -282,6 +291,8 @@ static void *je_malloc(void *pool, size_t size) {
         return NULL;
     }
 
+    VALGRIND_DO_MEMPOOL_ALLOC(pool, ptr, size);
+
     return ptr;
 }
 
@@ -290,6 +301,7 @@ static umf_result_t je_free(void *pool, void *ptr) {
     assert(pool);
 
     if (ptr != NULL) {
+        VALGRIND_DO_MEMPOOL_FREE(pool, ptr);
         dallocx(ptr, MALLOCX_TCACHE_NONE);
     }
 
@@ -305,6 +317,8 @@ static void *je_calloc(void *pool, size_t num, size_t size) {
         return NULL;
     }
 
+    utils_annotate_memory_defined(ptr, num * size);
+
     memset(ptr, 0, csize); // TODO: device memory is not accessible by host
     return ptr;
 }
@@ -314,6 +328,7 @@ static void *je_realloc(void *pool, void *ptr, size_t size) {
     if (size == 0 && ptr != NULL) {
         dallocx(ptr, MALLOCX_TCACHE_NONE);
         TLS_last_allocation_error = UMF_RESULT_SUCCESS;
+        VALGRIND_DO_MEMPOOL_FREE(pool, ptr);
         return NULL;
     } else if (ptr == NULL) {
         return je_malloc(pool, size);
@@ -329,6 +344,14 @@ static void *je_realloc(void *pool, void *ptr, size_t size) {
         return NULL;
     }
 
+    if (new_ptr != ptr) {
+        VALGRIND_DO_MEMPOOL_ALLOC(pool, new_ptr, size);
+        VALGRIND_DO_MEMPOOL_FREE(pool, ptr);
+
+        // memory was copied from old ptr so it's now defined
+        utils_annotate_memory_defined(new_ptr, size);
+    }
+
     return new_ptr;
 }
 
@@ -346,6 +369,8 @@ static void *je_aligned_alloc(void *pool, size_t size, size_t alignment) {
         return NULL;
     }
 
+    VALGRIND_DO_MEMPOOL_ALLOC(pool, ptr, size);
+
     return ptr;
 }
 
@@ -392,6 +417,8 @@ static umf_result_t je_initialize(umf_memory_provider_handle_t provider,
 
     *out_pool = (umf_memory_pool_handle_t)pool;
 
+    VALGRIND_DO_CREATE_MEMPOOL(pool, 0, 0);
+
     return UMF_RESULT_SUCCESS;
 
 err_free_pool:
@@ -407,6 +434,8 @@ static void je_finalize(void *pool) {
     mallctl(cmd, NULL, 0, NULL, 0);
     pool_by_arena_index[je_pool->arena_index] = NULL;
     umf_ba_global_free(je_pool);
+
+    VALGRIND_DO_DESTROY_MEMPOOL(pool);
 }
 
 static size_t je_malloc_usable_size(void *pool, void *ptr) {

test/common/pool.hpp

@@ -45,6 +45,7 @@ bool isReallocSupported(umf_memory_pool_handle_t hPool) {
     static constexpr size_t allocSize = 8;
     bool supported = false;
     auto *ptr = umfPoolMalloc(hPool, allocSize);
+    memset(ptr, 0, allocSize);
     auto *new_ptr = umfPoolRealloc(hPool, ptr, allocSize * 2);
 
     if (new_ptr) {

test/poolFixtures.hpp

@@ -115,6 +115,7 @@ TEST_P(umfPoolTest, reallocFree) {
     static constexpr size_t multiplier = 3;
     auto *ptr = umfPoolMalloc(pool.get(), allocSize);
     ASSERT_NE(ptr, nullptr);
+    memset(ptr, 0, allocSize);
     auto *new_ptr = umfPoolRealloc(pool.get(), ptr, allocSize * multiplier);
     ASSERT_NE(new_ptr, nullptr);
     std::memset(new_ptr, 0, allocSize * multiplier);

test/supp/umf_test-jemalloc_pool.supp

@@ -0,0 +1,8 @@
+{
+   realloc false-positive suppresion - see pool_jemalloc.c for more details
+   Memcheck:Addr8
+   fun:memmove
+   ...
+   fun:do_rallocx
+   fun:je_realloc
+}