Skip to content

Commit bf1999f

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

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)(void *user_data,
24+
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

+102-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,34 @@ static inline size_t ba_leak_pool_contains_pointer(void *ptr) {
353369
/*****************************************************************************/
354370

355371
void *malloc(size_t size) {
372+
if (Handler_malloc_pre) {
373+
Handler_malloc_pre(Handler_malloc_pre_user_data, &size);
374+
}
375+
376+
void *ptr = NULL;
356377
#ifndef _WIN32
357378
if (size < Size_threshold_value) {
358-
return System_malloc(size);
379+
ptr = System_malloc(size);
380+
goto handler_post;
359381
}
360382
#endif /* _WIN32 */
361383

362-
if (!was_called_from_umfPool && Proxy_pool) {
384+
umf_memory_pool_handle_t pool = Proxy_pool;
385+
if (!was_called_from_umfPool && pool) {
363386
was_called_from_umfPool = 1;
364-
void *ptr = umfPoolMalloc(Proxy_pool, size);
387+
ptr = umfPoolMalloc(pool, size);
365388
was_called_from_umfPool = 0;
366-
return ptr;
389+
goto handler_post;
367390
}
368391

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

372402
void *calloc(size_t nmemb, size_t size) {
@@ -387,15 +417,28 @@ void *calloc(size_t nmemb, size_t size) {
387417
}
388418

389419
void free(void *ptr) {
420+
umf_memory_pool_handle_t pool = NULL;
421+
if (Handler_free_pre) {
422+
pool = umfPoolByPtr(ptr);
423+
Handler_free_pre(Handler_free_pre_user_data, &ptr, pool);
424+
}
425+
390426
if (ptr == NULL) {
391427
return;
392428
}
393429

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

398-
if (Proxy_pool && (umfPoolByPtr(ptr) == Proxy_pool)) {
437+
if (Proxy_pool && pool == NULL) {
438+
pool = umfPoolByPtr(ptr);
439+
}
440+
441+
if (pool == Proxy_pool) {
399442
if (umfPoolFree(Proxy_pool, ptr) != UMF_RESULT_SUCCESS) {
400443
LOG_ERR("umfPoolFree() failed");
401444
}
@@ -448,20 +491,36 @@ void *realloc(void *ptr, size_t size) {
448491
}
449492

450493
void *aligned_alloc(size_t alignment, size_t size) {
494+
umf_memory_pool_handle_t pool = Proxy_pool;
495+
if (Handler_aligned_malloc_pre) {
496+
Handler_aligned_malloc_pre(Handler_aligned_malloc_pre_user_data, &pool,
497+
&size, &alignment);
498+
}
499+
500+
void *ptr = NULL;
451501
#ifndef _WIN32
452502
if (size < Size_threshold_value) {
453-
return System_aligned_alloc(alignment, size);
503+
ptr = System_aligned_alloc(alignment, size);
504+
goto handler_post;
454505
}
455506
#endif /* _WIN32 */
456507

457508
if (!was_called_from_umfPool && Proxy_pool) {
458509
was_called_from_umfPool = 1;
459-
void *ptr = umfPoolAlignedMalloc(Proxy_pool, size, alignment);
510+
ptr = umfPoolAlignedMalloc(Proxy_pool, size, alignment);
460511
was_called_from_umfPool = 0;
461-
return ptr;
512+
goto handler_post;
462513
}
463514

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

467526
#ifdef _WIN32
@@ -555,3 +614,35 @@ void *_aligned_offset_recalloc(void *ptr, size_t num, size_t size,
555614
}
556615

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

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)