From 996273c1153d1ece82ca5f54150e77949f7494bd Mon Sep 17 00:00:00 2001
From: Lukasz Dorau <lukasz.dorau@intel.com>
Date: Wed, 17 Jul 2024 16:33:07 +0200
Subject: [PATCH] Add disable_provider_free parameter to pool jemalloc

Add the disable_provider_free parameter to pool jemalloc.
It should be set to true if umfMemoryProviderFree()
should never be called.

This option will be needed by the future memory providers
that will not provide the free() op.

Signed-off-by: Lukasz Dorau <lukasz.dorau@intel.com>
---
 include/umf/pools/pool_jemalloc.h |  7 +++++++
 src/pool/pool_jemalloc.c          | 24 ++++++++++++++++++++++--
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/include/umf/pools/pool_jemalloc.h b/include/umf/pools/pool_jemalloc.h
index c30df6509..dfd75746b 100644
--- a/include/umf/pools/pool_jemalloc.h
+++ b/include/umf/pools/pool_jemalloc.h
@@ -14,8 +14,15 @@
 extern "C" {
 #endif
 
+#include <stdbool.h>
 #include <umf/memory_pool_ops.h>
 
+/// @brief Configuration of Jemalloc Pool
+typedef struct umf_jemalloc_pool_params_t {
+    /// Set to true if umfMemoryProviderFree() should never be called.
+    bool disable_provider_free;
+} umf_jemalloc_pool_params_t;
+
 umf_memory_pool_ops_t *umfJemallocPoolOps(void);
 
 #ifdef __cplusplus
diff --git a/src/pool/pool_jemalloc.c b/src/pool/pool_jemalloc.c
index ef72ca2f6..094ceeaf7 100644
--- a/src/pool/pool_jemalloc.c
+++ b/src/pool/pool_jemalloc.c
@@ -37,6 +37,8 @@
 typedef struct jemalloc_memory_pool_t {
     umf_memory_provider_handle_t provider;
     unsigned int arena_index; // index of jemalloc arena
+    // set to true if umfMemoryProviderFree() should never be called
+    bool disable_provider_free;
 } jemalloc_memory_pool_t;
 
 static __TLS umf_result_t TLS_last_allocation_error;
@@ -80,7 +82,9 @@ static void *arena_extent_alloc(extent_hooks_t *extent_hooks, void *new_addr,
     }
 
     if (new_addr != NULL && ptr != new_addr) {
-        umfMemoryProviderFree(pool->provider, ptr, size);
+        if (!pool->disable_provider_free) {
+            umfMemoryProviderFree(pool->provider, ptr, size);
+        }
         return NULL;
     }
 
@@ -114,6 +118,10 @@ static void arena_extent_destroy(extent_hooks_t *extent_hooks, void *addr,
 
     jemalloc_memory_pool_t *pool = get_pool_by_arena_index(arena_ind);
 
+    if (pool->disable_provider_free) {
+        return;
+    }
+
     umf_result_t ret;
     ret = umfMemoryProviderFree(pool->provider, addr, size);
     if (ret != UMF_RESULT_SUCCESS) {
@@ -136,6 +144,10 @@ static bool arena_extent_dalloc(extent_hooks_t *extent_hooks, void *addr,
 
     jemalloc_memory_pool_t *pool = get_pool_by_arena_index(arena_ind);
 
+    if (pool->disable_provider_free) {
+        return true; // opt-out from deallocation
+    }
+
     umf_result_t ret;
     ret = umfMemoryProviderFree(pool->provider, addr, size);
     if (ret != UMF_RESULT_SUCCESS) {
@@ -388,7 +400,9 @@ static umf_result_t op_initialize(umf_memory_provider_handle_t provider,
                                   void *params, void **out_pool) {
     assert(provider);
     assert(out_pool);
-    (void)params; // unused
+
+    umf_jemalloc_pool_params_t *je_params =
+        (umf_jemalloc_pool_params_t *)params;
 
     extent_hooks_t *pHooks = &arena_extent_hooks;
     size_t unsigned_size = sizeof(unsigned);
@@ -402,6 +416,12 @@ static umf_result_t op_initialize(umf_memory_provider_handle_t provider,
 
     pool->provider = provider;
 
+    if (je_params) {
+        pool->disable_provider_free = je_params->disable_provider_free;
+    } else {
+        pool->disable_provider_free = false;
+    }
+
     unsigned arena_index;
     err = je_mallctl("arenas.create", (void *)&arena_index, &unsigned_size,
                      NULL, 0);