@@ -68,7 +68,7 @@ typedef struct jemalloc_memory_pool_t {
6868 umf_memory_provider_handle_t provider ;
6969 size_t n_arenas ;
7070 char name [64 ];
71- unsigned int arena_index [] ;
71+ unsigned int * arena_index ;
7272} jemalloc_memory_pool_t ;
7373
7474static __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) {
437437
438438static umf_result_t op_initialize (umf_memory_provider_handle_t provider ,
439439 const void * params , void * * out_pool ) {
440- assert (provider );
440+ (void )params ;
441+ (void )provider ;
441442 assert (out_pool );
442443
444+ jemalloc_memory_pool_t * pool = umf_ba_global_alloc (sizeof (* pool ));
445+ if (!pool ) {
446+ return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY ;
447+ }
448+ memset (pool , 0 , sizeof (* pool ));
449+ * out_pool = pool ;
450+
451+ return UMF_RESULT_SUCCESS ;
452+ }
453+
454+ static umf_result_t op_post_initialize (umf_memory_provider_handle_t provider ,
455+ const void * params , void * pool ) {
456+ assert (provider );
457+ assert (pool );
458+
443459 extent_hooks_t * pHooks = & arena_extent_hooks ;
444460 size_t unsigned_size = sizeof (unsigned );
445461 int n_arenas_set_from_params = 0 ;
462+ jemalloc_memory_pool_t * je_pool = (jemalloc_memory_pool_t * )pool ;
446463 int err ;
447464 const umf_jemalloc_pool_params_t * jemalloc_params = params ;
448465
@@ -454,32 +471,35 @@ static umf_result_t op_initialize(umf_memory_provider_handle_t provider,
454471
455472 if (n_arenas == 0 ) {
456473 n_arenas = utils_get_num_cores () * 4 ;
457- if (n_arenas > MALLOCX_ARENA_MAX ) {
458- n_arenas = MALLOCX_ARENA_MAX ;
459- }
474+ n_arenas = utils_min (n_arenas , (size_t )MALLOCX_ARENA_MAX );
460475 }
461476
462477 if (n_arenas > MALLOCX_ARENA_MAX ) {
463478 LOG_ERR ("Number of arenas %zu exceeds the limit (%i)." , n_arenas ,
464479 MALLOCX_ARENA_MAX );
480+ umf_ba_global_free (je_pool );
465481 return UMF_RESULT_ERROR_INVALID_ARGUMENT ;
466482 }
467483
468- jemalloc_memory_pool_t * pool = umf_ba_global_alloc (
469- sizeof (* pool ) + n_arenas * sizeof (* pool -> arena_index ));
470- if (!pool ) {
484+ je_pool -> arena_index =
485+ umf_ba_global_alloc (n_arenas * sizeof (* je_pool -> arena_index ));
486+ if (!je_pool -> arena_index ) {
487+ LOG_ERR ("Could not allocate memory for arena indices." );
488+ umf_ba_global_free (je_pool );
471489 return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY ;
472490 }
473- memset (pool , 0 , sizeof (* pool ) + n_arenas * sizeof (* pool -> arena_index ));
491+
492+ memset (je_pool -> arena_index , 0 , n_arenas * sizeof (* je_pool -> arena_index ));
493+
474494 const char * pool_name = DEFAULT_NAME ;
475495 if (jemalloc_params ) {
476496 pool_name = jemalloc_params -> name ;
477497 }
478498
479- snprintf (pool -> name , sizeof (pool -> name ), "%s" , pool_name );
499+ snprintf (je_pool -> name , sizeof (je_pool -> name ), "%s" , pool_name );
480500
481- pool -> provider = provider ;
482- pool -> n_arenas = n_arenas ;
501+ je_pool -> provider = provider ;
502+ je_pool -> n_arenas = n_arenas ;
483503
484504 size_t num_created = 0 ;
485505 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,
504524 break ;
505525 }
506526
507- pool -> arena_index [num_created ++ ] = arena_index ;
527+ je_pool -> arena_index [num_created ++ ] = arena_index ;
508528 if (arena_index >= MALLOCX_ARENA_MAX ) {
509529 LOG_ERR ("Number of arenas exceeds the limit." );
510530 goto err_cleanup ;
511531 }
512532
513- pool_by_arena_index [arena_index ] = pool ;
533+ pool_by_arena_index [arena_index ] = je_pool ;
514534
515535 // Setup extent_hooks for the newly created arena.
516536 char cmd [64 ];
@@ -521,21 +541,24 @@ static umf_result_t op_initialize(umf_memory_provider_handle_t provider,
521541 goto err_cleanup ;
522542 }
523543 }
524- * out_pool = (umf_memory_pool_handle_t )pool ;
525544
526- VALGRIND_DO_CREATE_MEMPOOL (pool , 0 , 0 );
545+ VALGRIND_DO_CREATE_MEMPOOL (je_pool , 0 , 0 );
527546
528547 return UMF_RESULT_SUCCESS ;
529548
530549err_cleanup :
531550 // Destroy any arenas that were successfully created.
532551 for (size_t i = 0 ; i < num_created ; i ++ ) {
533552 char cmd [64 ];
534- unsigned arena = pool -> arena_index [i ];
553+ unsigned arena = je_pool -> arena_index [i ];
535554 snprintf (cmd , sizeof (cmd ), "arena.%u.destroy" , arena );
536555 (void )je_mallctl (cmd , NULL , 0 , NULL , 0 );
537556 }
538- umf_ba_global_free (pool );
557+ if (je_pool -> arena_index ) {
558+ umf_ba_global_free (je_pool -> arena_index );
559+ je_pool -> arena_index = NULL ;
560+ }
561+ umf_ba_global_free (je_pool );
539562 return UMF_RESULT_ERROR_MEMORY_PROVIDER_SPECIFIC ;
540563}
541564
@@ -552,6 +575,9 @@ static umf_result_t op_finalize(void *pool) {
552575 ret = UMF_RESULT_ERROR_UNKNOWN ;
553576 }
554577 }
578+ if (je_pool -> arena_index ) {
579+ umf_ba_global_free (je_pool -> arena_index );
580+ }
555581 umf_ba_global_free (je_pool );
556582
557583 VALGRIND_DO_DESTROY_MEMPOOL (pool );
@@ -600,6 +626,7 @@ static umf_memory_pool_ops_t UMF_JEMALLOC_POOL_OPS = {
600626 .free = op_free ,
601627 .get_last_allocation_error = op_get_last_allocation_error ,
602628 .get_name = op_get_name ,
629+ .ext_post_initialize = op_post_initialize ,
603630};
604631
605632const umf_memory_pool_ops_t * umfJemallocPoolOps (void ) {
0 commit comments