Implement umfMemoryProviderAllocation[Split/Merge]

igchor · ldorau · commit 10d56ce58def · 2024-01-17T22:27:42.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 *leftPtr, void *rightPtr,
+    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 *leftPtr,
+                                     void *rightPtr, 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 *leftPtr, void *rightPtr,
+                                 size_t totalSize) {
+    if ((uintptr_t)leftPtr >= (uintptr_t)rightPtr) {
+        fprintf(
+            stderr,
+            "umfMemoryProviderAllocationMerge failed. check failed: %p >= %p\n",
+            leftPtr, rightPtr);
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    if ((uintptr_t)rightPtr - (uintptr_t)leftPtr > totalSize) {
+        fprintf(stderr,
+                "umfMemoryProviderAllocationMerge failed. check failed: %p - "
+                "%p > %zu\n",
+                rightPtr, leftPtr, totalSize);
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    umf_result_t res = hProvider->ops.allocation_merge(
+        hProvider->provider_priv, leftPtr, rightPtr, 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
@@ -122,7 +122,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 *leftPtr,
+                                        void *rightPtr, size_t totalSize) {
+    (void)provider;
+    (void)leftPtr;
+    (void)rightPtr;
+    (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
@@ -119,6 +119,134 @@ 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_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;
+
+    tracker_value_t *value =
+        (tracker_value_t *)critnib_get((critnib *)p->hTracker, (uintptr_t)ptr);
+    if (!value) {
+        free(splitValue);
+        fprintf(stderr, "tracking split: no such value\n");
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+    if (value->size != totalSize) {
+        free(splitValue);
+        fprintf(stderr, "tracking split: %zu != %zu\n", value->size, totalSize);
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    umf_result_t ret = umfMemoryProviderAllocationSplit(p->hUpstream, ptr,
+                                                        totalSize, leftSize);
+    if (ret != UMF_RESULT_SUCCESS) {
+        free(splitValue);
+        fprintf(stderr,
+                "tracking split: umfMemoryProviderAllocationSplit failed\n");
+        return ret;
+    }
+
+    void *rightPtr = (void *)(((uintptr_t)ptr) + leftSize);
+    size_t rightSize = totalSize - leftSize;
+
+    // We'll have duplicate entry for the range [rightPtr, 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, rightPtr, rightSize);
+    if (ret != UMF_RESULT_SUCCESS) {
+        free(splitValue);
+        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.
+        return ret;
+    }
+
+    int cret = critnib_insert((critnib *)p->hTracker, (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);
+
+    return UMF_RESULT_SUCCESS;
+}
+
+static umf_result_t trackingAllocationMerge(void *provider, void *leftPtr,
+                                            void *rightPtr, size_t totalSize) {
+    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;
+    }
+
+    tracker_value_t *leftValue = (tracker_value_t *)critnib_get(
+        (critnib *)p->hTracker, (uintptr_t)leftPtr);
+    if (!leftValue) {
+        free(mergedValue);
+        fprintf(stderr, "tracking merge: no left value\n");
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+    tracker_value_t *rightValue = (tracker_value_t *)critnib_get(
+        (critnib *)p->hTracker, (uintptr_t)rightPtr);
+    if (!rightValue) {
+        free(mergedValue);
+        fprintf(stderr, "tracking merge: no right value\n");
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+    if (leftValue->pool != rightValue->pool) {
+        free(mergedValue);
+        fprintf(stderr, "tracking merge: pool mismatch\n");
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+    if (leftValue->size + rightValue->size != totalSize) {
+        // TODO: should we allow leftValue->size + rightValue->size > totalSize?
+        fprintf(stderr, "tracking merge: leftValue->size + rightValue->size != "
+                        "totalSize\n");
+        free(mergedValue);
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    umf_result_t ret = umfMemoryProviderAllocationMerge(p->hUpstream, leftPtr,
+                                                        rightPtr, totalSize);
+    if (ret != UMF_RESULT_SUCCESS) {
+        free(mergedValue);
+        fprintf(stderr,
+                "tracking merge: umfMemoryProviderAllocationMerge failed\n");
+        return ret;
+    }
+
+    // We'll have duplicate entry for the range [rightPtr, rightValue->size] but this is fine,
+    // the value is the same anyway and we forbid splitting/removing that range concurrently
+    int cret = critnib_insert((critnib *)p->hTracker, (uintptr_t)leftPtr,
+                              (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((critnib *)p->hTracker, (uintptr_t)rightPtr);
+    assert(erasedRightValue == rightValue);
+    free(erasedRightValue);
+
+    return UMF_RESULT_SUCCESS;
+}
+
 static umf_result_t trackingFree(void *hProvider, void *ptr, size_t size) {
     umf_result_t ret;
     umf_tracking_memory_provider_t *p =
@@ -216,7 +344,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,

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