Add trackerShrinkEntry() and trackerGrowEntry()

Add trackerShrinkEntry() that shrinks (or removes) an entry
in the tracker and returns the totalSize of the original entry.

Add trackerGrowEntry() that grows (or adds) an entry
in the tracker to its original size.

Signed-off-by: Lukasz Dorau <[email protected]>

Author: ldorau

src/provider/provider_fixed_memory.c

@@ -21,6 +21,7 @@
 #include "coarse.h"
 #include "libumf.h"
 #include "memory_pool_internal.h"
+#include "provider_tracking.h"
 #include "utils_common.h"
 #include "utils_concurrency.h"
 #include "utils_log.h"
@@ -114,6 +115,15 @@ static umf_result_t fixed_initialize(void *params, void **provider) {
             LOG_ERR("cannot get the tracking provider of the pool %p", pool);
             return ret;
         }
+
+        ret = trackerShrinkEntry(trackingProvider, in_params->ptr,
+                                 in_params->size, &ptr_orig_size);
+        if (ret != UMF_RESULT_SUCCESS) {
+            LOG_ERR(
+                "cannot shrink the allocation %p in the tracker to size %zu",
+                in_params->ptr, in_params->size);
+            return ret;
+        }
     }
 
     fixed_memory_provider_t *fixed_provider =
@@ -173,6 +183,20 @@ static umf_result_t fixed_initialize(void *params, void **provider) {
 static void fixed_finalize(void *provider) {
     fixed_memory_provider_t *fixed_provider = provider;
     coarse_delete(fixed_provider->coarse);
+
+    if (fixed_provider->trackingProvider &&
+        (fixed_provider->ptr_orig_size >= fixed_provider->size)) {
+        umf_result_t ret = trackerGrowEntry(
+            fixed_provider->trackingProvider, fixed_provider->base,
+            fixed_provider->size, fixed_provider->ptr_orig_size);
+        if (ret != UMF_RESULT_SUCCESS) {
+            LOG_ERR("cannot grow the allocation %p in the tracker (from size "
+                    "%zu to size %zu)",
+                    fixed_provider->base, fixed_provider->size,
+                    fixed_provider->ptr_orig_size);
+        }
+    }
+
     umf_ba_global_free(fixed_provider);
 }
 

src/provider/provider_tracking.c

@@ -267,6 +267,61 @@ static umf_result_t trackingAllocationSplit(void *hProvider, void *ptr,
     return ret;
 }
 
+// shrink (or remove) an entry in the tracker and return the totalSize of the original entry
+umf_result_t trackerShrinkEntry(void *hProvider, void *ptr, size_t shrinkSize,
+                                size_t *totalSize) {
+    umf_result_t ret = UMF_RESULT_ERROR_UNKNOWN;
+    umf_tracking_memory_provider_t *provider =
+        (umf_tracking_memory_provider_t *)hProvider;
+
+    int r = utils_mutex_lock(&provider->hTracker->splitMergeMutex);
+    if (r) {
+        return UMF_RESULT_ERROR_UNKNOWN;
+    }
+
+    tracker_alloc_info_t *value = (tracker_alloc_info_t *)critnib_get(
+        provider->hTracker->alloc_segments_map, (uintptr_t)ptr);
+    if (!value) {
+        LOG_ERR("region for shrinking is not found in the tracker");
+        ret = UMF_RESULT_ERROR_INVALID_ARGUMENT;
+        goto err;
+    }
+    if (shrinkSize > value->size) {
+        LOG_ERR("requested size %zu to shrink exceeds the tracked size %zu",
+                shrinkSize, value->size);
+        ret = UMF_RESULT_ERROR_INVALID_ARGUMENT;
+        goto err;
+    }
+
+    if (shrinkSize < value->size) {
+        void *highPtr = (void *)(((uintptr_t)ptr) + shrinkSize);
+        size_t secondSize = value->size - shrinkSize;
+        ret = umfMemoryTrackerAdd(provider->hTracker, provider->pool, highPtr,
+                                  secondSize);
+        if (ret != UMF_RESULT_SUCCESS) {
+            LOG_ERR("failed to add the new shrunk region to the tracker, ptr = "
+                    "%p, size = %zu, ret = %d",
+                    highPtr, secondSize, ret);
+            goto err;
+        }
+    }
+
+    *totalSize = value->size;
+
+    void *erasedValue =
+        critnib_remove(provider->hTracker->alloc_segments_map, (uintptr_t)ptr);
+    assert(erasedValue == value);
+    umf_ba_free(provider->hTracker->alloc_info_allocator, erasedValue);
+
+    utils_mutex_unlock(&provider->hTracker->splitMergeMutex);
+
+    return UMF_RESULT_SUCCESS;
+
+err:
+    utils_mutex_unlock(&provider->hTracker->splitMergeMutex);
+    return ret;
+}
+
 static umf_result_t trackingAllocationMerge(void *hProvider, void *lowPtr,
                                             void *highPtr, size_t totalSize) {
     umf_result_t ret = UMF_RESULT_ERROR_UNKNOWN;
@@ -353,6 +408,68 @@ static umf_result_t trackingAllocationMerge(void *hProvider, void *lowPtr,
     return ret;
 }
 
+// grow (or add) an entry in the tracker to its original size
+umf_result_t trackerGrowEntry(void *hProvider, void *ptr, size_t growSize,
+                              size_t origSize) {
+    umf_result_t ret = UMF_RESULT_ERROR_UNKNOWN;
+    umf_tracking_memory_provider_t *provider =
+        (umf_tracking_memory_provider_t *)hProvider;
+
+    if (growSize > origSize) {
+        LOG_ERR("Invalid grow size %zu (larger than the original size %zu)",
+                growSize, origSize);
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    int r = utils_mutex_lock(&provider->hTracker->splitMergeMutex);
+    if (r) {
+        return UMF_RESULT_ERROR_UNKNOWN;
+    }
+
+    void *highPtr = (void *)(((uintptr_t)ptr) + growSize);
+    tracker_alloc_info_t *highValue = NULL;
+
+    if (growSize < origSize) {
+        highValue = (tracker_alloc_info_t *)critnib_get(
+            provider->hTracker->alloc_segments_map, (uintptr_t)highPtr);
+        if (!highValue) {
+            LOG_ERR("cannot find the tracker entry to grow %p", highPtr);
+            ret = UMF_RESULT_ERROR_INVALID_ARGUMENT;
+            goto err;
+        }
+        if (growSize + highValue->size != origSize) {
+            LOG_ERR("Grow size plus the current size does not equal the "
+                    "original size");
+            ret = UMF_RESULT_ERROR_INVALID_ARGUMENT;
+            goto err;
+        }
+    }
+
+    ret =
+        umfMemoryTrackerAdd(provider->hTracker, provider->pool, ptr, origSize);
+    if (ret != UMF_RESULT_SUCCESS) {
+        LOG_ERR("failed to add the new grown region to the tracker, ptr = %p, "
+                "size = %zu, ret = %d",
+                ptr, origSize, ret);
+        goto err;
+    }
+
+    if (growSize < origSize) {
+        void *erasedhighValue = critnib_remove(
+            provider->hTracker->alloc_segments_map, (uintptr_t)highPtr);
+        assert(erasedhighValue == highValue);
+        umf_ba_free(provider->hTracker->alloc_info_allocator, erasedhighValue);
+    }
+
+    utils_mutex_unlock(&provider->hTracker->splitMergeMutex);
+
+    return UMF_RESULT_SUCCESS;
+
+err:
+    utils_mutex_unlock(&provider->hTracker->splitMergeMutex);
+    return ret;
+}
+
 static umf_result_t trackingFree(void *hProvider, void *ptr, size_t size) {
     umf_result_t ret;
     umf_result_t ret_remove = UMF_RESULT_ERROR_UNKNOWN;

src/provider/provider_tracking.h

@@ -55,6 +55,12 @@ void umfTrackingMemoryProviderGetUpstreamProvider(
     umf_memory_provider_handle_t hTrackingProvider,
     umf_memory_provider_handle_t *hUpstream);
 
+umf_result_t trackerShrinkEntry(void *hProvider, void *ptr, size_t shrinkSize,
+                                size_t *totalSize);
+
+umf_result_t trackerGrowEntry(void *hProvider, void *ptr, size_t growSize,
+                              size_t origSize);
+
 #ifdef __cplusplus
 }
 #endif