Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Change TLSF memory format for alignment #9

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Change TLSF memory format for alignment #9

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
188 changes: 56 additions & 132 deletions src/tlsf/tlsf.cpp
View file
Open in desktop
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
Loading