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

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

Author: igchor

source/common/unified_malloc_framework/CMakeLists.txt

@@ -6,10 +6,15 @@
 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 (MSVC)
+    set(UMF_SOURCES ${UMF_SOURCES} src/memory_tracker_windows.cpp)
+endif()
+
 if(UMF_BUILD_SHARED_LIBRARY)
     message(WARNING "Unified Malloc Framework is still an early work in progress."
                      "There are no API/ABI backward compatibility guarantees. There will be breakages."

source/common/unified_malloc_framework/src/memory_tracker.cpp renamed to source/common/unified_malloc_framework/src/memory_tracker.c

@@ -9,112 +9,97 @@
  */
 
 #include "memory_tracker.h"
+#include "critnib/critnib.h"
+
+#include <umf/memory_pool.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)
+critnib *TRACKER = NULL;
+void __attribute__((constructor)) createLibTracker() {
+    TRACKER = critnib_new();
+}
+void __attribute__((destructor)) deleteLibTracker() { critnib_delete(TRACKER); }
 
-        if (size == 0) {
-            return UMF_RESULT_SUCCESS;
-        }
+umf_memory_tracker_handle_t umfMemoryTrackerGet(void) {
+    return (umf_memory_tracker_handle_t)TRACKER;
+}
+#endif
 
-        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_t {
+    umf_memory_pool_handle_t 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,
+                    umf_memory_pool_handle_t pool, const void *ptr,
+                    size_t size) {
+    assert(ptr);
 
-        map.erase(reinterpret_cast<uintptr_t>(ptr));
+    struct tracker_value_t *value =
+        (struct tracker_value_t *)malloc(sizeof(struct tracker_value_t));
+    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);
-
-        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;
-        }
+    free(value);
 
-        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);
-}
+    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;
+    }
 
-extern "C" {
+    free(value);
 
-#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;
-    }
-    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;
+    return UMF_RESULT_SUCCESS;
 }
 
-void __attribute__((destructor)) deleteLibTracker() { delete tracker; }
-#endif
+umf_memory_pool_handle_t
+umfMemoryTrackerGetPool(umf_memory_tracker_handle_t hTracker, const void *ptr) {
+    assert(ptr);
 
-umf_memory_tracker_handle_t umfMemoryTrackerGet(void) { return tracker; }
+    uintptr_t rkey;
+    struct tracker_value_t *rvalue;
+    int found = critnib_find((critnib *)hTracker, (uintptr_t)ptr, FIND_LE,
+                             (void *)&rkey, (void **)&rvalue);
+    if (!found) {
+        return NULL;
+    }
 
-void *umfMemoryTrackerGetPool(umf_memory_tracker_handle_t hTracker,
-                              const void *ptr) {
-    return hTracker->find(ptr);
+    return (rkey + rvalue->size >= (uintptr_t)ptr) ? rvalue->pool : NULL;
 }
 
 struct umf_tracking_memory_provider_t {
@@ -136,7 +121,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 +144,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 +254,3 @@ void umfTrackingMemoryProviderGetUpstreamProvider(
         (umf_tracking_memory_provider_t *)hTrackingProvider;
     *hUpstream = p->hUpstream;
 }
-}

source/common/unified_malloc_framework/src/memory_tracker.h

@@ -22,8 +22,8 @@ extern "C" {
 typedef struct umf_memory_tracker_t *umf_memory_tracker_handle_t;
 
 umf_memory_tracker_handle_t umfMemoryTrackerGet(void);
-void *umfMemoryTrackerGetPool(umf_memory_tracker_handle_t hTracker,
-                              const void *ptr);
+umf_memory_pool_handle_t
+umfMemoryTrackerGetPool(umf_memory_tracker_handle_t hTracker, const void *ptr);
 
 // Creates a memory provider that tracks each allocation/deallocation through umf_memory_tracker_handle_t and
 // forwards all requests to hUpstream memory Provider. hUpstream lifetime should be managed by the user of this function.

source/common/unified_malloc_framework/src/memory_tracker_windows.cpp

@@ -0,0 +1,35 @@
+/*
+ *
+ * Copyright (C) 2023 Intel Corporation
+ *
+ * Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM Exceptions.
+ * See LICENSE.TXT
+ * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+ *
+ */
+
+#include <windows.h>
+
+#if 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;
+}
+#else
+struct tracker_t {
+    tracker_t() { map = critnib_new(); }
+    ~tracker_t() { critnib_remove(map); }
+    critnib *map;
+};
+tracker_t TRACKER_INSTANCE;
+critnib *TRACKER = TRACKER_INSTANCE.map;
+#endif
+
+umf_memory_tracker_handle_t umfMemoryTrackerGet(void) {
+    return (umf_memory_tracker_handle_t)TRACKER;
+}

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);