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

Code before:
#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;
/* 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;

/* 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);
    abort();
}

static void (*zmalloc_oom_handler)(size_t) = zmalloc_default_oom;

#ifdef HAVE_MALLOC_SIZE
void *extend_to_usable(void *ptr, size_t size) {
    UNUSED(size);
    return ptr;
}
#endif

Review comment (suggestion):
Are you assuming threads are lunch and never terminate and relaunched? cause otherwise this can easily reach max over time

Code after:
#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);
    } else {
        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();
    if (unlikely(thread_index >= MAX_THREADS_NUM)) {
        atomic_fetch_sub_explicit(&used_memory_for_additional_threads, size, memory_order_relaxed);
    } else {
        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);
    abort();
}