[umf] replace memory tracker with a critnib-based impl

This implementation is lock-free and does not require
C++ runtime.

source/common/unified_malloc_framework/CMakeLists.txt

@@ -6,8 +6,9 @@
 set(UMF_SOURCES
     src/memory_pool.c
     src/memory_provider.c
-    src/memory_tracker.cpp
+    src/memory_tracker.c
     src/memory_provider_get_last_failed.cpp
+    src/critnib/critnib.c
 )
 
 if(UMF_BUILD_SHARED_LIBRARY)

source/common/unified_malloc_framework/src/memory_tracker.cpp → source/common/unified_malloc_framework/src/memory_tracker.c

@@ -9,112 +9,117 @@
  */
 
 #include "memory_tracker.h"
+#include "critnib/critnib.h"
+
 #include <umf/memory_provider.h>
 #include <umf/memory_provider_ops.h>
 
-#include <cassert>
-#include <map>
-#include <mutex>
-#include <shared_mutex>
+#include <assert.h>
+#include <errno.h>
 #include <stdlib.h>
 
 #ifdef _WIN32
 #include <windows.h>
 #endif
 
-// TODO: reimplement in C and optimize...
-struct umf_memory_tracker_t {
-    enum umf_result_t add(void *pool, const void *ptr, size_t size) {
-        std::unique_lock<std::shared_mutex> lock(mtx);
+#if defined(_WIN32) && defined(UMF_SHARED_LIBRARY)
+critnib *TRACKER = NULL;
+BOOL APIENTRY DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved) {
+    if (fdwReason == DLL_PROCESS_DETACH) {
+        critnib_delete(TRACKER);
+    } else if (fdwReason == DLL_PROCESS_ATTACH) {
+        TRACKER = critnib_new();
+    }
+    return TRUE;
+}
+#elif defined(_WIN32)
+struct tracker_t {
+    tracker_t() { map = critnib_new(); }
+    ~tracker_t() { critnib_remove(map); }
+    critnib *map;
+};
+tracker_t TRACKER_INSTANCE;
+critnib *TRACKER = TRACKER_INSTANCE.map;
+#else
+critnib *TRACKER = NULL;
+void __attribute__((constructor)) createLibTracker() {
+    TRACKER = critnib_new();
+}
+void __attribute__((destructor)) deleteLibTracker() { critnib_delete(TRACKER); }
+#endif
 
-        if (size == 0) {
-            return UMF_RESULT_SUCCESS;
-        }
+umf_memory_tracker_handle_t umfMemoryTrackerGet(void) {
+    return (umf_memory_tracker_handle_t)TRACKER;
+}
 
-        auto ret =
-            map.try_emplace(reinterpret_cast<uintptr_t>(ptr), size, pool);
-        return ret.second ? UMF_RESULT_SUCCESS : UMF_RESULT_ERROR_UNKNOWN;
-    }
+struct tracker_value {
+    void *pool;
+    size_t size;
+};
 
-    enum umf_result_t remove(const void *ptr, size_t size) {
-        std::unique_lock<std::shared_mutex> lock(mtx);
+static enum umf_result_t
+umfMemoryTrackerAdd(umf_memory_tracker_handle_t hTracker, void *pool,
+                    const void *ptr, size_t size) {
+    assert(ptr);
 
-        map.erase(reinterpret_cast<uintptr_t>(ptr));
+    struct tracker_value *value =
+        (struct tracker_value *)malloc(sizeof(struct tracker_value));
+    value->pool = pool;
+    value->size = size;
 
-        // TODO: handle removing part of the range
-        (void)size;
+    int ret = critnib_insert((critnib *)hTracker, (uintptr_t)ptr, value, 0);
 
+    if (ret == 0) {
         return UMF_RESULT_SUCCESS;
     }
 
-    void *find(const void *ptr) {
-        std::shared_lock<std::shared_mutex> lock(mtx);
+    free(value);
 
-        auto intptr = reinterpret_cast<uintptr_t>(ptr);
-        auto it = map.upper_bound(intptr);
-        if (it == map.begin()) {
-            return nullptr;
-        }
-
-        --it;
-
-        auto address = it->first;
-        auto size = it->second.first;
-        auto pool = it->second.second;
-
-        if (intptr >= address && intptr < address + size) {
-            return pool;
-        }
-
-        return nullptr;
+    if (ret == ENOMEM) {
+        return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
     }
 
-  private:
-    std::shared_mutex mtx;
-    std::map<uintptr_t, std::pair<size_t, void *>> map;
-};
-
-static enum umf_result_t
-umfMemoryTrackerAdd(umf_memory_tracker_handle_t hTracker, void *pool,
-                    const void *ptr, size_t size) {
-    return hTracker->add(pool, ptr, size);
+    // This should not happen
+    // TODO: add logging here
+    return UMF_RESULT_ERROR_UNKNOWN;
 }
 
 static enum umf_result_t
 umfMemoryTrackerRemove(umf_memory_tracker_handle_t hTracker, const void *ptr,
                        size_t size) {
-    return hTracker->remove(ptr, size);
-}
-
-extern "C" {
-
-#if defined(_WIN32) && defined(UMF_SHARED_LIBRARY)
-umf_memory_tracker_t *tracker = nullptr;
-BOOL APIENTRY DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved) {
-    if (fdwReason == DLL_PROCESS_DETACH) {
-        delete tracker;
-    } else if (fdwReason == DLL_PROCESS_ATTACH) {
-        tracker = new umf_memory_tracker_t;
+    assert(ptr);
+
+    // TODO: there is no support for removing partial ranges (or multipe entires
+    // in a single remove call) yet.
+    // Every umfMemoryTrackerAdd(..., ptr, ...) should have a corresponsding
+    // umfMemoryTrackerRemove call with the same ptr value.
+    (void)size;
+
+    void *value = critnib_remove((critnib *)hTracker, (uintptr_t)ptr);
+    if (!value) {
+        // This should not happen
+        // TODO: add logging here
+        return UMF_RESULT_ERROR_UNKNOWN;
     }
-    return TRUE;
-}
-#elif defined(_WIN32)
-umf_memory_tracker_t trackerInstance;
-umf_memory_tracker_t *tracker = &trackerInstance;
-#else
-umf_memory_tracker_t *tracker = nullptr;
-void __attribute__((constructor)) createLibTracker() {
-    tracker = new umf_memory_tracker_t;
-}
 
-void __attribute__((destructor)) deleteLibTracker() { delete tracker; }
-#endif
+    free(value);
 
-umf_memory_tracker_handle_t umfMemoryTrackerGet(void) { return tracker; }
+    return UMF_RESULT_SUCCESS;
+}
 
 void *umfMemoryTrackerGetPool(umf_memory_tracker_handle_t hTracker,
                               const void *ptr) {
-    return hTracker->find(ptr);
+    assert(ptr);
+
+    uintptr_t rkey;
+    struct tracker_value *rvalue;
+    int found = critnib_find((critnib *)hTracker, (uintptr_t)ptr, FIND_LE,
+                             (void *)&rkey, (void **)&rvalue);
+    if (!found) {
+        return NULL;
+    }
+
+    return (rkey + rvalue->size >= (uintptr_t)ptr) ? rvalue->pool : NULL;
 }
 
 struct umf_tracking_memory_provider_t {
@@ -136,7 +141,7 @@ static enum umf_result_t trackingAlloc(void *hProvider, size_t size,
     }
 
     ret = umfMemoryProviderAlloc(p->hUpstream, size, alignment, ptr);
-    if (ret != UMF_RESULT_SUCCESS) {
+    if (ret != UMF_RESULT_SUCCESS || !*ptr) {
         return ret;
     }
 
@@ -159,9 +164,11 @@ static enum umf_result_t trackingFree(void *hProvider, void *ptr, size_t size) {
     // to avoid a race condition. If the order would be different, other thread
     // could allocate the memory at address `ptr` before a call to umfMemoryTrackerRemove
     // resulting in inconsistent state.
-    ret = umfMemoryTrackerRemove(p->hTracker, ptr, size);
-    if (ret != UMF_RESULT_SUCCESS) {
-        return ret;
+    if (ptr) {
+        ret = umfMemoryTrackerRemove(p->hTracker, ptr, size);
+        if (ret != UMF_RESULT_SUCCESS) {
+            return ret;
+        }
     }
 
     ret = umfMemoryProviderFree(p->hUpstream, ptr, size);
@@ -267,4 +274,3 @@ void umfTrackingMemoryProviderGetUpstreamProvider(
         (umf_tracking_memory_provider_t *)hTrackingProvider;
     *hUpstream = p->hUpstream;
 }
-}

test/unified_malloc_framework/common/provider.c

@@ -23,7 +23,7 @@ static enum umf_result_t nullAlloc(void *provider, size_t size,
     (void)provider;
     (void)size;
     (void)alignment;
-    (void)ptr;
+    *ptr = NULL;
     return UMF_RESULT_SUCCESS;
 }
 

test/unified_malloc_framework/memoryProviderAPI.cpp

@@ -23,7 +23,8 @@ TEST_F(test, memoryProviderTrace) {
 
     size_t call_count = 0;
 
-    auto ret = umfMemoryProviderAlloc(tracingProvider.get(), 0, 0, nullptr);
+    void *ptr;
+    auto ret = umfMemoryProviderAlloc(tracingProvider.get(), 0, 0, &ptr);
     ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
     ASSERT_EQ(calls["alloc"], 1);
     ASSERT_EQ(calls.size(), ++call_count);