add post-initialize function to pools and providers

Split between initialize and post-initialize function is necessary
for properly handling CTL defaults.

Author: KFilipek

include/umf/memory_pool_ops.h

@@ -22,7 +22,7 @@ extern "C" {
 /// @brief Version of the Memory Pool ops structure.
 /// NOTE: This is equal to the latest UMF version, in which the ops structure
 /// has been modified.
-#define UMF_POOL_OPS_VERSION_CURRENT UMF_MAKE_VERSION(1, 0)
+#define UMF_POOL_OPS_VERSION_CURRENT UMF_MAKE_VERSION(1, 1)
 
 ///
 /// @brief This structure comprises function pointers used by corresponding umfPool*
@@ -166,6 +166,18 @@ typedef struct umf_memory_pool_ops_t {
                             const char *name, void *arg, size_t size,
                             umf_ctl_query_type_t queryType, va_list args);
 
+    ///
+    /// @brief Post-initializes memory pool.
+    /// @param provider memory provider that will be used for coarse-grain allocations.
+    ///        Should contain at least one memory provider.
+    /// @param numProvider number of elements in the providers array
+    /// @param params pool-specific params, or NULL for defaults
+    /// @param pool [out] returns pointer to the pool
+    /// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
+    ///
+    umf_result_t (*ext_post_initialize)(umf_memory_provider_handle_t provider,
+                                        const void *params, void *pool);
+
 } umf_memory_pool_ops_t;
 
 #ifdef __cplusplus

include/umf/memory_provider_ops.h

@@ -21,7 +21,7 @@ extern "C" {
 /// @brief Version of the Memory Provider ops structure.
 /// NOTE: This is equal to the latest UMF version, in which the ops structure
 /// has been modified.
-#define UMF_PROVIDER_OPS_VERSION_CURRENT UMF_MAKE_VERSION(1, 0)
+#define UMF_PROVIDER_OPS_VERSION_CURRENT UMF_MAKE_VERSION(1, 1)
 
 ///
 /// @brief This structure comprises function pointers used by corresponding
@@ -278,6 +278,14 @@ typedef struct umf_memory_provider_ops_t {
                             const char *name, void *arg, size_t size,
                             umf_ctl_query_type_t queryType, va_list args);
 
+    ///
+    /// @brief Post-initializes memory provider.
+    /// @param params provider-specific params, or NULL for defaults
+    /// @param provider pointer to the provider
+    /// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
+    ///
+    umf_result_t (*ext_post_initialize)(const void *params, void *provider);
+
 } umf_memory_provider_ops_t;
 
 #ifdef __cplusplus

src/memory_pool.c

@@ -200,11 +200,24 @@ static umf_result_t umfPoolCreateInternal(const umf_memory_pool_ops_t *ops,
     }
 
     umf_result_t ret = UMF_RESULT_SUCCESS;
-
     if (ops->version != UMF_POOL_OPS_VERSION_CURRENT) {
         LOG_WARN("Memory Pool ops version \"%d\" is different than the current "
                  "version \"%d\"",
                  ops->version, UMF_POOL_OPS_VERSION_CURRENT);
+
+        // Create a new ops compatible structure with the current version
+        umf_memory_pool_ops_t compatible_ops;
+        memset(&compatible_ops, 0, sizeof(compatible_ops));
+        if (UMF_MINOR_VERSION(ops->version) == 0) {
+            memcpy(&compatible_ops, ops,
+                   offsetof(umf_memory_pool_ops_t, ext_post_initialize));
+        } else {
+            LOG_ERR("Memory Pool ops unknown version, which \"%d\" is not "
+                    "supported",
+                    ops->version);
+            return UMF_RESULT_ERROR_NOT_SUPPORTED;
+        }
+        ops = &compatible_ops;
     }
 
     umf_memory_pool_handle_t pool =
@@ -261,7 +274,16 @@ static umf_result_t umfPoolCreateInternal(const umf_memory_pool_ops_t *ops,
         }
     }
 
+    if (ops->ext_post_initialize != NULL) {
+        ret = ops->ext_post_initialize(pool->provider, params, pool->pool_priv);
+        if (ret != UMF_RESULT_SUCCESS) {
+            LOG_ERR("Failed to post-initialize pool");
+            goto err_pool_init;
+        }
+    }
+
     *hPool = pool;
+
     LOG_INFO("Memory pool created: %p", (void *)pool);
     return UMF_RESULT_SUCCESS;
 

src/memory_provider.c

@@ -119,6 +119,13 @@ static umf_result_t umfDefaultCloseIPCHandle(void *provider, void *ptr,
     return UMF_RESULT_ERROR_NOT_SUPPORTED;
 }
 
+static umf_result_t umfDefaultPostInitialize(const void *params,
+                                             void *provider) {
+    (void)params;
+    (void)provider;
+    return UMF_RESULT_SUCCESS;
+}
+
 static umf_result_t
 umfDefaultCtlHandle(void *provider, umf_ctl_query_source_t operationType,
                     const char *name, void *arg, size_t size,
@@ -153,6 +160,10 @@ void assignOpsExtDefaults(umf_memory_provider_ops_t *ops) {
     if (!ops->ext_ctl) {
         ops->ext_ctl = umfDefaultCtlHandle;
     }
+
+    if (!ops->ext_post_initialize) {
+        ops->ext_post_initialize = umfDefaultPostInitialize;
+    }
 }
 
 void assignOpsIpcDefaults(umf_memory_provider_ops_t *ops) {
@@ -250,6 +261,16 @@ umf_result_t umfMemoryProviderCreate(const umf_memory_provider_ops_t *ops,
 
     provider->provider_priv = provider_priv;
 
+    if (provider->ops.ext_post_initialize != NULL) {
+        ret = provider->ops.ext_post_initialize(params, provider_priv);
+        if (ret != UMF_RESULT_SUCCESS) {
+            LOG_ERR("Failed to post-initialize provider");
+            provider->ops.finalize(provider_priv);
+            umf_ba_global_free(provider);
+            return ret;
+        }
+    }
+
     *hProvider = provider;
 
     return UMF_RESULT_SUCCESS;

src/pool/pool_disjoint.c

@@ -35,39 +35,39 @@ static UTIL_ONCE_FLAG ctl_initialized = UTIL_ONCE_FLAG_INIT;
 
 // Disable name ctl for 1.0 release
 #if 0
-static umf_result_t CTL_READ_HANDLER(name)(void *ctx,
-                                           umf_ctl_query_source_t source,
-                                           void *arg, size_t size,
-                                           umf_ctl_index_utlist_t *indexes) {
+static umf_result_t CTL_READ_HANDLER(name)(void* ctx,
+    umf_ctl_query_source_t source,
+    void* arg, size_t size,
+    umf_ctl_index_utlist_t* indexes) {
     (void)source, (void)indexes;
 
-    disjoint_pool_t *pool = (disjoint_pool_t *)ctx;
+    disjoint_pool_t* pool = (disjoint_pool_t*)ctx;
 
     if (arg == NULL) {
         return UMF_RESULT_ERROR_INVALID_ARGUMENT;
     }
 
     if (size > 0) {
-        strncpy((char *)arg, pool->params.name, size - 1);
-        ((char *)arg)[size - 1] = '\0';
+        strncpy((char*)arg, pool->params.name, size - 1);
+        ((char*)arg)[size - 1] = '\0';
     }
 
     return UMF_RESULT_SUCCESS;
 }
 
 static const struct ctl_argument CTL_ARG(name) = CTL_ARG_STRING(255);
 
-static umf_result_t CTL_WRITE_HANDLER(name)(void *ctx,
-                                            umf_ctl_query_source_t source,
-                                            void *arg, size_t size,
-                                            umf_ctl_index_utlist_t *indexes) {
+static umf_result_t CTL_WRITE_HANDLER(name)(void* ctx,
+    umf_ctl_query_source_t source,
+    void* arg, size_t size,
+    umf_ctl_index_utlist_t* indexes) {
     (void)source, (void)indexes, (void)size;
-    disjoint_pool_t *pool = (disjoint_pool_t *)ctx;
+    disjoint_pool_t* pool = (disjoint_pool_t*)ctx;
     if (arg == NULL) {
         return UMF_RESULT_ERROR_INVALID_ARGUMENT;
     }
 
-    strncpy(pool->params.name, (char *)arg, sizeof(pool->params.name) - 1);
+    strncpy(pool->params.name, (char*)arg, sizeof(pool->params.name) - 1);
     pool->params.name[sizeof(pool->params.name) - 1] = '\0';
 
     return UMF_RESULT_SUCCESS;
@@ -758,9 +758,21 @@ umf_result_t disjoint_pool_initialize(umf_memory_provider_handle_t provider,
     disjoint_pool->provider = provider;
     disjoint_pool->params = *dp_params;
 
+    *ppPool = (void *)disjoint_pool;
+
+    return UMF_RESULT_SUCCESS;
+}
+
+umf_result_t
+disjoint_pool_post_initialize(umf_memory_provider_handle_t provider,
+                              const void *params, void *ppPool) {
+    (void)params;
+    disjoint_pool_t *disjoint_pool = (disjoint_pool_t *)ppPool;
+
     disjoint_pool->known_slabs = critnib_new(free_slab, NULL);
     if (disjoint_pool->known_slabs == NULL) {
-        goto err_free_disjoint_pool;
+        umf_ba_global_free(disjoint_pool);
+        return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
     }
 
     // Generate buckets sized such as: 64, 96, 128, 192, ..., CutOff.
@@ -821,8 +833,6 @@ umf_result_t disjoint_pool_initialize(umf_memory_provider_handle_t provider,
         disjoint_pool->provider_min_page_size = 0;
     }
 
-    *ppPool = (void *)disjoint_pool;
-
     return UMF_RESULT_SUCCESS;
 
 err_free_buckets:
@@ -838,10 +848,6 @@ umf_result_t disjoint_pool_initialize(umf_memory_provider_handle_t provider,
 
 err_free_known_slabs:
     critnib_delete(disjoint_pool->known_slabs);
-
-err_free_disjoint_pool:
-    umf_ba_global_free(disjoint_pool);
-
     return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
 }
 
@@ -1146,7 +1152,7 @@ static umf_memory_pool_ops_t UMF_DISJOINT_POOL_OPS = {
     .get_last_allocation_error = disjoint_pool_get_last_allocation_error,
     .get_name = disjoint_pool_get_name,
     .ext_ctl = disjoint_pool_ctl,
-};
+    .ext_post_initialize = disjoint_pool_post_initialize};
 
 const umf_memory_pool_ops_t *umfDisjointPoolOps(void) {
     return &UMF_DISJOINT_POOL_OPS;

src/pool/pool_jemalloc.c

@@ -68,7 +68,7 @@ typedef struct jemalloc_memory_pool_t {
     umf_memory_provider_handle_t provider;
     size_t n_arenas;
     char name[64];
-    unsigned int arena_index[];
+    unsigned int *arena_index;
 } jemalloc_memory_pool_t;
 
 static __TLS umf_result_t TLS_last_allocation_error;
@@ -437,12 +437,29 @@ static void *op_aligned_alloc(void *pool, size_t size, size_t alignment) {
 
 static umf_result_t op_initialize(umf_memory_provider_handle_t provider,
                                   const void *params, void **out_pool) {
-    assert(provider);
+    (void)params;
+    (void)provider;
     assert(out_pool);
 
+    jemalloc_memory_pool_t *pool = umf_ba_global_alloc(sizeof(*pool));
+    if (!pool) {
+        return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
+    }
+    memset(pool, 0, sizeof(*pool));
+    *out_pool = pool;
+
+    return UMF_RESULT_SUCCESS;
+}
+
+static umf_result_t op_post_initialize(umf_memory_provider_handle_t provider,
+                                       const void *params, void *pool) {
+    assert(provider);
+    assert(pool);
+
     extent_hooks_t *pHooks = &arena_extent_hooks;
     size_t unsigned_size = sizeof(unsigned);
     int n_arenas_set_from_params = 0;
+    jemalloc_memory_pool_t *je_pool = (jemalloc_memory_pool_t *)pool;
     int err;
     const umf_jemalloc_pool_params_t *jemalloc_params = params;
 
@@ -454,32 +471,35 @@ static umf_result_t op_initialize(umf_memory_provider_handle_t provider,
 
     if (n_arenas == 0) {
         n_arenas = utils_get_num_cores() * 4;
-        if (n_arenas > MALLOCX_ARENA_MAX) {
-            n_arenas = MALLOCX_ARENA_MAX;
-        }
+        n_arenas = utils_min(n_arenas, (size_t)MALLOCX_ARENA_MAX);
     }
 
     if (n_arenas > MALLOCX_ARENA_MAX) {
         LOG_ERR("Number of arenas %zu exceeds the limit (%i).", n_arenas,
                 MALLOCX_ARENA_MAX);
+        umf_ba_global_free(je_pool);
         return UMF_RESULT_ERROR_INVALID_ARGUMENT;
     }
 
-    jemalloc_memory_pool_t *pool = umf_ba_global_alloc(
-        sizeof(*pool) + n_arenas * sizeof(*pool->arena_index));
-    if (!pool) {
+    je_pool->arena_index =
+        umf_ba_global_alloc(n_arenas * sizeof(*je_pool->arena_index));
+    if (!je_pool->arena_index) {
+        LOG_ERR("Could not allocate memory for arena indices.");
+        umf_ba_global_free(je_pool);
         return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
     }
-    memset(pool, 0, sizeof(*pool) + n_arenas * sizeof(*pool->arena_index));
+
+    memset(je_pool->arena_index, 0, n_arenas * sizeof(*je_pool->arena_index));
+
     const char *pool_name = DEFAULT_NAME;
     if (jemalloc_params) {
         pool_name = jemalloc_params->name;
     }
 
-    snprintf(pool->name, sizeof(pool->name), "%s", pool_name);
+    snprintf(je_pool->name, sizeof(je_pool->name), "%s", pool_name);
 
-    pool->provider = provider;
-    pool->n_arenas = n_arenas;
+    je_pool->provider = provider;
+    je_pool->n_arenas = n_arenas;
 
     size_t num_created = 0;
     for (size_t i = 0; i < n_arenas; i++) {
@@ -504,13 +524,13 @@ static umf_result_t op_initialize(umf_memory_provider_handle_t provider,
             break;
         }
 
-        pool->arena_index[num_created++] = arena_index;
+        je_pool->arena_index[num_created++] = arena_index;
         if (arena_index >= MALLOCX_ARENA_MAX) {
             LOG_ERR("Number of arenas exceeds the limit.");
             goto err_cleanup;
         }
 
-        pool_by_arena_index[arena_index] = pool;
+        pool_by_arena_index[arena_index] = je_pool;
 
         // Setup extent_hooks for the newly created arena.
         char cmd[64];
@@ -521,21 +541,24 @@ static umf_result_t op_initialize(umf_memory_provider_handle_t provider,
             goto err_cleanup;
         }
     }
-    *out_pool = (umf_memory_pool_handle_t)pool;
 
-    VALGRIND_DO_CREATE_MEMPOOL(pool, 0, 0);
+    VALGRIND_DO_CREATE_MEMPOOL(je_pool, 0, 0);
 
     return UMF_RESULT_SUCCESS;
 
 err_cleanup:
     // Destroy any arenas that were successfully created.
     for (size_t i = 0; i < num_created; i++) {
         char cmd[64];
-        unsigned arena = pool->arena_index[i];
+        unsigned arena = je_pool->arena_index[i];
         snprintf(cmd, sizeof(cmd), "arena.%u.destroy", arena);
         (void)je_mallctl(cmd, NULL, 0, NULL, 0);
     }
-    umf_ba_global_free(pool);
+    if (je_pool->arena_index) {
+        umf_ba_global_free(je_pool->arena_index);
+        je_pool->arena_index = NULL;
+    }
+    umf_ba_global_free(je_pool);
     return UMF_RESULT_ERROR_MEMORY_PROVIDER_SPECIFIC;
 }
 
@@ -552,6 +575,9 @@ static umf_result_t op_finalize(void *pool) {
             ret = UMF_RESULT_ERROR_UNKNOWN;
         }
     }
+    if (je_pool->arena_index) {
+        umf_ba_global_free(je_pool->arena_index);
+    }
     umf_ba_global_free(je_pool);
 
     VALGRIND_DO_DESTROY_MEMPOOL(pool);
@@ -600,6 +626,7 @@ static umf_memory_pool_ops_t UMF_JEMALLOC_POOL_OPS = {
     .free = op_free,
     .get_last_allocation_error = op_get_last_allocation_error,
     .get_name = op_get_name,
+    .ext_post_initialize = op_post_initialize,
 };
 
 const umf_memory_pool_ops_t *umfJemallocPoolOps(void) {