Skip to content

feat: add custom memory allocators #3373

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

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

feat: add custom memory allocators #3373

Show file tree
Hide file tree
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
2 changes: 1 addition & 1 deletion Makefile.am
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -135,7 +135,7 @@ endif

### Tests (make check)

TESTS = tests/mantest tests/jqtest tests/shtest tests/utf8test tests/base64test tests/uritest
TESTS = tests/mantest tests/jqtest tests/shtest tests/utf8test tests/base64test tests/uritest tests/alloctest
if !WIN32
TESTS += tests/optionaltest
endif
Expand Down
207 changes: 198 additions & 9 deletions src/jv_alloc.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,19 @@
#include <stdlib.h>
#include <string.h>
#include "jv.h"
#include "jv_alloc.h"

struct nomem_handler {
jv_nomem_handler_f handler;
void *data;
};

struct custom_allocator {
jv_malloc_t malloc_fn;
jv_free_t free_fn;
jv_realloc_t realloc_fn;
};

#if !defined(HAVE_PTHREAD_KEY_CREATE) || \
!defined(HAVE_PTHREAD_ONCE) || \
!defined(HAVE_ATEXIT)
Expand All @@ -18,11 +25,13 @@ struct nomem_handler {
#ifdef _MSC_VER
/* Visual C++: yes */
static __declspec(thread) struct nomem_handler nomem_handler;
static __declspec(thread) struct custom_allocator custom_allocator;
#define USE_TLS
#else
#ifdef HAVE___THREAD
/* GCC and friends: yes */
static __thread struct nomem_handler nomem_handler;
static __thread struct custom_allocator custom_allocator;
#define USE_TLS
#endif /* HAVE___THREAD */
#endif /* _MSC_VER */
Expand All @@ -47,17 +56,26 @@ static void memory_exhausted(void) {
#include <pthread.h>

static pthread_key_t nomem_handler_key;
static pthread_key_t custom_allocator_key;
static pthread_once_t mem_once = PTHREAD_ONCE_INIT;

/* tsd_fini is called on application exit */
/* it clears the nomem_handler allocated in the main thread */
/* it clears the nomem_handler and custom_allocator allocated in the main thread */
static void tsd_fini(void) {
struct nomem_handler *nomem_handler;
struct custom_allocator *custom_alloc;

nomem_handler = pthread_getspecific(nomem_handler_key);
if (nomem_handler) {
(void) pthread_setspecific(nomem_handler_key, NULL);
free(nomem_handler);
}

custom_alloc = pthread_getspecific(custom_allocator_key);
if (custom_alloc) {
(void) pthread_setspecific(custom_allocator_key, NULL);
free(custom_alloc);
}
}

/* The tsd_fini_thread is a destructor set by calling */
Expand All @@ -67,11 +85,20 @@ static void tsd_fini_thread(void *nomem_handler) {
free(nomem_handler);
}

/* Destructor for custom allocator thread-specific data */
static void tsd_fini_custom_alloc(void *custom_alloc) {
free(custom_alloc);
}

static void tsd_init(void) {
if (pthread_key_create(&nomem_handler_key, tsd_fini_thread) != 0) {
fprintf(stderr, "jq: error: cannot create thread specific key");
abort();
}
if (pthread_key_create(&custom_allocator_key, tsd_fini_custom_alloc) != 0) {
fprintf(stderr, "jq: error: cannot create thread specific key");
abort();
}
if (atexit(tsd_fini) != 0) {
fprintf(stderr, "jq: error: cannot set an exit handler");
abort();
Expand All @@ -89,6 +116,17 @@ static void tsd_init_nomem_handler(void)
}
}

static void tsd_init_custom_allocator(void)
{
if (pthread_getspecific(custom_allocator_key) == NULL) {
struct custom_allocator *custom_alloc = calloc(1, sizeof(struct custom_allocator));
if (pthread_setspecific(custom_allocator_key, custom_alloc) != 0) {
fprintf(stderr, "jq: error: cannot set thread specific data");
abort();
}
}
}

void jv_nomem_handler(jv_nomem_handler_f handler, void *data) {
pthread_once(&mem_once, tsd_init); // cannot fail
tsd_init_nomem_handler();
Expand Down Expand Up @@ -124,6 +162,8 @@ static void memory_exhausted(void) {
/* No thread-local storage of any kind that we know how to handle */

static struct nomem_handler nomem_handler;
static struct custom_allocator custom_allocator;

void jv_nomem_handler(jv_nomem_handler_f handler, void *data) {
nomem_handler.handler = handler;
nomem_handler.data = data;
Expand All @@ -138,21 +178,65 @@ static void memory_exhausted(void) {
#endif /* USE_TLS */


// Helper functions to get custom allocator based on threading model
#ifdef USE_TLS
static struct custom_allocator* get_custom_allocator(void) {
return &custom_allocator;
}
#else
#ifdef HAVE_PTHREAD_KEY_CREATE
static struct custom_allocator* get_custom_allocator(void) {
pthread_once(&mem_once, tsd_init);
tsd_init_custom_allocator();
return pthread_getspecific(custom_allocator_key);
}
#else
static struct custom_allocator* get_custom_allocator(void) {
return &custom_allocator;
}
#endif
#endif

void* jv_mem_alloc(size_t sz) {
void* p = malloc(sz);
struct custom_allocator* custom_alloc = get_custom_allocator();
void* p;

if (custom_alloc && custom_alloc->malloc_fn) {
p = custom_alloc->malloc_fn(sz);
} else {
p = malloc(sz);
}

if (!p) {
memory_exhausted();
}
return p;
}

void* jv_mem_alloc_unguarded(size_t sz) {
return malloc(sz);
struct custom_allocator* custom_alloc = get_custom_allocator();

if (custom_alloc && custom_alloc->malloc_fn) {
return custom_alloc->malloc_fn(sz);
} else {
return malloc(sz);
}
}

void* jv_mem_calloc(size_t nemb, size_t sz) {
assert(nemb > 0 && sz > 0);
void* p = calloc(nemb, sz);
struct custom_allocator* custom_alloc = get_custom_allocator();
void* p;

if (custom_alloc && custom_alloc->malloc_fn) {
p = custom_alloc->malloc_fn(nemb * sz);
if (p) {
memset(p, 0, nemb * sz);
}
} else {
p = calloc(nemb, sz);
}

if (!p) {
memory_exhausted();
}
Expand All @@ -161,29 +245,134 @@ void* jv_mem_calloc(size_t nemb, size_t sz) {

void* jv_mem_calloc_unguarded(size_t nemb, size_t sz) {
assert(nemb > 0 && sz > 0);
return calloc(nemb, sz);
struct custom_allocator* custom_alloc = get_custom_allocator();

if (custom_alloc && custom_alloc->malloc_fn) {
void* p = custom_alloc->malloc_fn(nemb * sz);
if (p) {
memset(p, 0, nemb * sz);
}
return p;
} else {
return calloc(nemb, sz);
}
}

char* jv_mem_strdup(const char *s) {
char *p = strdup(s);
struct custom_allocator* custom_alloc = get_custom_allocator();
char *p;

if (custom_alloc && custom_alloc->malloc_fn) {
size_t len = strlen(s) + 1;
p = custom_alloc->malloc_fn(len);
if (p) {
memcpy(p, s, len);
}
} else {
p = strdup(s);
}

if (!p) {
memory_exhausted();
}
return p;
}

char* jv_mem_strdup_unguarded(const char *s) {
return strdup(s);
struct custom_allocator* custom_alloc = get_custom_allocator();

if (custom_alloc && custom_alloc->malloc_fn) {
size_t len = strlen(s) + 1;
char *p = custom_alloc->malloc_fn(len);
if (p) {
memcpy(p, s, len);
}
return p;
} else {
return strdup(s);
}
}

void jv_mem_free(void* p) {
free(p);
struct custom_allocator* custom_alloc = get_custom_allocator();

if (custom_alloc && custom_alloc->free_fn) {
custom_alloc->free_fn(p);
} else {
free(p);
}
}

void* jv_mem_realloc(void* p, size_t sz) {
p = realloc(p, sz);
struct custom_allocator* custom_alloc = get_custom_allocator();

if (custom_alloc && custom_alloc->realloc_fn) {
p = custom_alloc->realloc_fn(p, sz);
} else {
p = realloc(p, sz);
}

if (!p) {
memory_exhausted();
}
return p;
}

#ifdef USE_TLS
// Thread-local storage implementation
void jv_set_alloc_funcs(jv_malloc_t malloc_fn, jv_free_t free_fn) {
custom_allocator.malloc_fn = malloc_fn;
custom_allocator.free_fn = free_fn;
custom_allocator.realloc_fn = NULL;
}

void jv_set_alloc_funcs_ex(jv_malloc_t malloc_fn, jv_free_t free_fn, jv_realloc_t realloc_fn) {
custom_allocator.malloc_fn = malloc_fn;
custom_allocator.free_fn = free_fn;
custom_allocator.realloc_fn = realloc_fn;
}

#else

#ifdef HAVE_PTHREAD_KEY_CREATE
// Pthread-based implementation
void jv_set_alloc_funcs(jv_malloc_t malloc_fn, jv_free_t free_fn) {
pthread_once(&mem_once, tsd_init);
tsd_init_custom_allocator();

struct custom_allocator *custom_alloc = pthread_getspecific(custom_allocator_key);
if (custom_alloc) {
custom_alloc->malloc_fn = malloc_fn;
custom_alloc->free_fn = free_fn;
custom_alloc->realloc_fn = NULL;
}
}

void jv_set_alloc_funcs_ex(jv_malloc_t malloc_fn, jv_free_t free_fn, jv_realloc_t realloc_fn) {
pthread_once(&mem_once, tsd_init);
tsd_init_custom_allocator();

struct custom_allocator *custom_alloc = pthread_getspecific(custom_allocator_key);
if (custom_alloc) {
custom_alloc->malloc_fn = malloc_fn;
custom_alloc->free_fn = free_fn;
custom_alloc->realloc_fn = realloc_fn;
}
}

#else
// Global variable implementation
void jv_set_alloc_funcs(jv_malloc_t malloc_fn, jv_free_t free_fn) {
custom_allocator.malloc_fn = malloc_fn;
custom_allocator.free_fn = free_fn;
custom_allocator.realloc_fn = NULL;
}

void jv_set_alloc_funcs_ex(jv_malloc_t malloc_fn, jv_free_t free_fn, jv_realloc_t realloc_fn) {
custom_allocator.malloc_fn = malloc_fn;
custom_allocator.free_fn = free_fn;
custom_allocator.realloc_fn = realloc_fn;
}

#endif /* HAVE_PTHREAD_KEY_CREATE */
#endif /* USE_TLS */
7 changes: 7 additions & 0 deletions src/jv_alloc.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,10 @@

#include <stddef.h>

typedef void* (*jv_malloc_t)(size_t size);
typedef void (*jv_free_t)(void* ptr);
typedef void* (*jv_realloc_t)(void* ptr, size_t size);

void* jv_mem_alloc(size_t);
void* jv_mem_alloc_unguarded(size_t);
void* jv_mem_calloc(size_t, size_t);
Expand All @@ -12,4 +16,7 @@ char* jv_mem_strdup_unguarded(const char *);
void jv_mem_free(void*);
__attribute__((warn_unused_result)) void* jv_mem_realloc(void*, size_t);

void jv_set_alloc_funcs(jv_malloc_t malloc_fn, jv_free_t free_fn);
void jv_set_alloc_funcs_ex(jv_malloc_t malloc_fn, jv_free_t free_fn, jv_realloc_t realloc_fn);

#endif
Loading
Loading