Skip to content

Commit 48d5146

Browse files
vinser52vinser52
committed
Update jemalloc pool params implementation
1 parent 882a4c8 commit 48d5146

File tree

4 files changed

+217
-10
lines changed

4 files changed

+217
-10
lines changed

‎examples/dram_and_fsdax/dram_and_fsdax.c

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -67,19 +67,35 @@ static umf_memory_pool_handle_t create_fsdax_pool(const char *path) {
6767
// so it should be used with a pool manager that will take over
6868
// the managing of the provided memory - for example the jemalloc pool
6969
// with the `disable_provider_free` parameter set to true.
70-
umf_jemalloc_pool_params_t pool_params;
71-
pool_params.disable_provider_free = true;
70+
umf_jemalloc_pool_params_handle_t pool_params;
71+
umf_result = umfJemallocPoolParamsCreate(&pool_params);
72+
if (umf_result != UMF_RESULT_SUCCESS) {
73+
fprintf(stderr, "Failed to create jemalloc params!\n");
74+
umfMemoryProviderDestroy(provider_fsdax);
75+
return NULL;
76+
}
77+
umf_result = umfJemallocPoolParamsSetKeepAllMemory(pool_params, true);
78+
if (umf_result != UMF_RESULT_SUCCESS) {
79+
fprintf(stderr, "Failed to set KeepAllMemory!\n");
80+
umfMemoryProviderDestroy(provider_fsdax);
81+
return NULL;
82+
}
7283

7384
// Create an FSDAX memory pool
7485
umf_result =
75-
umfPoolCreate(umfJemallocPoolOps(), provider_fsdax, &pool_params,
86+
umfPoolCreate(umfJemallocPoolOps(), provider_fsdax, pool_params,
7687
UMF_POOL_CREATE_FLAG_OWN_PROVIDER, &pool_fsdax);
7788
if (umf_result != UMF_RESULT_SUCCESS) {
7889
fprintf(stderr, "Failed to create an FSDAX memory pool!\n");
7990
umfMemoryProviderDestroy(provider_fsdax);
8091
return NULL;
8192
}
8293

94+
umf_result = umfJemallocPoolParamsDestroy(pool_params);
95+
if (umf_result != UMF_RESULT_SUCCESS) {
96+
fprintf(stderr, "Failed to destroy jemalloc params!\n");
97+
}
98+
8399
return pool_fsdax;
84100
}
85101

‎include/umf/pools/pool_jemalloc.h

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,31 @@ extern "C" {
1717
#include <stdbool.h>
1818
#include <umf/memory_pool_ops.h>
1919

20-
/// @brief Configuration of Jemalloc Pool
21-
typedef struct umf_jemalloc_pool_params_t {
22-
/// Set to true if umfMemoryProviderFree() should never be called.
23-
bool disable_provider_free;
24-
} umf_jemalloc_pool_params_t;
20+
struct umf_jemalloc_pool_params_t;
21+
22+
/// @brief handle to the parameters of the jemalloc pool.
23+
typedef struct umf_jemalloc_pool_params_t *umf_jemalloc_pool_params_handle_t;
24+
25+
/// @brief Create a struct to store parameters of jemalloc pool.
26+
/// @param hParams [out] handle to the newly created parameters struct.
27+
/// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
28+
umf_result_t
29+
umfJemallocPoolParamsCreate(umf_jemalloc_pool_params_handle_t *hParams);
30+
31+
/// @brief Destroy parameters struct.
32+
/// @param hParams handle to the parameters of the jemalloc pool.
33+
/// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
34+
umf_result_t
35+
umfJemallocPoolParamsDestroy(umf_jemalloc_pool_params_handle_t hParams);
36+
37+
/// @brief Set if \p umfMemoryProviderFree() should never be called.
38+
/// @param hParams handle to the parameters of the jemalloc pool.
39+
/// @param keepAllMemory \p true if the jemalloc pool should not call
40+
/// \p umfMemoryProviderFree, \p false otherwise.
41+
/// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
42+
umf_result_t
43+
umfJemallocPoolParamsSetKeepAllMemory(umf_jemalloc_pool_params_handle_t hParams,
44+
bool keepAllMemory);
2545

2646
umf_memory_pool_ops_t *umfJemallocPoolOps(void);
2747

‎src/pool/pool_jemalloc.c

Lines changed: 54 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,12 @@ typedef struct jemalloc_memory_pool_t {
4141
bool disable_provider_free;
4242
} jemalloc_memory_pool_t;
4343

44+
// Configuration of Jemalloc Pool
45+
typedef struct umf_jemalloc_pool_params_t {
46+
/// Set to true if umfMemoryProviderFree() should never be called.
47+
bool disable_provider_free;
48+
} umf_jemalloc_pool_params_t;
49+
4450
static __TLS umf_result_t TLS_last_allocation_error;
4551

4652
static jemalloc_memory_pool_t *pool_by_arena_index[MALLCTL_ARENAS_ALL];
@@ -53,6 +59,52 @@ static jemalloc_memory_pool_t *get_pool_by_arena_index(unsigned arena_ind) {
5359
return pool_by_arena_index[arena_ind];
5460
}
5561

62+
umf_result_t
63+
umfJemallocPoolParamsCreate(umf_jemalloc_pool_params_handle_t *hParams) {
64+
if (!hParams) {
65+
LOG_ERR("jemalloc pool params handle is NULL");
66+
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
67+
}
68+
69+
umf_jemalloc_pool_params_t *params_data =
70+
umf_ba_global_alloc(sizeof(*params_data));
71+
if (!params_data) {
72+
LOG_ERR("cannot allocate memory for jemalloc poolparams");
73+
return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
74+
}
75+
76+
params_data->disable_provider_free = false;
77+
78+
*hParams = (umf_jemalloc_pool_params_handle_t)params_data;
79+
80+
return UMF_RESULT_SUCCESS;
81+
}
82+
83+
umf_result_t
84+
umfJemallocPoolParamsDestroy(umf_jemalloc_pool_params_handle_t hParams) {
85+
if (!hParams) {
86+
LOG_ERR("jemalloc pool params handle is NULL");
87+
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
88+
}
89+
90+
umf_ba_global_free(hParams);
91+
92+
return UMF_RESULT_SUCCESS;
93+
}
94+
95+
umf_result_t
96+
umfJemallocPoolParamsSetKeepAllMemory(umf_jemalloc_pool_params_handle_t hParams,
97+
bool keepAllMemory) {
98+
if (!hParams) {
99+
LOG_ERR("jemalloc pool params handle is NULL");
100+
return UMF_RESULT_ERROR_INVALID_ARGUMENT;
101+
}
102+
103+
hParams->disable_provider_free = keepAllMemory;
104+
105+
return UMF_RESULT_SUCCESS;
106+
}
107+
56108
// arena_extent_alloc - an extent allocation function conforms to the extent_alloc_t type and upon
57109
// success returns a pointer to size bytes of mapped memory on behalf of arena arena_ind such that
58110
// the extent's base address is a multiple of alignment, as well as setting *zero to indicate
@@ -401,8 +453,8 @@ static umf_result_t op_initialize(umf_memory_provider_handle_t provider,
401453
assert(provider);
402454
assert(out_pool);
403455

404-
umf_jemalloc_pool_params_t *je_params =
405-
(umf_jemalloc_pool_params_t *)params;
456+
umf_jemalloc_pool_params_handle_t je_params =
457+
(umf_jemalloc_pool_params_handle_t)params;
406458

407459
extent_hooks_t *pHooks = &arena_extent_hooks;
408460
size_t unsigned_size = sizeof(unsigned);

‎test/pools/jemalloc_pool.cpp

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,3 +42,122 @@ TEST_F(test, metadataNotAllocatedUsingProvider) {
4242
[pool = pool.get()](void *ptr) { umfPoolFree(pool, ptr); });
4343
}
4444
}
45+
46+
using jemallocPoolParams = bool;
47+
struct umfJemallocPoolParamsTest
48+
: umf_test::test,
49+
::testing::WithParamInterface<jemallocPoolParams> {
50+
51+
struct validation_params_t {
52+
bool keep_all_memory;
53+
};
54+
55+
struct provider_validator : public umf_test::provider_ba_global {
56+
using base_provider = umf_test::provider_ba_global;
57+
58+
umf_result_t initialize(validation_params_t *params) {
59+
EXPECT_NE(params, nullptr);
60+
expected_params = params;
61+
return UMF_RESULT_SUCCESS;
62+
}
63+
umf_result_t free(void *ptr, size_t size) {
64+
EXPECT_EQ(expected_params->keep_all_memory, false);
65+
return base_provider::free(ptr, size);
66+
}
67+
68+
validation_params_t *expected_params;
69+
};
70+
71+
static constexpr umf_memory_provider_ops_t VALIDATOR_PROVIDER_OPS =
72+
umf::providerMakeCOps<provider_validator, validation_params_t>();
73+
74+
umfJemallocPoolParamsTest() : expected_params{false}, params(nullptr) {}
75+
void SetUp() override {
76+
test::SetUp();
77+
expected_params.keep_all_memory = this->GetParam();
78+
umf_result_t ret = umfJemallocPoolParamsCreate(&params);
79+
ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
80+
ret = umfJemallocPoolParamsSetKeepAllMemory(
81+
params, expected_params.keep_all_memory);
82+
ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
83+
}
84+
85+
void TearDown() override {
86+
umfJemallocPoolParamsDestroy(params);
87+
test::TearDown();
88+
}
89+
90+
umf::pool_unique_handle_t makePool() {
91+
umf_memory_provider_handle_t hProvider = nullptr;
92+
umf_memory_pool_handle_t hPool = nullptr;
93+
94+
auto ret = umfMemoryProviderCreate(&VALIDATOR_PROVIDER_OPS,
95+
&expected_params, &hProvider);
96+
EXPECT_EQ(ret, UMF_RESULT_SUCCESS);
97+
98+
ret = umfPoolCreate(umfJemallocPoolOps(), hProvider, params,
99+
UMF_POOL_CREATE_FLAG_OWN_PROVIDER, &hPool);
100+
EXPECT_EQ(ret, UMF_RESULT_SUCCESS);
101+
102+
return umf::pool_unique_handle_t(hPool, &umfPoolDestroy);
103+
}
104+
105+
void allocFreeFlow() {
106+
static const size_t ALLOC_SIZE = 128;
107+
static const size_t NUM_ALLOCATIONS = 100;
108+
std::vector<void *> ptrs;
109+
110+
auto pool = makePool();
111+
ASSERT_NE(pool, nullptr);
112+
113+
for (size_t i = 0; i < NUM_ALLOCATIONS; ++i) {
114+
auto *ptr = umfPoolMalloc(pool.get(), ALLOC_SIZE);
115+
ASSERT_NE(ptr, nullptr);
116+
ptrs.push_back(ptr);
117+
}
118+
119+
for (size_t i = 0; i < NUM_ALLOCATIONS; ++i) {
120+
auto ret = umfPoolFree(pool.get(), ptrs[i]);
121+
ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
122+
}
123+
124+
// Now pool can call free during pool destruction
125+
expected_params.keep_all_memory = false;
126+
}
127+
128+
validation_params_t expected_params;
129+
umf_jemalloc_pool_params_handle_t params;
130+
};
131+
132+
TEST_P(umfJemallocPoolParamsTest, allocFree) { allocFreeFlow(); }
133+
134+
TEST_P(umfJemallocPoolParamsTest, updateParams) {
135+
expected_params.keep_all_memory = !expected_params.keep_all_memory;
136+
umf_result_t ret = umfJemallocPoolParamsSetKeepAllMemory(
137+
params, expected_params.keep_all_memory);
138+
ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
139+
140+
allocFreeFlow();
141+
}
142+
143+
TEST_P(umfJemallocPoolParamsTest, invalidParams) {
144+
umf_result_t ret = umfJemallocPoolParamsCreate(nullptr);
145+
ASSERT_EQ(ret, UMF_RESULT_ERROR_INVALID_ARGUMENT);
146+
147+
ret = umfJemallocPoolParamsSetKeepAllMemory(nullptr, true);
148+
ASSERT_EQ(ret, UMF_RESULT_ERROR_INVALID_ARGUMENT);
149+
150+
ret = umfJemallocPoolParamsSetKeepAllMemory(nullptr, false);
151+
ASSERT_EQ(ret, UMF_RESULT_ERROR_INVALID_ARGUMENT);
152+
153+
ret = umfJemallocPoolParamsDestroy(nullptr);
154+
ASSERT_EQ(ret, UMF_RESULT_ERROR_INVALID_ARGUMENT);
155+
}
156+
157+
GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(umfJemallocPoolParamsTest);
158+
159+
/* TODO: enable this test after the issue #903 is fixed.
160+
(https://github.com/oneapi-src/unified-memory-framework/issues/903)
161+
INSTANTIATE_TEST_SUITE_P(jemallocPoolTest, umfJemallocPoolParamsTest,
162+
testing::Values(false, true));
163+
*/

0 commit comments

Comments
 (0)