Skip to content

Commit

Permalink
Change TLSF memory format for alignment
Browse files Browse the repository at this point in the history
  • Loading branch information
@sykwer
sykwer committed Dec 25, 2023
1 parent 22c8536 commit be7f2ab
Showing 1 changed file with 56 additions and 132 deletions.
188 changes: 56 additions & 132 deletions src/tlsf/tlsf.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,14 @@
#include <pthread.h>
#include <stdio.h>
#include <unistd.h>
#include <pthread.h>
#include <sys/mman.h>
#include <string.h>
#include <malloc.h>

#include <cstdint>
#include <unordered_map>
#include <string>
#include <memory>

#include "tlsf/tlsf.h"

Expand All @@ -19,8 +19,7 @@

static char * mempool_ptr;
static size_t INITIAL_MEMPOOL_SIZE = 100 * 1000 * 1000; // default: 100MB
static size_t ADDITIONAL_MEMPOOL_SIZE = 100 * 1000 * 1000; // default: 100MB
static std::unordered_map<void *, void *> * aligned2orig;
static size_t ADDITIONAL_MEMPOOL_SIZE = 50 * 1000 * 1000; // default: 50MB

static pthread_mutex_t init_mtx = PTHREAD_MUTEX_INITIALIZER;
static pthread_cond_t init_cond = PTHREAD_COND_INITIALIZER; // mempool_initialized == true
Expand All @@ -44,14 +43,12 @@ static void initialize_mempool()
MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
memset(mempool_ptr, 0, INITIAL_MEMPOOL_SIZE);
init_memory_pool(INITIAL_MEMPOOL_SIZE, mempool_ptr); // tlsf library function

// aligned2orig.reserve(10000000);
}

// Do not use printf
void check_mempool_initialized()
{
if (mempool_initialized) {return;}
if (mempool_initialized) return;

pthread_mutex_lock(&init_mtx);

Expand All @@ -60,7 +57,6 @@ void check_mempool_initialized()
pthread_mutex_unlock(&init_mtx);

initialize_mempool();
aligned2orig = new std::unordered_map<void *, void *>();

mempool_initialized = true;
pthread_cond_signal(&init_cond);
Expand All @@ -73,64 +69,67 @@ void check_mempool_initialized()
}
}

template<class F>
static void * tlsf_allocate_internal(F allocate)
static void * tlsf_malloc_wrapped(size_t size)
{
pthread_mutex_lock(&tlsf_mtx);

void * ret = allocate();
void * ret = tlsf_malloc(size);
pthread_mutex_unlock(&tlsf_mtx);

size_t multiplier = 1;
while (ret == NULL) {
char * addr = (char *) mmap(
NULL, multiplier * ADDITIONAL_MEMPOOL_SIZE, PROT_READ | PROT_WRITE,
MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);

pthread_mutex_lock(&tlsf_mtx);
add_new_area(addr, multiplier * ADDITIONAL_MEMPOOL_SIZE, mempool_ptr); // tlsf library function
pthread_mutex_unlock(&tlsf_mtx);

fprintf(
stderr, "TLSF memory pool exhausted: %lu bytes additionally mmaped.\n",
multiplier * ADDITIONAL_MEMPOOL_SIZE);

ret = allocate();
pthread_mutex_lock(&tlsf_mtx);
ret = tlsf_malloc(size);;
pthread_mutex_unlock(&tlsf_mtx);

multiplier *= 2;
}

pthread_mutex_unlock(&tlsf_mtx);
return ret;
}

static void * tlsf_malloc_wrapped(size_t size)
{
return tlsf_allocate_internal([size] {return tlsf_malloc(size);});
}

static void * tlsf_calloc_wrapped(size_t num, size_t size)
{
return tlsf_allocate_internal([num, size] {return tlsf_calloc(num, size);});
}

static void * tlsf_realloc_wrapped(void * ptr, size_t new_size)
{
return tlsf_allocate_internal([ptr, new_size] {return tlsf_realloc(ptr, new_size);});
}

static void tlsf_free_wrapped(void * ptr)
{
void** p = reinterpret_cast<void**>(ptr);

pthread_mutex_lock(&tlsf_mtx);
tlsf_free(ptr);
tlsf_free(p[-1]);
pthread_mutex_unlock(&tlsf_mtx);
}

static void * tlsf_aligned_malloc(size_t alignment, size_t size)
{
void * addr = tlsf_malloc_wrapped(alignment + size);
void * aligned =
reinterpret_cast<void *>(reinterpret_cast<uint64_t>(addr) + alignment -
reinterpret_cast<uint64_t>(addr) % alignment);
(*aligned2orig)[aligned] = addr;
size_t total_size = size + alignment + sizeof(void*) + sizeof(size_t);
void* raw_memory = tlsf_malloc_wrapped(total_size);

void* aligned_memory = static_cast<char*>(raw_memory) + sizeof(void*) + sizeof(size_t);
size_t remain = total_size - sizeof(void*) + sizeof(size_t);
void* result = std::align(alignment, size, aligned_memory /* updated */, remain /* updated */);

if (!result) {
tlsf_free_wrapped(raw_memory);
throw std::bad_alloc();
}

void** p = reinterpret_cast<void**>(aligned_memory);
size_t* p2 = reinterpret_cast<size_t*>(p - 1);
p[-1] = raw_memory;
p2[-1] = alignment;

//printf("In tlsf_aligned_malloc: orig=%p -> aligned=%p\n", addr, aligned);

return aligned;
return aligned_memory;
}

using namespace heaphook;
Expand All @@ -139,41 +138,26 @@ class TlfsAllocator : public GlobalAllocator
{
void * do_alloc(size_t size, size_t alignment) override
{
if (alignment == 1) {
static malloc_type original_malloc =
static malloc_type original_malloc =
reinterpret_cast<malloc_type>(dlsym(RTLD_NEXT, "malloc"));
static __thread bool malloc_no_hook = false;

if (malloc_no_hook) {
if (mempool_initialized) {
return tlsf_malloc_wrapped(size);
} else {
return original_malloc(size);
}
}

malloc_no_hook = true;
check_mempool_initialized();
void * ret = tlsf_malloc_wrapped(size);
malloc_no_hook = false;
return ret;
} else {
static aligned_alloc_type original_aligned_alloc =
static aligned_alloc_type original_aligned_alloc =
reinterpret_cast<aligned_alloc_type>(dlsym(RTLD_NEXT, "aligned_alloc"));
static __thread bool aligned_alloc_no_hook = false;

if (aligned_alloc_no_hook /*|| pthread_self() == logging_thread*/) {
if (mempool_initialized) {return tlsf_aligned_malloc(alignment, size);} else {
return original_aligned_alloc(alignment, size);
}
}
static __thread bool malloc_no_hook = false;

aligned_alloc_no_hook = true;
check_mempool_initialized();
void * ret = tlsf_aligned_malloc(alignment, size);
aligned_alloc_no_hook = false;
return ret;
if (malloc_no_hook) {
if (mempool_initialized) {
return tlsf_aligned_malloc(alignment, size);
} else {
return alignment == 1 ? original_malloc(size) : original_aligned_alloc(alignment, size);
}
}

malloc_no_hook = true;
check_mempool_initialized();
void * ret = tlsf_aligned_malloc(alignment, size);
malloc_no_hook = false;
return ret;
}

void do_dealloc(void * ptr) override
Expand All @@ -193,13 +177,6 @@ class TlfsAllocator : public GlobalAllocator

free_no_hook = true;
check_mempool_initialized();

auto it = aligned2orig->find(ptr);
if (it != aligned2orig->end()) {
ptr = it->second;
aligned2orig->erase(it);
}

tlsf_free_wrapped(ptr);
free_no_hook = false;
}
Expand All @@ -217,69 +194,16 @@ class TlfsAllocator : public GlobalAllocator
// |--------------------| -+
// buf_ptr --> | buffer |
// |--------------------|
void * buf_ptr = ptr;
size_t buf_addr = reinterpret_cast<size_t>(buf_ptr);
auto it = aligned2orig->find(ptr);
if (it != aligned2orig->end()) {
// If block is aligned, ptr does not point to the beginning of the block.
void * block_ptr = it->second;
size_t block_addr = reinterpret_cast<size_t>(block_ptr);
size_t block_size = (*reinterpret_cast<size_t *>(block_addr - 8)) & (~0b1111ull);
size_t unused_size = buf_addr - block_addr;
return block_size - unused_size;
} else {
size_t block_size = (*reinterpret_cast<size_t *>(buf_addr - 8)) & (~0b1111ull);
return block_size;
}
}

void * do_alloc_zeroed(size_t size) override
{
static calloc_type original_calloc = reinterpret_cast<calloc_type>(dlsym(RTLD_NEXT, "calloc"));
static __thread bool calloc_no_hook = false;

if (calloc_no_hook) {
if (mempool_initialized) {
return tlsf_calloc_wrapped(size, 1);
} else {
return original_calloc(size, 1);
}
}

calloc_no_hook = true;
check_mempool_initialized();
void * ret = tlsf_calloc_wrapped(size, 1);
calloc_no_hook = false;

return ret;
}
void ** p = reinterpret_cast<void**>(ptr);
void* block_ptr = p[-1];
size_t block_addr = reinterpret_cast<size_t>(block_ptr);

void * do_realloc(void * ptr, size_t new_size) override
{
static realloc_type original_realloc =
reinterpret_cast<realloc_type>(dlsym(RTLD_NEXT, "realloc"));
static __thread bool realloc_no_hook = false;
size_t* p2 = reinterpret_cast<size_t*>(p - 1);
size_t alignment = p2[-1];

if (realloc_no_hook) {
if (mempool_initialized) {
return tlsf_realloc_wrapped(ptr, new_size);
} else {
return original_realloc(ptr, new_size);
}
}

realloc_no_hook = true;
check_mempool_initialized();

auto it = aligned2orig->find(ptr);
if (it != aligned2orig->end()) {
ptr = it->second;
aligned2orig->erase(ptr);
}

void * ret = tlsf_realloc_wrapped(ptr, new_size);
realloc_no_hook = false;
return ret;
size_t block_size = (*reinterpret_cast<size_t *>(block_addr - 8)) & (~0b1111ull);
return block_size - sizeof(void*) - sizeof(size_t) - alignment;
}
};

Expand Down

0 comments on commit be7f2ab

Please sign in to comment.