Implement umfMemoryProviderAllocation[Split/Merge]

igchor · igchor · commit e9b8495e0d10 · 2024-01-17T23:15:07.000+01:00
That is needed to support jemalloc (avoid leaking memory).
diff --git a/include/umf/memory_provider.h b/include/umf/memory_provider.h
@@ -156,6 +156,13 @@ umfMemoryProviderGetName(umf_memory_provider_handle_t hProvider);
 ///
 UMF_EXPORT umf_memory_provider_handle_t umfGetLastFailedMemoryProvider(void);
 
+UMF_EXPORT umf_result_t
+umfMemoryProviderAllocationSplit(umf_memory_provider_handle_t hProvider,
+                                 void *ptr, size_t totalSize, size_t leftSize);
+UMF_EXPORT umf_result_t
+umfMemoryProviderAllocationMerge(umf_memory_provider_handle_t hProvider,
+                                 void *lowPtr, void *highPtr, size_t totalSize);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/include/umf/memory_provider_ops.h b/include/umf/memory_provider_ops.h
@@ -139,6 +139,11 @@ typedef struct umf_memory_provider_ops_t {
     /// @return pointer to a string containing the name of the \p provider
     ///
     const char *(*get_name)(void *provider);
+
+    umf_result_t (*allocation_split)(void *provider, void *ptr,
+                                     size_t totalSize, size_t leftSize);
+    umf_result_t (*allocation_merge)(void *provider, void *lowPtr,
+                                     void *highPtr, size_t totalSize);
 } umf_memory_provider_ops_t;
 
 #ifdef __cplusplus
diff --git a/src/memory_provider.c b/src/memory_provider.c
@@ -11,6 +11,7 @@
 #include <umf/memory_provider.h>
 
 #include <assert.h>
+#include <stdio.h>
 #include <stdlib.h>
 
 typedef struct umf_memory_provider_t {
@@ -129,3 +130,46 @@ const char *umfMemoryProviderGetName(umf_memory_provider_handle_t hProvider) {
 umf_memory_provider_handle_t umfGetLastFailedMemoryProvider(void) {
     return *umfGetLastFailedMemoryProviderPtr();
 }
+
+umf_result_t
+umfMemoryProviderAllocationSplit(umf_memory_provider_handle_t hProvider,
+                                 void *ptr, size_t totalSize, size_t leftSize) {
+    if (leftSize >= totalSize) {
+        fprintf(stderr,
+                "umfMemoryProviderAllocationSplit failed. check failed: %zu >= "
+                "%zu\n",
+                leftSize, totalSize);
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    umf_result_t res = hProvider->ops.allocation_split(
+        hProvider->provider_priv, ptr, totalSize, leftSize);
+    checkErrorAndSetLastProvider(res, hProvider);
+    return res;
+}
+
+umf_result_t
+umfMemoryProviderAllocationMerge(umf_memory_provider_handle_t hProvider,
+                                 void *lowPtr, void *highPtr,
+                                 size_t totalSize) {
+    if ((uintptr_t)lowPtr >= (uintptr_t)highPtr) {
+        fprintf(
+            stderr,
+            "umfMemoryProviderAllocationMerge failed. check failed: %p >= %p\n",
+            lowPtr, highPtr);
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    if ((uintptr_t)highPtr - (uintptr_t)lowPtr > totalSize) {
+        fprintf(stderr,
+                "umfMemoryProviderAllocationMerge failed. check failed: %p - "
+                "%p > %zu\n",
+                highPtr, lowPtr, totalSize);
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    umf_result_t res = hProvider->ops.allocation_merge(
+        hProvider->provider_priv, lowPtr, highPtr, totalSize);
+    checkErrorAndSetLastProvider(res, hProvider);
+    return res;
+}
diff --git a/src/pool/pool_jemalloc.c b/src/pool/pool_jemalloc.c
@@ -210,15 +210,14 @@ static bool arena_extent_split(extent_hooks_t *extent_hooks, void *addr,
                                size_t size, size_t size_a, size_t size_b,
                                bool committed, unsigned arena_ind) {
     (void)extent_hooks; // unused
-    (void)addr;         // unused
-    (void)size;         // unused
-    (void)size_a;       // unused
-    (void)size_b;       // unused
     (void)committed;    // unused
-    (void)arena_ind;    // unused
+    (void)size_b;
 
-    // TODO: add this function to the provider API to support Windows and USM
-    return true; // true means failure (unsupported)
+    assert(size_a + size_b == size);
+
+    jemalloc_memory_pool_t *pool = get_pool_by_arena_index(arena_ind);
+    return umfMemoryProviderAllocationSplit(pool->provider, addr, size,
+                                            size_a) != UMF_RESULT_SUCCESS;
 }
 
 // arena_extent_merge - an extent merge function conforms to the extent_merge_t type and optionally
@@ -231,15 +230,12 @@ static bool arena_extent_merge(extent_hooks_t *extent_hooks, void *addr_a,
                                size_t size_a, void *addr_b, size_t size_b,
                                bool committed, unsigned arena_ind) {
     (void)extent_hooks; // unused
-    (void)addr_a;       // unused
-    (void)addr_b;       // unused
-    (void)size_a;       // unused
-    (void)size_b;       // unused
     (void)committed;    // unused
-    (void)arena_ind;    // unused
 
-    // TODO: add this function to the provider API to support Windows and USM
-    return true; // true means failure (unsupported)
+    jemalloc_memory_pool_t *pool = get_pool_by_arena_index(arena_ind);
+    return umfMemoryProviderAllocationMerge(pool->provider, addr_a, addr_b,
+                                            size_a + size_b) !=
+           UMF_RESULT_SUCCESS;
 }
 
 // The extent_hooks_t structure comprises function pointers which are described individually below.
diff --git a/src/pool/pool_scalable.c b/src/pool/pool_scalable.c
@@ -121,7 +121,7 @@ static void tbb_raw_free_wrapper(intptr_t pool_id, void *ptr, size_t bytes) {
         TLS_last_free_error = ret;
         fprintf(
             stderr,
-            "Memory provider failed to free memory, addr = %p, size = %lu\n",
+            "Memory provider failed to free memory, addr = %p, size = %zu\n",
             ptr, bytes);
     }
 }
diff --git a/src/provider/provider_os_memory.c b/src/provider/provider_os_memory.c
@@ -449,6 +449,26 @@ static const char *os_get_name(void *provider) {
     return "OS";
 }
 
+static umf_result_t os_allocation_split(void *provider, void *ptr,
+                                        size_t totalSize, size_t leftSize) {
+    (void)provider;
+    (void)ptr;
+    (void)totalSize;
+    (void)leftSize;
+    // nop
+    return UMF_RESULT_SUCCESS;
+}
+
+static umf_result_t os_allocation_merge(void *provider, void *lowPtr,
+                                        void *highPtr, size_t totalSize) {
+    (void)provider;
+    (void)lowPtr;
+    (void)highPtr;
+    (void)totalSize;
+    // nop
+    return UMF_RESULT_SUCCESS;
+}
+
 umf_memory_provider_ops_t UMF_OS_MEMORY_PROVIDER_OPS = {
     .version = UMF_VERSION_CURRENT,
     .initialize = os_initialize,
@@ -461,4 +481,5 @@ umf_memory_provider_ops_t UMF_OS_MEMORY_PROVIDER_OPS = {
     .purge_lazy = os_purge_lazy,
     .purge_force = os_purge_force,
     .get_name = os_get_name,
-};
+    .allocation_split = os_allocation_split,
+    .allocation_merge = os_allocation_merge};
diff --git a/src/provider/provider_tracking.c b/src/provider/provider_tracking.c
@@ -9,13 +9,15 @@
 
 #include "provider_tracking.h"
 #include "critnib.h"
+#include "utils_concurrency.h"
 
 #include <umf/memory_pool.h>
 #include <umf/memory_provider.h>
 #include <umf/memory_provider_ops.h>
 
 #include <assert.h>
 #include <errno.h>
+#include <stdio.h>
 #include <stdlib.h>
 
 typedef struct tracker_value_t {
@@ -32,7 +34,7 @@ static umf_result_t umfMemoryTrackerAdd(umf_memory_tracker_handle_t hTracker,
     value->pool = pool;
     value->size = size;
 
-    int ret = critnib_insert((critnib *)hTracker, (uintptr_t)ptr, value, 0);
+    int ret = critnib_insert(hTracker->map, (uintptr_t)ptr, value, 0);
 
     if (ret == 0) {
         return UMF_RESULT_SUCCESS;
@@ -59,7 +61,7 @@ static umf_result_t umfMemoryTrackerRemove(umf_memory_tracker_handle_t hTracker,
     // umfMemoryTrackerRemove call with the same ptr value.
     (void)size;
 
-    void *value = critnib_remove((critnib *)hTracker, (uintptr_t)ptr);
+    void *value = critnib_remove(hTracker->map, (uintptr_t)ptr);
     if (!value) {
         // This should not happen
         // TODO: add logging here
@@ -77,7 +79,7 @@ umfMemoryTrackerGetPool(umf_memory_tracker_handle_t hTracker, const void *ptr) {
 
     uintptr_t rkey;
     tracker_value_t *rvalue;
-    int found = critnib_find((critnib *)hTracker, (uintptr_t)ptr, FIND_LE,
+    int found = critnib_find(hTracker->map, (uintptr_t)ptr, FIND_LE,
                              (void *)&rkey, (void **)&rvalue);
     if (!found) {
         return NULL;
@@ -119,6 +121,158 @@ static umf_result_t trackingAlloc(void *hProvider, size_t size,
     return ret;
 }
 
+static umf_result_t trackingAllocationSplit(void *provider, void *ptr,
+                                            size_t totalSize, size_t leftSize) {
+    umf_result_t ret = UMF_RESULT_ERROR_UNKNOWN;
+    umf_tracking_memory_provider_t *p =
+        (umf_tracking_memory_provider_t *)provider;
+
+    tracker_value_t *splitValue =
+        (tracker_value_t *)malloc(sizeof(tracker_value_t));
+    splitValue->pool = p->pool;
+    splitValue->size = leftSize;
+
+    int r = util_mutex_lock(p->hTracker->splitMergeMutex);
+    if (r) {
+        goto err_lock;
+    }
+
+    tracker_value_t *value =
+        (tracker_value_t *)critnib_get(p->hTracker->map, (uintptr_t)ptr);
+    if (!value) {
+        fprintf(stderr, "tracking split: no such value\n");
+        ret = UMF_RESULT_ERROR_INVALID_ARGUMENT;
+        goto err;
+    }
+    if (value->size != totalSize) {
+        fprintf(stderr, "tracking split: %zu != %zu\n", value->size, totalSize);
+        ret = UMF_RESULT_ERROR_INVALID_ARGUMENT;
+        goto err;
+    }
+
+    ret = umfMemoryProviderAllocationSplit(p->hUpstream, ptr, totalSize,
+                                           leftSize);
+    if (ret != UMF_RESULT_SUCCESS) {
+        fprintf(stderr,
+                "tracking split: umfMemoryProviderAllocationSplit failed\n");
+        goto err;
+    }
+
+    void *highPtr = (void *)(((uintptr_t)ptr) + leftSize);
+    size_t rightSize = totalSize - leftSize;
+
+    // We'll have duplicate entry for the range [highPtr, rightValue->size] but this is fine,
+    // the value is the same anyway and we forbid splitting/removing that range concurrently
+    ret = umfMemoryTrackerAdd(p->hTracker, p->pool, highPtr, rightSize);
+    if (ret != UMF_RESULT_SUCCESS) {
+        fprintf(stderr, "tracking split: umfMemoryTrackerAdd failed\n");
+        // TODO: what now? should we rollback the split? This is can only happen due to ENOMEM
+        // so it's unlikely but probably the best solution would be to try to preallocate everything
+        // (value and critnib nodes) before calling umfMemoryProviderAllocationSplit.
+        goto err;
+    }
+
+    int cret = critnib_insert(p->hTracker->map, (uintptr_t)ptr,
+                              (void *)splitValue, 1 /* update */);
+    // this cannot fail since we know the element exists (nothing to allocate)
+    assert(cret == 0);
+    (void)cret;
+
+    // free the original value
+    free(value);
+    util_mutex_unlock(p->hTracker->splitMergeMutex);
+
+    return UMF_RESULT_SUCCESS;
+
+err:
+    util_mutex_unlock(p->hTracker->splitMergeMutex);
+err_lock:
+    free(splitValue);
+    return ret;
+}
+
+static umf_result_t trackingAllocationMerge(void *provider, void *lowPtr,
+                                            void *highPtr, size_t totalSize) {
+    umf_result_t ret = UMF_RESULT_ERROR_UNKNOWN;
+    umf_tracking_memory_provider_t *p =
+        (umf_tracking_memory_provider_t *)provider;
+
+    tracker_value_t *mergedValue =
+        (tracker_value_t *)malloc(sizeof(tracker_value_t));
+    mergedValue->pool = p->pool;
+    mergedValue->size = totalSize;
+
+    if (!mergedValue) {
+        return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
+    }
+
+    int r = util_mutex_lock(p->hTracker->splitMergeMutex);
+    if (r) {
+        goto err_lock;
+    }
+
+    tracker_value_t *leftValue =
+        (tracker_value_t *)critnib_get(p->hTracker->map, (uintptr_t)lowPtr);
+    if (!leftValue) {
+        fprintf(stderr, "tracking merge: no left value\n");
+        ret = UMF_RESULT_ERROR_INVALID_ARGUMENT;
+        goto err;
+    }
+    tracker_value_t *rightValue =
+        (tracker_value_t *)critnib_get(p->hTracker->map, (uintptr_t)highPtr);
+    if (!rightValue) {
+        fprintf(stderr, "tracking merge: no right value\n");
+        ret = UMF_RESULT_ERROR_INVALID_ARGUMENT;
+        goto err;
+    }
+    if (leftValue->pool != rightValue->pool) {
+        fprintf(stderr, "tracking merge: pool mismatch\n");
+        ret = UMF_RESULT_ERROR_INVALID_ARGUMENT;
+        goto err;
+    }
+    if (leftValue->size + rightValue->size != totalSize) {
+        fprintf(stderr, "tracking merge: leftValue->size + rightValue->size != "
+                        "totalSize\n");
+        ret = UMF_RESULT_ERROR_INVALID_ARGUMENT;
+        goto err;
+    }
+
+    ret = umfMemoryProviderAllocationMerge(p->hUpstream, lowPtr, highPtr,
+                                           totalSize);
+    if (ret != UMF_RESULT_SUCCESS) {
+        fprintf(stderr,
+                "tracking merge: umfMemoryProviderAllocationMerge failed\n");
+        goto err;
+    }
+
+    // We'll have duplicate entry for the range [highPtr, rightValue->size] but this is fine,
+    // the value is the same anyway and we forbid splitting/removing that range concurrently
+    int cret = critnib_insert(p->hTracker->map, (uintptr_t)lowPtr,
+                              (void *)mergedValue, 1 /* update */);
+    // this cannot fail since we know the element exists (nothing to allocate)
+    assert(cret == 0);
+    (void)cret;
+
+    // free old value that we just replaced with mergedValue
+    free(leftValue);
+
+    void *erasedRightValue =
+        critnib_remove(p->hTracker->map, (uintptr_t)highPtr);
+    assert(erasedRightValue == rightValue);
+
+    free(erasedRightValue);
+
+    util_mutex_unlock(p->hTracker->splitMergeMutex);
+
+    return UMF_RESULT_SUCCESS;
+
+err:
+    util_mutex_unlock(p->hTracker->splitMergeMutex);
+err_lock:
+    free(mergedValue);
+    return ret;
+}
+
 static umf_result_t trackingFree(void *hProvider, void *ptr, size_t size) {
     umf_result_t ret;
     umf_tracking_memory_provider_t *p =
@@ -214,7 +368,8 @@ umf_memory_provider_ops_t UMF_TRACKING_MEMORY_PROVIDER_OPS = {
     .purge_force = trackingPurgeForce,
     .purge_lazy = trackingPurgeLazy,
     .get_name = trackingName,
-};
+    .allocation_split = trackingAllocationSplit,
+    .allocation_merge = trackingAllocationMerge};
 
 umf_result_t umfTrackingMemoryProviderCreate(
     umf_memory_provider_handle_t hUpstream, umf_memory_pool_handle_t hPool,
diff --git a/src/provider/provider_tracking.h b/src/provider/provider_tracking.h
diff --git a/src/provider/provider_tracking_linux.c b/src/provider/provider_tracking_linux.c
diff --git a/src/provider/provider_tracking_windows.c b/src/provider/provider_tracking_windows.c

Original file line number	Diff line number	Diff line change
`@@ -121,7 +121,7 @@ static void tbb_raw_free_wrapper(intptr_t pool_id, void *ptr, size_t bytes) {`
`121`	`121`	`TLS_last_free_error = ret;`
`122`	`122`	`fprintf(`
`123`	`123`	`stderr,`
`124`		`- "Memory provider failed to free memory, addr = %p, size = %lu\n",`
	`124`	`+ "Memory provider failed to free memory, addr = %p, size = %zu\n",`
`125`	`125`	`ptr, bytes);`
`126`	`126`	`}`
`127`	`127`	`}`