Skip to content

Commit 2f50f1e

Browse files
melvercopybara-github
authored andcommitted
alloc-token: Add stubs for 2-token configuration
This adds stubs to TCMalloc to wire up the __alloc_token allocation functions when building with Clang's Alloction token feature [1]. The currently support configuration requires building with: ``` clang -fsanitize=alloc-token \ -fsanitize-alloc-token-fast-abi \ -falloc-token-max=2 \ -DALLOC_TOKEN_MAX=2 \ -mllvm -alloc-token-fallback=$FALLBACK \ -DALLOC_TOKEN_FALLBACK=$FALLBACK ... ``` This is a preparatory step; when used, this has no effect at runtime yet. [1] https://clang.llvm.org/docs/AllocToken.html PiperOrigin-RevId: 822131093 Change-Id: Ibd5f62acb853f7d163879f6bf9be7bccd690eeed
1 parent 4111da1 commit 2f50f1e

File tree

1 file changed

+104
-0
lines changed

1 file changed

+104
-0
lines changed

tcmalloc/libc_override.h

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
#include <new>
3434

3535
#include "tcmalloc/alloc_at_least.h"
36+
#include "tcmalloc/malloc_extension.h" // IWYU pragma: keep
3637
#include "tcmalloc/tcmalloc.h" // IWYU pragma: keep
3738

3839
#define TCMALLOC_ALIAS(tc_fn) \
@@ -217,6 +218,109 @@ alloc_result_t aligned_alloc_at_least(size_t alignment,
217218
size_t min_size) TCMALLOC_NOTHROW
218219
TCMALLOC_ALIAS(TCMallocInternalAlignedAllocAtLeast);
219220

221+
//
222+
// Partitioned Variants for Clang's -fsanitize=alloc-token.
223+
//
224+
#ifdef __SANITIZE_ALLOC_TOKEN__
225+
#define DEFINE_ALLOC_TOKEN_NEW(id) \
226+
void* __alloc_token_##id##__Znwm(size_t) \
227+
TCMALLOC_ALIAS(TCMallocInternalNew); \
228+
void* __alloc_token_##id##__Znam(size_t) \
229+
TCMALLOC_ALIAS(TCMallocInternalNewArray); \
230+
void* __alloc_token_##id##__ZnwmRKSt9nothrow_t(size_t, \
231+
const std::nothrow_t&) \
232+
TCMALLOC_ALIAS(TCMallocInternalNewNothrow); \
233+
void* __alloc_token_##id##__ZnamRKSt9nothrow_t(size_t, \
234+
const std::nothrow_t&) \
235+
TCMALLOC_ALIAS(TCMallocInternalNewArrayNothrow); \
236+
void* __alloc_token_##id##__ZnwmSt11align_val_t(size_t, std::align_val_t) \
237+
TCMALLOC_ALIAS(TCMallocInternalNewAligned); \
238+
void* __alloc_token_##id##__ZnamSt11align_val_t(size_t, std::align_val_t) \
239+
TCMALLOC_ALIAS(TCMallocInternalNewArrayAligned); \
240+
void* __alloc_token_##id##__ZnwmSt11align_val_tRKSt9nothrow_t( \
241+
size_t, std::align_val_t, const std::nothrow_t&) \
242+
TCMALLOC_ALIAS(TCMallocInternalNewAlignedNothrow); \
243+
void* __alloc_token_##id##__ZnamSt11align_val_tRKSt9nothrow_t( \
244+
size_t, std::align_val_t, const std::nothrow_t&) \
245+
TCMALLOC_ALIAS(TCMallocInternalNewArrayAlignedNothrow);
246+
247+
#ifndef TCMALLOC_INTERNAL_METHODS_ONLY
248+
#define DEFINE_ALLOC_TOKEN_NEW_EXTENSION(id) \
249+
void* __alloc_token_##id##__Znwm12__hot_cold_t(size_t, __hot_cold_t) \
250+
TCMALLOC_ALIAS(_Znwm12__hot_cold_t); \
251+
void* __alloc_token_##id##__ZnwmRKSt9nothrow_t12__hot_cold_t( \
252+
size_t, const std::nothrow_t&, \
253+
__hot_cold_t) noexcept TCMALLOC_ALIAS(_ZnwmRKSt9nothrow_t12__hot_cold_t); \
254+
void* __alloc_token_##id##__ZnwmSt11align_val_t12__hot_cold_t( \
255+
size_t, std::align_val_t, __hot_cold_t) \
256+
TCMALLOC_ALIAS(_ZnwmSt11align_val_t12__hot_cold_t); \
257+
void* __alloc_token_##id##__ZnwmSt11align_val_tRKSt9nothrow_t12__hot_cold_t( \
258+
size_t, std::align_val_t, const std::nothrow_t&, \
259+
__hot_cold_t) noexcept TCMALLOC_ALIAS(_ZnwmSt11align_val_tRKSt9nothrow_t12__hot_cold_t); \
260+
void* __alloc_token_##id##__Znam12__hot_cold_t(size_t, __hot_cold_t) \
261+
TCMALLOC_ALIAS(_Znam12__hot_cold_t); \
262+
void* __alloc_token_##id##__ZnamRKSt9nothrow_t12__hot_cold_t( \
263+
size_t, const std::nothrow_t&, \
264+
__hot_cold_t) noexcept TCMALLOC_ALIAS(_ZnamRKSt9nothrow_t12__hot_cold_t); \
265+
void* __alloc_token_##id##__ZnamSt11align_val_t12__hot_cold_t( \
266+
size_t, std::align_val_t, __hot_cold_t) \
267+
TCMALLOC_ALIAS(_ZnamSt11align_val_t12__hot_cold_t); \
268+
void* __alloc_token_##id##__ZnamSt11align_val_tRKSt9nothrow_t12__hot_cold_t( \
269+
size_t, std::align_val_t, const std::nothrow_t&, \
270+
__hot_cold_t) noexcept TCMALLOC_ALIAS(_ZnamSt11align_val_tRKSt9nothrow_t12__hot_cold_t); \
271+
__sized_ptr_t __alloc_token_##id##___size_returning_new(size_t) \
272+
TCMALLOC_ALIAS(__size_returning_new); \
273+
__sized_ptr_t __alloc_token_##id##___size_returning_new_aligned( \
274+
size_t, std::align_val_t) TCMALLOC_ALIAS(__size_returning_new_aligned); \
275+
__sized_ptr_t __alloc_token_##id##___size_returning_new_hot_cold( \
276+
size_t, __hot_cold_t) TCMALLOC_ALIAS(__size_returning_new_hot_cold); \
277+
__sized_ptr_t __alloc_token_##id##___size_returning_new_aligned_hot_cold( \
278+
size_t, std::align_val_t, __hot_cold_t) \
279+
TCMALLOC_ALIAS(__size_returning_new_aligned_hot_cold);
280+
#else
281+
#define ALLOC_TOKEN_NEW_EXTENSION(id)
282+
#endif // TCMALLOC_INTERNAL_METHODS_ONLY
283+
284+
#define DEFINE_ALLOC_TOKEN_STDLIB(id) \
285+
void* __alloc_token_##id##_malloc(size_t) \
286+
TCMALLOC_NOTHROW TCMALLOC_ALIAS(TCMallocInternalMalloc); \
287+
void* __alloc_token_##id##_realloc(void*, size_t) \
288+
TCMALLOC_NOTHROW TCMALLOC_ALIAS(TCMallocInternalRealloc); \
289+
void* __alloc_token_##id##_reallocarray(void*, size_t, size_t) \
290+
TCMALLOC_NOTHROW TCMALLOC_ALIAS(TCMallocInternalReallocArray); \
291+
void* __alloc_token_##id##_calloc(size_t, size_t) \
292+
TCMALLOC_NOTHROW TCMALLOC_ALIAS(TCMallocInternalCalloc); \
293+
void* __alloc_token_##id##_memalign(size_t, size_t) \
294+
TCMALLOC_NOTHROW TCMALLOC_ALIAS(TCMallocInternalMemalign); \
295+
void* __alloc_token_##id##_aligned_alloc(size_t, size_t) \
296+
TCMALLOC_NOTHROW TCMALLOC_ALIAS(TCMallocInternalAlignedAlloc); \
297+
void* __alloc_token_##id##_valloc(size_t) \
298+
TCMALLOC_NOTHROW TCMALLOC_ALIAS(TCMallocInternalValloc); \
299+
void* __alloc_token_##id##_pvalloc(size_t) \
300+
TCMALLOC_NOTHROW TCMALLOC_ALIAS(TCMallocInternalPvalloc); \
301+
int __alloc_token_##id##_posix_memalign(void**, size_t, size_t) \
302+
TCMALLOC_NOTHROW TCMALLOC_ALIAS(TCMallocInternalPosixMemalign);
303+
304+
#define DEFINE_ALLOC_TOKEN_VARIANTS(id) \
305+
DEFINE_ALLOC_TOKEN_NEW(id) \
306+
DEFINE_ALLOC_TOKEN_NEW_EXTENSION(id) \
307+
DEFINE_ALLOC_TOKEN_STDLIB(id)
308+
309+
#ifndef ALLOC_TOKEN_MAX
310+
#error "Define ALLOC_TOKEN_MAX to match -falloc-token-max=<max number of IDs>"
311+
#endif
312+
static_assert(ALLOC_TOKEN_MAX == 2);
313+
DEFINE_ALLOC_TOKEN_VARIANTS(0)
314+
DEFINE_ALLOC_TOKEN_VARIANTS(1)
315+
#ifdef ALLOC_TOKEN_FALLBACK
316+
// Define the functions for the fallback token ID if overridden with -mllvm
317+
// -alloc-token-fallback=N; should fall outside the range of normal token IDs.
318+
static_assert(ALLOC_TOKEN_FALLBACK >= ALLOC_TOKEN_MAX);
319+
DEFINE_ALLOC_TOKEN_VARIANTS(ALLOC_TOKEN_FALLBACK)
320+
#endif
321+
322+
#endif // __SANITIZE_ALLOC_TOKEN__
323+
220324
} // extern "C"
221325

222326
// NOLINTEND(misc-definitions-in-headers)

0 commit comments

Comments
 (0)