Skip to content

Commit cf4932b

Browse files
bratpiorkabratpiorka
committed
add free handle POC
1 parent ac9094d commit cf4932b

11 files changed

+331
-37
lines changed

Diff for: ‎include/umf/proxy_lib_handlers.h

+54
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
/*
2+
* Copyright (C) 2024 Intel Corporation
3+
*
4+
* Under the Apache License v2.0 with LLVM Exceptions. See LICENSE.TXT.
5+
* SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
*/
7+
8+
#ifndef UMF_PROXY_LIB_HANDLERS_H
9+
#define UMF_PROXY_LIB_HANDLERS_H 1
10+
11+
#include <umf/memory_pool.h>
12+
13+
#ifdef __cplusplus
14+
extern "C" {
15+
#endif
16+
17+
// TODO - improve and cleanup comments
18+
19+
// malloc API handlers
20+
// NOTE - in malloc/aligned_malloc pre handlers the default pool could be
21+
// changed along with the requested size and alignment
22+
// - in free pre handler user can update the ptr
23+
typedef void (*umf_proxy_lib_handler_malloc_pre_t)(
24+
void *user_data, umf_memory_pool_handle_t *pool, size_t *size);
25+
typedef void (*umf_proxy_lib_handler_aligned_malloc_pre_t)(
26+
void *user_data, umf_memory_pool_handle_t *pool, size_t *size,
27+
size_t *alignment);
28+
typedef void (*umf_proxy_lib_handler_free_pre_t)(void *user_data, void **ptr,
29+
umf_memory_pool_handle_t pool);
30+
31+
void umfSetProxyLibHandlerMallocPre(umf_proxy_lib_handler_malloc_pre_t handler,
32+
void *user_data);
33+
void umfSetProxyLibHandlerAlignedMallocPre(
34+
umf_proxy_lib_handler_aligned_malloc_pre_t handler, void *user_data);
35+
void umfSetProxyLibHandlerFreePre(umf_proxy_lib_handler_free_pre_t handler,
36+
void *user_data);
37+
38+
// NOTE - in the malloc/aligned_malloc post handlers the pointer to allocated
39+
// data could be changed by the user
40+
typedef void (*umf_proxy_lib_handler_malloc_post_t)(
41+
void *user_data, void **ptr, umf_memory_pool_handle_t pool);
42+
typedef void (*umf_proxy_lib_handler_aligned_malloc_post_t)(
43+
void *user_data, void **ptr, umf_memory_pool_handle_t pool);
44+
45+
void umfSetProxyLibHandlerMallocPost(
46+
umf_proxy_lib_handler_malloc_post_t handler, void *user_data);
47+
void umfSetProxyLibHandlerAlignedMallocPost(
48+
umf_proxy_lib_handler_aligned_malloc_post_t handler, void *user_data);
49+
50+
#ifdef __cplusplus
51+
}
52+
#endif
53+
54+
#endif /* UMF_PROXY_LIB_HANDLERS_H */

Diff for: ‎src/proxy_lib/proxy_lib.c

+103-11
Original file line numberDiff line numberDiff line change
@@ -50,9 +50,9 @@
5050
#include <umf/memory_pool.h>
5151
#include <umf/memory_provider.h>
5252
#include <umf/providers/provider_os_memory.h>
53+
#include <umf/proxy_lib_handlers.h>
5354

5455
#include "base_alloc_linear.h"
55-
#include "proxy_lib.h"
5656
#include "utils_common.h"
5757
#include "utils_load_library.h"
5858
#include "utils_log.h"
@@ -135,6 +135,22 @@ static __TLS int was_called_from_umfPool = 0;
135135
// TODO remove this WA when the issue is fixed.
136136
static __TLS int was_called_from_malloc_usable_size = 0;
137137

138+
// malloc API handlers
139+
static umf_proxy_lib_handler_malloc_pre_t Handler_malloc_pre = NULL;
140+
static umf_proxy_lib_handler_aligned_malloc_pre_t Handler_aligned_malloc_pre =
141+
NULL;
142+
static umf_proxy_lib_handler_free_pre_t Handler_free_pre = NULL;
143+
144+
static umf_proxy_lib_handler_malloc_post_t Handler_malloc_post = NULL;
145+
static umf_proxy_lib_handler_aligned_malloc_post_t Handler_aligned_malloc_post =
146+
NULL;
147+
148+
static void *Handler_malloc_pre_user_data = NULL;
149+
static void *Handler_aligned_malloc_pre_user_data = NULL;
150+
static void *Handler_free_pre_user_data = NULL;
151+
static void *Handler_malloc_post_user_data = NULL;
152+
static void *Handler_aligned_malloc_post_user_data = NULL;
153+
138154
/*****************************************************************************/
139155
/*** The constructor and destructor of the proxy library *********************/
140156
/*****************************************************************************/
@@ -353,20 +369,35 @@ static inline size_t ba_leak_pool_contains_pointer(void *ptr) {
353369
/*****************************************************************************/
354370

355371
void *malloc(size_t size) {
372+
umf_memory_pool_handle_t pool = Proxy_pool;
373+
if (Handler_malloc_pre) {
374+
Handler_malloc_pre(Handler_malloc_pre_user_data, &pool, &size);
375+
}
376+
377+
void *ptr = NULL;
356378
#ifndef _WIN32
357379
if (size < Size_threshold_value) {
358-
return System_malloc(size);
380+
ptr = System_malloc(size);
381+
goto handler_post;
359382
}
383+
360384
#endif /* _WIN32 */
361385

362-
if (!was_called_from_umfPool && Proxy_pool) {
386+
if (!was_called_from_umfPool && pool) {
363387
was_called_from_umfPool = 1;
364-
void *ptr = umfPoolMalloc(Proxy_pool, size);
388+
ptr = umfPoolMalloc(pool, size);
365389
was_called_from_umfPool = 0;
366-
return ptr;
390+
goto handler_post;
367391
}
368392

369-
return ba_leak_malloc(size);
393+
ptr = ba_leak_malloc(size);
394+
395+
handler_post:
396+
if (Handler_malloc_post) {
397+
Handler_malloc_post(Handler_malloc_post_user_data, &ptr, pool);
398+
}
399+
400+
return ptr;
370401
}
371402

372403
void *calloc(size_t nmemb, size_t size) {
@@ -387,15 +418,28 @@ void *calloc(size_t nmemb, size_t size) {
387418
}
388419

389420
void free(void *ptr) {
421+
umf_memory_pool_handle_t pool = NULL;
422+
if (Handler_free_pre) {
423+
pool = umfPoolByPtr(ptr);
424+
Handler_free_pre(Handler_free_pre_user_data, &ptr, pool);
425+
}
426+
390427
if (ptr == NULL) {
391428
return;
392429
}
393430

431+
// NOTE: for system allocations made during UMF and Proxy Lib
432+
// initialisation, we never call free handlers, as they should handle
433+
// only user-made allocations
394434
if (ba_leak_free(ptr) == 0) {
395435
return;
396436
}
397437

398-
if (Proxy_pool && (umfPoolByPtr(ptr) == Proxy_pool)) {
438+
if (Proxy_pool && pool == NULL) {
439+
pool = umfPoolByPtr(ptr);
440+
}
441+
442+
if (pool == Proxy_pool) {
399443
if (umfPoolFree(Proxy_pool, ptr) != UMF_RESULT_SUCCESS) {
400444
LOG_ERR("umfPoolFree() failed");
401445
}
@@ -448,20 +492,36 @@ void *realloc(void *ptr, size_t size) {
448492
}
449493

450494
void *aligned_alloc(size_t alignment, size_t size) {
495+
umf_memory_pool_handle_t pool = Proxy_pool;
496+
if (Handler_aligned_malloc_pre) {
497+
Handler_aligned_malloc_pre(Handler_aligned_malloc_pre_user_data, &pool,
498+
&size, &alignment);
499+
}
500+
501+
void *ptr = NULL;
451502
#ifndef _WIN32
452503
if (size < Size_threshold_value) {
453-
return System_aligned_alloc(alignment, size);
504+
ptr = System_aligned_alloc(alignment, size);
505+
goto handler_post;
454506
}
455507
#endif /* _WIN32 */
456508

457509
if (!was_called_from_umfPool && Proxy_pool) {
458510
was_called_from_umfPool = 1;
459-
void *ptr = umfPoolAlignedMalloc(Proxy_pool, size, alignment);
511+
ptr = umfPoolAlignedMalloc(Proxy_pool, size, alignment);
460512
was_called_from_umfPool = 0;
461-
return ptr;
513+
goto handler_post;
462514
}
463515

464-
return ba_leak_aligned_alloc(alignment, size);
516+
ptr = ba_leak_aligned_alloc(alignment, size);
517+
518+
handler_post:
519+
if (Handler_aligned_malloc_post) {
520+
Handler_aligned_malloc_post(Handler_aligned_malloc_post_user_data, &ptr,
521+
pool);
522+
}
523+
524+
return ptr;
465525
}
466526

467527
#ifdef _WIN32
@@ -555,3 +615,35 @@ void *_aligned_offset_recalloc(void *ptr, size_t num, size_t size,
555615
}
556616

557617
#endif
618+
619+
// malloc API handlers
620+
621+
void umfSetProxyLibHandlerMallocPre(umf_proxy_lib_handler_malloc_pre_t handler,
622+
void *user_data) {
623+
Handler_malloc_pre = handler;
624+
Handler_malloc_pre_user_data = user_data;
625+
}
626+
627+
void umfSetProxyLibHandlerAlignedMallocPre(
628+
umf_proxy_lib_handler_aligned_malloc_pre_t handler, void *user_data) {
629+
Handler_aligned_malloc_pre = handler;
630+
Handler_aligned_malloc_pre_user_data = user_data;
631+
}
632+
633+
void umfSetProxyLibHandlerFreePre(umf_proxy_lib_handler_free_pre_t handler,
634+
void *user_data) {
635+
Handler_free_pre = handler;
636+
Handler_free_pre_user_data = user_data;
637+
}
638+
639+
void umfSetProxyLibHandlerMallocPost(
640+
umf_proxy_lib_handler_malloc_post_t handler, void *user_data) {
641+
Handler_malloc_post = handler;
642+
Handler_malloc_post_user_data = user_data;
643+
}
644+
645+
void umfSetProxyLibHandlerAlignedMallocPost(
646+
umf_proxy_lib_handler_aligned_malloc_post_t handler, void *user_data) {
647+
Handler_aligned_malloc_post = handler;
648+
Handler_aligned_malloc_post_user_data = user_data;
649+
}

Diff for: ‎src/proxy_lib/proxy_lib.def

+6
Original file line numberDiff line numberDiff line change
@@ -21,3 +21,9 @@ EXPORTS
2121
_aligned_offset_malloc
2222
_aligned_offset_realloc
2323
_aligned_offset_recalloc
24+
; handlers
25+
umfSetProxyLibHandlerMallocPre
26+
umfSetProxyLibHandlerAlignedMallocPre
27+
umfSetProxyLibHandlerFreePre
28+
umfSetProxyLibHandlerMallocPost
29+
umfSetProxyLibHandlerAlignedMallocPost

Diff for: ‎src/proxy_lib/proxy_lib.h

-22
This file was deleted.

Diff for: ‎src/proxy_lib/proxy_lib.map

+6
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,12 @@
1212
malloc;
1313
malloc_usable_size;
1414
realloc;
15+
# handlers
16+
umfSetProxyLibHandlerMallocPre;
17+
umfSetProxyLibHandlerAlignedMallocPre;
18+
umfSetProxyLibHandlerFreePre;
19+
umfSetProxyLibHandlerMallocPost;
20+
umfSetProxyLibHandlerAlignedMallocPost;
1521
local:
1622
*;
1723
};

Diff for: ‎src/proxy_lib/proxy_lib_linux.c

+4-1
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,10 @@
77
*
88
*/
99

10-
#include "proxy_lib.h"
10+
#include <umf/proxy_lib_handlers.h>
11+
12+
void proxy_lib_create_common(void);
13+
void proxy_lib_destroy_common(void);
1114

1215
// The priority 102 is used, because the constructor should be called as the second one
1316
// (just after the first constructor of the base allocator with priority 101)

Diff for: ‎src/proxy_lib/proxy_lib_windows.c

+4-1
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,10 @@
99

1010
#include <Windows.h>
1111

12-
#include "proxy_lib.h"
12+
#include <umf/proxy_lib_handlers.h>
13+
14+
void proxy_lib_create_common(void);
15+
void proxy_lib_destroy_common(void);
1316

1417
static void proxy_lib_create(void) { proxy_lib_create_common(); }
1518

Diff for: ‎test/CMakeLists.txt

+7-2
Original file line numberDiff line numberDiff line change
@@ -447,14 +447,19 @@ add_umf_test(
447447
if(UMF_PROXY_LIB_ENABLED AND UMF_BUILD_SHARED_LIBRARY)
448448
add_umf_test(
449449
NAME proxy_lib_basic
450-
SRCS ${BA_SOURCES_FOR_TEST} test_proxy_lib.cpp
450+
SRCS ${BA_SOURCES_FOR_TEST} proxy_lib/proxy_lib.cpp
451+
LIBS ${UMF_UTILS_FOR_TEST} umf_proxy)
452+
453+
add_umf_test(
454+
NAME proxy_lib_handlers
455+
SRCS ${BA_SOURCES_FOR_TEST} proxy_lib/proxy_lib_handlers.cpp
451456
LIBS ${UMF_UTILS_FOR_TEST} umf_proxy)
452457

453458
# TODO enable this test on Windows
454459
if(LINUX)
455460
add_umf_test(
456461
NAME test_proxy_lib_size_threshold
457-
SRCS ${BA_SOURCES_FOR_TEST} test_proxy_lib_size_threshold.cpp
462+
SRCS ${BA_SOURCES_FOR_TEST} proxy_lib/proxy_lib_size_threshold.cpp
458463
LIBS ${UMF_UTILS_FOR_TEST} umf_proxy)
459464
set_property(TEST umf-test_proxy_lib_size_threshold
460465
PROPERTY ENVIRONMENT UMF_PROXY="size.threshold=64")

Diff for: ‎test/test_proxy_lib.cpp renamed to ‎test/proxy_lib/proxy_lib.cpp

File renamed without changes.

0 commit comments

Comments
 (0)