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

Code before:
    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();
}

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

/* Try allocating memory, and return NULL if failed.
 * '*usable' is set to the usable size if non NULL. */
static inline void *ztrymalloc_usable_internal(size_t size, size_t *usable) {
    /* Possible overflow, return NULL, so that the caller can panic or handle a failed allocation. */

Review comment (suggestion):
```suggestion
    if (unlikely(thread_index >= MAX_THREADS_NUM)) {
```

Code after:
#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();
}

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

/* Try allocating memory, and return NULL if failed.
 * '*usable' is set to the usable size if non NULL. */
static inline void *ztrymalloc_usable_internal(size_t size, size_t *usable) {