x

Author: bratpiorka

include/umf/memory_pool.h

@@ -190,6 +190,15 @@ umf_result_t umfPoolSetTag(umf_memory_pool_handle_t hPool, void *tag,
 /// @return UMF_RESULT_SUCCESS on success.
 umf_result_t umfPoolGetTag(umf_memory_pool_handle_t hPool, void **tag);
 
+///
+/// @brief Trims memory pool to keep at least \p minBytesToKeep bytes of memory
+///        if possible.
+/// @param hPool specified memory pool
+/// @param minBytesToKeep minimum number of bytes to keep in the pool
+/// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
+umf_result_t umfPoolTrimMemory(umf_memory_pool_handle_t hPool,
+                               size_t minBytesToKeep);
+
 #ifdef __cplusplus
 }
 #endif

include/umf/memory_pool_ops.h

@@ -161,6 +161,16 @@ typedef struct umf_memory_pool_ops_t {
     /// @return A constant character string representing the pool's name.
     ///
     const char *(*ext_get_name)(void *pool);
+
+    ///
+    /// @brief Trims memory of the pool, removing resources that are not needed
+    ///        to keep the pool operational.
+    /// @param pool pointer to the memory pool
+    /// @param minBytesToKeep minimum number of bytes to keep in the pool if
+    ///        possible.
+    /// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
+    ///
+    umf_result_t (*ext_trim_memory)(void *pool, size_t minBytesToKeep);
 } umf_memory_pool_ops_t;
 
 #ifdef __cplusplus

src/libumf.def

@@ -144,3 +144,4 @@ EXPORTS
     umfJemallocPoolParamsDestroy
     umfJemallocPoolParamsSetNumArenas
     umfPoolGetName
+    umfPoolTrimMemory

src/libumf.map

@@ -144,4 +144,5 @@ UMF_0.12 {
         umfJemallocPoolParamsDestroy;
         umfJemallocPoolParamsSetNumArenas;
         umfPoolGetName;
+        umfPoolTrimMemory;
 } UMF_0.11;

src/memory_pool.c

@@ -240,6 +240,10 @@ static umf_result_t umfPoolCreateInternal(const umf_memory_pool_ops_t *ops,
 }
 
 umf_result_t umfPoolDestroy(umf_memory_pool_handle_t hPool) {
+    if (hPool == NULL) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
     if (umf_ba_global_is_destroyed()) {
         return UMF_RESULT_ERROR_UNKNOWN;
     }
@@ -417,3 +421,13 @@ umf_result_t umfPoolGetTag(umf_memory_pool_handle_t hPool, void **tag) {
     utils_mutex_unlock(&hPool->lock);
     return UMF_RESULT_SUCCESS;
 }
+
+umf_result_t umfPoolTrimMemory(umf_memory_pool_handle_t hPool,
+                               size_t minBytesToKeep) {
+    UMF_CHECK((hPool != NULL), UMF_RESULT_ERROR_INVALID_ARGUMENT);
+    if (hPool->ops.ext_trim_memory == NULL) {
+        return UMF_RESULT_ERROR_NOT_SUPPORTED;
+    }
+
+    return hPool->ops.ext_trim_memory(hPool->pool_priv, minBytesToKeep);
+}

src/pool/pool_disjoint.c

@@ -1031,6 +1031,47 @@ const char *disjoint_pool_get_name(void *pool) {
     return hPool->params.name;
 }
 
+umf_result_t disjoint_pool_trim_memory(void *pool, size_t minBytesToKeep) {
+    disjoint_pool_t *hPool = (disjoint_pool_t *)pool;
+    if (hPool == NULL) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    for (size_t i = 0; i < hPool->buckets_num; i++) {
+        bucket_t *bucket = hPool->buckets[i];
+        utils_mutex_lock(&bucket->bucket_lock);
+
+        int skip = (int)minBytesToKeep;
+
+        // remove empty slabs from the pool
+        slab_list_item_t *it = NULL, *tmp = NULL;
+        LL_FOREACH_SAFE(bucket->available_slabs, it, tmp) {
+            slab_t *slab = it->val;
+            if (slab->num_chunks_allocated == 0) {
+                // skip first minBytesToKeep bytes from each bucket
+                if (skip > 0) {
+                    skip -= slab->slab_size;
+                    continue;
+                }
+
+                // remove slab
+                pool_unregister_slab(hPool, slab);
+                DL_DELETE(bucket->available_slabs, it);
+                assert(bucket->available_slabs_num > 0);
+                bucket->available_slabs_num--;
+                destroy_slab(slab);
+
+                // update stats
+                bucket_update_stats(bucket, 0, -1);
+            }
+        }
+
+        utils_mutex_unlock(&bucket->bucket_lock);
+    }
+
+    return UMF_RESULT_SUCCESS;
+}
+
 static umf_memory_pool_ops_t UMF_DISJOINT_POOL_OPS = {
     .version = UMF_VERSION_CURRENT,
     .initialize = disjoint_pool_initialize,
@@ -1044,6 +1085,7 @@ static umf_memory_pool_ops_t UMF_DISJOINT_POOL_OPS = {
     .get_last_allocation_error = disjoint_pool_get_last_allocation_error,
     .ext_get_name = disjoint_pool_get_name,
     .ext_ctl = disjoint_pool_ctl,
+    .ext_trim_memory = disjoint_pool_trim_memory,
 };
 
 const umf_memory_pool_ops_t *umfDisjointPoolOps(void) {

src/pool/pool_jemalloc.c

@@ -552,6 +552,23 @@ static umf_result_t op_get_last_allocation_error(void *pool) {
     return TLS_last_allocation_error;
 }
 
+static umf_result_t op_trim_memory(void *pool, size_t minBytesToKeep) {
+    (void)minBytesToKeep; // unused - TODO?
+
+    jemalloc_memory_pool_t *je_pool = (jemalloc_memory_pool_t *)pool;
+    for (size_t i = 0; i < je_pool->n_arenas; i++) {
+        char cmd[64];
+        unsigned arena = je_pool->arena_index[i];
+        snprintf(cmd, sizeof(cmd), "arena.%u.purge", arena);
+        if (je_mallctl(cmd, NULL, NULL, NULL, 0)) {
+            LOG_ERR("Could not purge jemalloc arena %u", arena);
+            return UMF_RESULT_ERROR_UNKNOWN;
+        }
+    }
+
+    return UMF_RESULT_SUCCESS;
+}
+
 static umf_memory_pool_ops_t UMF_JEMALLOC_POOL_OPS = {
     .version = UMF_POOL_OPS_VERSION_CURRENT,
     .initialize = op_initialize,
@@ -563,6 +580,7 @@ static umf_memory_pool_ops_t UMF_JEMALLOC_POOL_OPS = {
     .malloc_usable_size = op_malloc_usable_size,
     .free = op_free,
     .get_last_allocation_error = op_get_last_allocation_error,
+    .ext_trim_memory = op_trim_memory,
 };
 
 const umf_memory_pool_ops_t *umfJemallocPoolOps(void) {

src/pool/pool_proxy.c

@@ -125,6 +125,18 @@ static umf_result_t proxy_get_last_allocation_error(void *pool) {
     return TLS_last_allocation_error;
 }
 
+static const char *proxy_get_name(void *pool) {
+    (void)pool; // not used
+    return "proxy";
+}
+
+static umf_result_t proxy_trim_memory(void *pool, size_t minBytesToKeep) {
+    (void)pool;
+    (void)minBytesToKeep;
+
+    return UMF_RESULT_SUCCESS;
+}
+
 static umf_memory_pool_ops_t UMF_PROXY_POOL_OPS = {
     .version = UMF_POOL_OPS_VERSION_CURRENT,
     .initialize = proxy_pool_initialize,
@@ -135,7 +147,10 @@ static umf_memory_pool_ops_t UMF_PROXY_POOL_OPS = {
     .aligned_malloc = proxy_aligned_malloc,
     .malloc_usable_size = proxy_malloc_usable_size,
     .free = proxy_free,
-    .get_last_allocation_error = proxy_get_last_allocation_error};
+    .get_last_allocation_error = proxy_get_last_allocation_error,
+    .ext_get_name = proxy_get_name,
+    .ext_trim_memory = proxy_trim_memory,
+};
 
 const umf_memory_pool_ops_t *umfProxyPoolOps(void) {
     return &UMF_PROXY_POOL_OPS;

src/pool/pool_scalable.c

@@ -13,15 +13,15 @@
 #include <stdio.h>
 #include <string.h>
 
-#include <ctl/ctl.h>
-#include <memory_pool_internal.h>
 #include <umf/memory_pool.h>
 #include <umf/memory_pool_ops.h>
 #include <umf/memory_provider.h>
 #include <umf/pools/pool_scalable.h>
 
 #include "base_alloc_global.h"
+#include "ctl/ctl.h"
 #include "libumf.h"
+#include "memory_pool_internal.h"
 #include "pool_scalable_internal.h"
 #include "utils_common.h"
 #include "utils_concurrency.h"
@@ -60,6 +60,7 @@ typedef struct tbb_callbacks_t {
     bool (*pool_destroy)(void *);
     void *(*pool_identify)(void *object);
     size_t (*pool_msize)(void *, void *);
+    int (*pool_allocation_command)(int, void *);
 #ifdef _WIN32
     HMODULE lib_handle;
 #else
@@ -431,6 +432,14 @@ static const char *scalable_get_name(void *pool) {
     return "scalable";
 }
 
+static umf_result_t scalable_trim_memory(void *pool, size_t minBytesToKeep) {
+    (void)pool;           // unused
+    (void)minBytesToKeep; // unused
+
+    //scalable_allocation_command?
+    return UMF_RESULT_SUCCESS;
+}
+
 static umf_memory_pool_ops_t UMF_SCALABLE_POOL_OPS = {
     .version = UMF_POOL_OPS_VERSION_CURRENT,
     .initialize = tbb_pool_initialize,
@@ -444,6 +453,7 @@ static umf_memory_pool_ops_t UMF_SCALABLE_POOL_OPS = {
     .get_last_allocation_error = tbb_get_last_allocation_error,
     .ext_ctl = pool_ctl,
     .ext_get_name = scalable_get_name,
+    .ext_trim_memory = scalable_trim_memory,
 };
 
 const umf_memory_pool_ops_t *umfScalablePoolOps(void) {

test/common/pool.hpp

@@ -117,6 +117,10 @@ typedef struct pool_base_t {
     umf_result_t get_last_allocation_error() noexcept {
         return UMF_RESULT_SUCCESS;
     }
+    umf_result_t
+    ext_trim_memory([[maybe_unused]] size_t minBytesToKeep) noexcept {
+        return UMF_RESULT_SUCCESS;
+    }
 } pool_base_t;
 
 struct malloc_pool : public pool_base_t {

test/memoryPoolAPI.cpp

@@ -13,7 +13,9 @@
 
 #include <umf/memory_provider.h>
 #include <umf/pools/pool_disjoint.h>
+#include <umf/pools/pool_jemalloc.h>
 #include <umf/pools/pool_proxy.h>
+#include <umf/pools/pool_scalable.h>
 
 #ifdef UMF_PROXY_LIB_ENABLED
 #include <umf/proxy_lib_new_delete.h>
@@ -303,7 +305,12 @@ INSTANTIATE_TEST_SUITE_P(
                             &BA_GLOBAL_PROVIDER_OPS, nullptr, nullptr},
         poolCreateExtParams{umfDisjointPoolOps(), defaultDisjointPoolConfig,
                             defaultDisjointPoolConfigDestroy,
-                            &BA_GLOBAL_PROVIDER_OPS, nullptr, nullptr}));
+                            &BA_GLOBAL_PROVIDER_OPS, nullptr, nullptr}),
+    //  poolCreateExtParams{umfScalablePoolOps(), nullptr, nullptr,
+    //                      &BA_GLOBAL_PROVIDER_OPS, nullptr, nullptr},
+    //  poolCreateExtParams{umfJemallocPoolOps(), nullptr, nullptr,
+    //                     &BA_GLOBAL_PROVIDER_OPS, nullptr, nullptr}),
+    poolCreateExtParamsNameGen);
 
 INSTANTIATE_TEST_SUITE_P(mallocMultiPoolTest, umfMultiPoolTest,
                          ::testing::Values(poolCreateExtParams{
@@ -532,4 +539,11 @@ INSTANTIATE_TEST_SUITE_P(
         umf_test::withGeneratedArgs(umfPoolCalloc),
         umf_test::withGeneratedArgs(umfPoolRealloc),
         umf_test::withGeneratedArgs(umfPoolMallocUsableSize),
-        umf_test::withGeneratedArgs(umfPoolGetLastAllocationError)));
+        umf_test::withGeneratedArgs(umfPoolGetLastAllocationError),
+        umf_test::withGeneratedArgs(umfPoolByPtr),
+        umf_test::withGeneratedArgs(umfPoolGetMemoryProvider),
+        umf_test::withGeneratedArgs(umfPoolGetTag),
+        umf_test::withGeneratedArgs(umfPoolSetTag),
+        umf_test::withGeneratedArgs(umfPoolCreate),
+        umf_test::withGeneratedArgs(umfPoolDestroy),
+        umf_test::withGeneratedArgs(umfPoolTrimMemory)));

test/poolFixtures.hpp

@@ -33,6 +33,24 @@ using poolCreateExtParams =
                pfnPoolParamsDestroy, const umf_memory_provider_ops_t *,
                pfnProviderParamsCreate, pfnProviderParamsDestroy>;
 
+std::string poolCreateExtParamsNameGen(
+    const testing::TestParamInfo<poolCreateExtParams> param) {
+
+    const umf_memory_pool_ops_t *pool_ops = std::get<0>(param.param);
+    const umf_memory_provider_ops_t *provider_ops = std::get<3>(param.param);
+
+    std::string poolName =
+        pool_ops->ext_get_name ? pool_ops->ext_get_name(NULL) : "unknown_pool";
+    std::string providerName = provider_ops->get_name
+                                   ? provider_ops->get_name(NULL)
+                                   : "unknown_provider";
+
+    std::string poolParams =
+        std::get<1>(param.param) ? std::string("_w_params") + std::to_string(param.index) : std::string("");
+
+    return poolName + poolParams + "_" + providerName;
+}
+
 umf_test::pool_unique_handle_t poolCreateExtUnique(poolCreateExtParams params) {
     auto [pool_ops, poolParamsCreate, poolParamsDestroy, provider_ops,
           providerParamsCreate, providerParamsDestroy] = params;
@@ -402,6 +420,23 @@ TEST_P(umfPoolTest, multiThreadedMallocFreeRandomSizes) {
     }
 }
 
+TEST_P(umfPoolTest, trimMemory) {
+    constexpr size_t size = 1024;
+
+    umf_memory_pool_handle_t hPool = pool.get();
+    char *ptr = (char *)umfPoolMalloc(hPool, size);
+    ASSERT_NE(ptr, nullptr);
+
+    umf_result_t ret = umfFree(ptr);
+    ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
+
+    // Call to umfPoolTrimMemory should purge the whole memory pool
+    ret = umfPoolTrimMemory(hPool, 0);
+    ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
+
+    // TODO add CTL to check that the memory was actually purged
+}
+
 TEST_P(umfMemTest, outOfMem) {
     static constexpr size_t allocSize = 4096;
     auto hPool = pool.get();