Repository: valkey-io/valkey
File: src/zmalloc.c (C)

Code before:

/* Explicitly override malloc/free etc when using tcmalloc. */
#if defined(USE_TCMALLOC)
#define malloc(size) tc_malloc(size)
#define calloc(count, size) tc_calloc(count, size)
#define realloc(ptr, size) tc_realloc(ptr, size)
#define free(ptr) tc_free(ptr)
/* Explicitly override malloc/free etc when using jemalloc. */
#elif defined(USE_JEMALLOC)
#define malloc(size) je_malloc(size)
#define calloc(count, size) je_calloc(count, size)
#define realloc(ptr, size) je_realloc(ptr, size)
#define free(ptr) je_free(ptr)
#define mallocx(size, flags) je_mallocx(size, flags)
#define dallocx(ptr, flags) je_dallocx(ptr, flags)
#endif

#if __STDC_NO_THREADS__
#define thread_local __thread
#else
#include <threads.h>
#endif

/* A thread-local storage which keep the current thread's index in the used_memory_thread array. */
static thread_local int thread_index = -1;
/* MAX_THREADS_NUM = IO_THREADS_MAX_NUM(128) + BIO threads(3) + main thread(1). */
#define MAX_THREADS_NUM 132
static unsigned long long used_memory_thread[MAX_THREADS_NUM];

static atomic_int total_active_threads;

/* Register the thread index in start_routine. */
static inline void zmalloc_register_thread_index(void) {
    thread_index = atomic_fetch_add_explicit(&total_active_threads, 1, memory_order_relaxed);
    /* TODO: Handle the case when exceed the MAX_THREADS_NUM (may rarely happen). */
    assert(total_active_threads < MAX_THREADS_NUM);
}

static inline void update_zmalloc_stat_alloc(size_t size) {
    if (unlikely(thread_index == -1)) zmalloc_register_thread_index();
    used_memory_thread[thread_index] += size;
}

static inline void update_zmalloc_stat_free(size_t size) {
    if (unlikely(thread_index == -1)) zmalloc_register_thread_index();
    used_memory_thread[thread_index] -= size;
}

static void zmalloc_default_oom(size_t size) {
    fprintf(stderr, "zmalloc: Out of memory trying to allocate %zu bytes\n", size);
    fflush(stderr);

Review comment (performance):
I think this is the only hold up for me. Can we move all these settings out of `networking.c` and into a new header say `os.h`

```c
  #define IO_THREADS_MAX_NUM 128
  #ifndef CACHE_LINE_SIZE
  #if defined(__aarch64__) && defined(__APPLE__)
  #define CACHE_LINE_SIZE 128
  #else
  #define CACHE_LINE_SIZE 64
  #endif
  #endif
  ``` 

Code after:

/* Explicitly override malloc/free etc when using tcmalloc. */
#if defined(USE_TCMALLOC)
#define malloc(size) tc_malloc(size)
#define calloc(count, size) tc_calloc(count, size)
#define realloc(ptr, size) tc_realloc(ptr, size)
#define free(ptr) tc_free(ptr)
/* Explicitly override malloc/free etc when using jemalloc. */
#elif defined(USE_JEMALLOC)
#define malloc(size) je_malloc(size)
#define calloc(count, size) je_calloc(count, size)
#define realloc(ptr, size) je_realloc(ptr, size)
#define free(ptr) je_free(ptr)
#define mallocx(size, flags) je_mallocx(size, flags)
#define dallocx(ptr, flags) je_dallocx(ptr, flags)
#endif

#if __STDC_NO_THREADS__
#define thread_local __thread
#else
#include <threads.h>
#endif

#define MAX_THREADS_NUM (IO_THREADS_MAX_NUM + 3 + 1)
/* A thread-local storage which keep the current thread's index in the used_memory_thread array. */
static thread_local int thread_index = -1;
/* Element in used_memory_thread array should only be written by a single thread which
 * distinguished by the thread-local storage thread_index. But when an element in
 * used_memory_thread array was written, it could be read by another thread simultaneously,
 * the reader will see the inconsistency memory on non x86 architecture potentially.
 * For the ARM and PowerPC platform, we can solve this issue by make the memory aligned.
 * For the other architecture, lets fall back to the atomic operation to keep safe. */
#if defined(__i386__) || defined(__x86_64__) || defined(__amd64__) || defined(__POWERPC__) || defined(__arm__) ||      \
    defined(__arm64__)
static __attribute__((aligned(sizeof(size_t)))) size_t used_memory_thread[MAX_THREADS_NUM];
#else
static _Atomic size_t used_memory_thread[MAX_THREADS_NUM];
#endif
static atomic_int total_active_threads = 0;
/* This is a simple protection. It's used only if some modules create a lot of threads. */
static atomic_size_t used_memory_for_additional_threads = 0;

/* Register the thread index in start_routine. */
static inline void zmalloc_register_thread_index(void) {
    thread_index = atomic_fetch_add_explicit(&total_active_threads, 1, memory_order_relaxed);
}

static inline void update_zmalloc_stat_alloc(size_t size) {
    if (unlikely(thread_index == -1)) zmalloc_register_thread_index();
    if (unlikely(thread_index >= MAX_THREADS_NUM)) {
        atomic_fetch_add_explicit(&used_memory_for_additional_threads, size, memory_order_relaxed);