From 6dbf797cbe930b6cb28304c05773229ef1196030 Mon Sep 17 00:00:00 2001 From: Sarthak Aggarwal Date: Thu, 16 Apr 2026 22:26:39 +0000 Subject: [PATCH 1/4] Add zmalloc_aligned() and fix SPMC queue buffer alignment Introduce zmalloc_aligned(alignment, size) for cache-line-aligned heap allocations with full Valkey memory tracking. On HAVE_MALLOC_SIZE platforms it wraps posix_memalign directly. On !HAVE_MALLOC_SIZE platforms it over-allocates and stores the raw pointer and a flag-tagged size before the aligned region so zfree() can transparently recover the original pointer. Use zmalloc_aligned() in spmcInit() to guarantee each spmcCell is cache-line aligned, fixing UBSan failures in sanitizer CI builds which force MALLOC=libc. Signed-off-by: Sarthak Aggarwal --- src/queues.c | 2 +- src/unit/test_queues.cpp | 2 + src/unit/test_zmalloc.cpp | 14 ++++++ src/zmalloc.c | 103 ++++++++++++++++++++++++++++++++++++-- src/zmalloc.h | 4 ++ 5 files changed, 119 insertions(+), 6 deletions(-) diff --git a/src/queues.c b/src/queues.c index e35ec5fc05f..63b30d5df85 100644 --- a/src/queues.c +++ b/src/queues.c @@ -104,7 +104,7 @@ inline size_t mpscDequeueBatch(mpscQueue *q, void **jobs_out, size_t max_jobs) { * ========================================================================== */ inline void spmcInit(spmcQueue *q) { - q->buffer = (spmcCell *)zmalloc(sizeof(spmcCell) * SPMC_QUEUE_SIZE); + q->buffer = (spmcCell *)zmalloc_aligned(CACHE_LINE_SIZE, sizeof(spmcCell) * SPMC_QUEUE_SIZE); atomic_init(&q->head, 0); q->tail = 0; q->head_cache = 0; diff --git a/src/unit/test_queues.cpp b/src/unit/test_queues.cpp index 8e9355c611b..05b7687a6ee 100644 --- a/src/unit/test_queues.cpp +++ b/src/unit/test_queues.cpp @@ -148,6 +148,8 @@ class SpmcQueueTest : public ::testing::Test { void SetUp() override { spmcInit(&q); + ASSERT_NE(q.buffer, nullptr); + EXPECT_EQ(reinterpret_cast(q.buffer) % CACHE_LINE_SIZE, 0u); } void TearDown() override { diff --git a/src/unit/test_zmalloc.cpp b/src/unit/test_zmalloc.cpp index a1abe963360..ea322577a82 100644 --- a/src/unit/test_zmalloc.cpp +++ b/src/unit/test_zmalloc.cpp @@ -6,6 +6,7 @@ #include "generated_wrappers.hpp" +#include #include extern "C" { @@ -48,3 +49,16 @@ TEST_F(ZmallocTest, TestZmallocAllocZeroByteAndFree) { ASSERT_EQ(zmalloc_used_memory(), used_memory_before); } + +TEST_F(ZmallocTest, TestZmallocAlignedAllocAndFree) { + size_t used_memory_before = zmalloc_used_memory(); + void *ptr = zmalloc_aligned(CACHE_LINE_SIZE, 123); + + ASSERT_NE(ptr, nullptr); + EXPECT_EQ(reinterpret_cast(ptr) % CACHE_LINE_SIZE, 0u); + EXPECT_GE(zmalloc_usable_size(ptr), 123u); + + zfree(ptr); + + ASSERT_EQ(zmalloc_used_memory(), used_memory_before); +} diff --git a/src/zmalloc.c b/src/zmalloc.c index 341da4d2341..bf954b120f3 100644 --- a/src/zmalloc.c +++ b/src/zmalloc.c @@ -72,6 +72,11 @@ void zlibc_free(void *ptr) { */ #define MALLOC_MIN_SIZE(x) ((x) > 0 ? (x) : sizeof(long)) +#ifndef HAVE_MALLOC_SIZE +#define ZMALLOC_ALIGNED_FLAG ((size_t)1 << (sizeof(size_t) * 8 - 1)) +#define ZMALLOC_ALIGNED_SIZE_MASK (~ZMALLOC_ALIGNED_FLAG) +#endif + /* Explicitly override malloc/free etc when using tcmalloc. */ #if defined(USE_TCMALLOC) #define malloc(size) tc_malloc(size) @@ -141,6 +146,30 @@ static void zmalloc_default_oom(size_t size) { static void (*zmalloc_oom_handler)(size_t) = zmalloc_default_oom; +static inline void zmallocValidateAlignment(size_t alignment) { + if (alignment < sizeof(void *) || (alignment & (alignment - 1)) != 0) { + panic("zmalloc_aligned alignment must be a power of two and at least sizeof(void *)"); + } +} + +#ifdef HAVE_MALLOC_SIZE +static inline int zmallocPosixMemalign(void **memptr, size_t alignment, size_t size) { +#ifdef USE_JEMALLOC + return je_posix_memalign(memptr, alignment, size); +#else + return posix_memalign(memptr, alignment, size); +#endif +} +#else +static inline int zmallocIsAlignedAllocation(size_t size) { + return (size & ZMALLOC_ALIGNED_FLAG) != 0; +} + +static inline size_t zmallocAlignedDataSize(size_t size) { + return size & ZMALLOC_ALIGNED_SIZE_MASK; +} +#endif + #ifdef HAVE_MALLOC_SIZE void *extend_to_usable(void *ptr, size_t size) { UNUSED(size); @@ -187,6 +216,58 @@ void *zmalloc(size_t size) { return ptr; } +/* Allocate aligned memory or panic. + * The returned pointer can be freed with zfree(). */ +void *zmalloc_aligned(size_t alignment, size_t size) { + size_t alloc_size = MALLOC_MIN_SIZE(size); + + zmallocValidateAlignment(alignment); + + if (alloc_size >= SIZE_MAX / 2) { + zmalloc_oom_handler(size); + return NULL; + } + +#ifdef HAVE_MALLOC_SIZE + void *ptr = NULL; + int ret = zmallocPosixMemalign(&ptr, alignment, alloc_size); + if (ret != 0 || !ptr) zmalloc_oom_handler(size); + update_zmalloc_stat_alloc(zmalloc_size(ptr)); + return ptr; +#else + if (alloc_size & ZMALLOC_ALIGNED_FLAG) { + zmalloc_oom_handler(size); + return NULL; + } + + size_t extra = alignment - 1; + if (extra > SIZE_MAX - PREFIX_SIZE - sizeof(void *)) { + zmalloc_oom_handler(size); + return NULL; + } + extra += PREFIX_SIZE + sizeof(void *); + if (alloc_size > SIZE_MAX - extra) { + zmalloc_oom_handler(size); + return NULL; + } + + unsigned char *raw = malloc(alloc_size + extra); + if (!raw) zmalloc_oom_handler(size); + + uintptr_t aligned = + ((uintptr_t)(raw + sizeof(void *) + PREFIX_SIZE + alignment - 1)) & ~((uintptr_t)alignment - 1); + void *ptr = (void *)aligned; + + /* Store the original allocation so zfree() can release the right pointer + * while keeping the public pointer aligned. */ + *((void **)((unsigned char *)ptr - PREFIX_SIZE - sizeof(void *))) = raw; + *((size_t *)((unsigned char *)ptr - PREFIX_SIZE)) = alloc_size | ZMALLOC_ALIGNED_FLAG; + + update_zmalloc_stat_alloc(alloc_size + PREFIX_SIZE); + return ptr; +#endif +} + /* Try allocating memory, and return NULL if failed. */ void *ztrymalloc(size_t size) { void *ptr = ztrymalloc_usable_internal(size, NULL); @@ -374,8 +455,8 @@ void *zrealloc_usable(void *ptr, size_t size, size_t *usable) { * information as the first bytes of every allocation. */ #ifndef HAVE_MALLOC_SIZE size_t zmalloc_size(void *ptr) { - void *realptr = (char *)ptr - PREFIX_SIZE; - size_t size = *((size_t *)realptr); + size_t size = *((size_t *)((char *)ptr - PREFIX_SIZE)); + size = zmallocAlignedDataSize(size); return size + PREFIX_SIZE; } size_t zmalloc_usable_size(void *ptr) { @@ -407,8 +488,15 @@ void zfree(void *ptr) { #ifdef HAVE_MALLOC_SIZE size_t size = zmalloc_size(ptr); #else - ptr = (char *)ptr - PREFIX_SIZE; - size_t data_size = *((size_t *)ptr); + unsigned char *prefix = (unsigned char *)ptr - PREFIX_SIZE; + size_t data_size = *((size_t *)prefix); + if (zmallocIsAlignedAllocation(data_size)) { + size_t size = zmallocAlignedDataSize(data_size) + PREFIX_SIZE; + void *raw = *((void **)(prefix - sizeof(void *))); + zfree_internal(raw, size); + return; + } + ptr = prefix; size_t size = data_size + PREFIX_SIZE; #endif @@ -420,7 +508,12 @@ void zfree_with_size(void *ptr, size_t size) { if (ptr == NULL) return; #ifndef HAVE_MALLOC_SIZE - ptr = (char *)ptr - PREFIX_SIZE; + unsigned char *prefix = (unsigned char *)ptr - PREFIX_SIZE; + if (zmallocIsAlignedAllocation(*((size_t *)prefix))) { + ptr = *((void **)(prefix - sizeof(void *))); + } else { + ptr = prefix; + } size += PREFIX_SIZE; #endif diff --git a/src/zmalloc.h b/src/zmalloc.h index db260bb51db..da6c79cf18f 100644 --- a/src/zmalloc.h +++ b/src/zmalloc.h @@ -108,11 +108,15 @@ #define zcalloc valkey_calloc #define zrealloc valkey_realloc #define zfree valkey_free +#define zmalloc_aligned valkey_malloc_aligned /* 'noinline' attribute is intended to prevent the `-Wstringop-overread` warning * when using gcc-12 later with LTO enabled. It may be removed once the * bug[https://gcc.gnu.org/bugzilla/show_bug.cgi?id=96503] is fixed. */ __attribute__((malloc, alloc_size(1), noinline)) void *zmalloc(size_t size); +/* Allocate memory with the requested alignment. The returned pointer can be freed + * with zfree(). */ +__attribute__((malloc, alloc_size(2), noinline)) void *zmalloc_aligned(size_t alignment, size_t size); __attribute__((malloc, alloc_size(1), noinline)) void *zcalloc(size_t size); __attribute__((malloc, alloc_size(1, 2), noinline)) void *zcalloc_num(size_t num, size_t size); __attribute__((alloc_size(2), noinline)) void *zrealloc(void *ptr, size_t size); From bb4f3451a581f36dd80aa1fe168d31be512404df Mon Sep 17 00:00:00 2001 From: Sarthak Aggarwal Date: Thu, 16 Apr 2026 17:39:56 -0700 Subject: [PATCH 2/4] Simplify zmalloc_aligned to zmalloc_cache_aligned Replace the general-purpose zmalloc_aligned(alignment, size) with a simpler zmalloc_cache_aligned(size) that always aligns to CACHE_LINE_SIZE. This is the only alignment any caller needs, so drop: - The !HAVE_MALLOC_SIZE manual alignment path (MSB flag, raw pointer storage, modified zfree/zfree_with_size/zmalloc_size) - The alignment validation helper - The posix_memalign wrapper The result is ~70 fewer lines with the same practical effect. On !HAVE_MALLOC_SIZE platforms (none that Valkey ships on), it panics at startup rather than silently corrupting memory. Signed-off-by: Sarthak Aggarwal --- src/queues.c | 2 +- src/unit/test_zmalloc.cpp | 4 +- src/zmalloc.c | 109 +++++++------------------------------- src/zmalloc.h | 8 +-- 4 files changed, 27 insertions(+), 96 deletions(-) diff --git a/src/queues.c b/src/queues.c index 63b30d5df85..d2146d5d9a9 100644 --- a/src/queues.c +++ b/src/queues.c @@ -104,7 +104,7 @@ inline size_t mpscDequeueBatch(mpscQueue *q, void **jobs_out, size_t max_jobs) { * ========================================================================== */ inline void spmcInit(spmcQueue *q) { - q->buffer = (spmcCell *)zmalloc_aligned(CACHE_LINE_SIZE, sizeof(spmcCell) * SPMC_QUEUE_SIZE); + q->buffer = (spmcCell *)zmalloc_cache_aligned(sizeof(spmcCell) * SPMC_QUEUE_SIZE); atomic_init(&q->head, 0); q->tail = 0; q->head_cache = 0; diff --git a/src/unit/test_zmalloc.cpp b/src/unit/test_zmalloc.cpp index ea322577a82..ed5655530c9 100644 --- a/src/unit/test_zmalloc.cpp +++ b/src/unit/test_zmalloc.cpp @@ -50,9 +50,9 @@ TEST_F(ZmallocTest, TestZmallocAllocZeroByteAndFree) { ASSERT_EQ(zmalloc_used_memory(), used_memory_before); } -TEST_F(ZmallocTest, TestZmallocAlignedAllocAndFree) { +TEST_F(ZmallocTest, TestZmallocCacheAlignedAllocAndFree) { size_t used_memory_before = zmalloc_used_memory(); - void *ptr = zmalloc_aligned(CACHE_LINE_SIZE, 123); + void *ptr = zmalloc_cache_aligned(123); ASSERT_NE(ptr, nullptr); EXPECT_EQ(reinterpret_cast(ptr) % CACHE_LINE_SIZE, 0u); diff --git a/src/zmalloc.c b/src/zmalloc.c index bf954b120f3..1ea6042d0b4 100644 --- a/src/zmalloc.c +++ b/src/zmalloc.c @@ -72,11 +72,6 @@ void zlibc_free(void *ptr) { */ #define MALLOC_MIN_SIZE(x) ((x) > 0 ? (x) : sizeof(long)) -#ifndef HAVE_MALLOC_SIZE -#define ZMALLOC_ALIGNED_FLAG ((size_t)1 << (sizeof(size_t) * 8 - 1)) -#define ZMALLOC_ALIGNED_SIZE_MASK (~ZMALLOC_ALIGNED_FLAG) -#endif - /* Explicitly override malloc/free etc when using tcmalloc. */ #if defined(USE_TCMALLOC) #define malloc(size) tc_malloc(size) @@ -146,29 +141,6 @@ static void zmalloc_default_oom(size_t size) { static void (*zmalloc_oom_handler)(size_t) = zmalloc_default_oom; -static inline void zmallocValidateAlignment(size_t alignment) { - if (alignment < sizeof(void *) || (alignment & (alignment - 1)) != 0) { - panic("zmalloc_aligned alignment must be a power of two and at least sizeof(void *)"); - } -} - -#ifdef HAVE_MALLOC_SIZE -static inline int zmallocPosixMemalign(void **memptr, size_t alignment, size_t size) { -#ifdef USE_JEMALLOC - return je_posix_memalign(memptr, alignment, size); -#else - return posix_memalign(memptr, alignment, size); -#endif -} -#else -static inline int zmallocIsAlignedAllocation(size_t size) { - return (size & ZMALLOC_ALIGNED_FLAG) != 0; -} - -static inline size_t zmallocAlignedDataSize(size_t size) { - return size & ZMALLOC_ALIGNED_SIZE_MASK; -} -#endif #ifdef HAVE_MALLOC_SIZE void *extend_to_usable(void *ptr, size_t size) { @@ -216,55 +188,26 @@ void *zmalloc(size_t size) { return ptr; } -/* Allocate aligned memory or panic. +/* Allocate CACHE_LINE_SIZE-aligned memory or panic. * The returned pointer can be freed with zfree(). */ -void *zmalloc_aligned(size_t alignment, size_t size) { - size_t alloc_size = MALLOC_MIN_SIZE(size); - - zmallocValidateAlignment(alignment); - - if (alloc_size >= SIZE_MAX / 2) { - zmalloc_oom_handler(size); - return NULL; - } - -#ifdef HAVE_MALLOC_SIZE +void *zmalloc_cache_aligned(size_t size) { +#ifndef HAVE_MALLOC_SIZE + /* The PREFIX_SIZE header that zmalloc uses on !HAVE_MALLOC_SIZE platforms + * would break the alignment guarantee. Every platform Valkey ships on + * (Linux/glibc, macOS, FreeBSD) defines HAVE_MALLOC_SIZE. */ + panic("zmalloc_cache_aligned requires HAVE_MALLOC_SIZE"); + return NULL; +#else void *ptr = NULL; - int ret = zmallocPosixMemalign(&ptr, alignment, alloc_size); + size = MALLOC_MIN_SIZE(size); +#ifdef USE_JEMALLOC + int ret = je_posix_memalign(&ptr, CACHE_LINE_SIZE, size); +#else + int ret = posix_memalign(&ptr, CACHE_LINE_SIZE, size); +#endif if (ret != 0 || !ptr) zmalloc_oom_handler(size); update_zmalloc_stat_alloc(zmalloc_size(ptr)); return ptr; -#else - if (alloc_size & ZMALLOC_ALIGNED_FLAG) { - zmalloc_oom_handler(size); - return NULL; - } - - size_t extra = alignment - 1; - if (extra > SIZE_MAX - PREFIX_SIZE - sizeof(void *)) { - zmalloc_oom_handler(size); - return NULL; - } - extra += PREFIX_SIZE + sizeof(void *); - if (alloc_size > SIZE_MAX - extra) { - zmalloc_oom_handler(size); - return NULL; - } - - unsigned char *raw = malloc(alloc_size + extra); - if (!raw) zmalloc_oom_handler(size); - - uintptr_t aligned = - ((uintptr_t)(raw + sizeof(void *) + PREFIX_SIZE + alignment - 1)) & ~((uintptr_t)alignment - 1); - void *ptr = (void *)aligned; - - /* Store the original allocation so zfree() can release the right pointer - * while keeping the public pointer aligned. */ - *((void **)((unsigned char *)ptr - PREFIX_SIZE - sizeof(void *))) = raw; - *((size_t *)((unsigned char *)ptr - PREFIX_SIZE)) = alloc_size | ZMALLOC_ALIGNED_FLAG; - - update_zmalloc_stat_alloc(alloc_size + PREFIX_SIZE); - return ptr; #endif } @@ -455,8 +398,8 @@ void *zrealloc_usable(void *ptr, size_t size, size_t *usable) { * information as the first bytes of every allocation. */ #ifndef HAVE_MALLOC_SIZE size_t zmalloc_size(void *ptr) { - size_t size = *((size_t *)((char *)ptr - PREFIX_SIZE)); - size = zmallocAlignedDataSize(size); + void *realptr = (char *)ptr - PREFIX_SIZE; + size_t size = *((size_t *)realptr); return size + PREFIX_SIZE; } size_t zmalloc_usable_size(void *ptr) { @@ -488,15 +431,8 @@ void zfree(void *ptr) { #ifdef HAVE_MALLOC_SIZE size_t size = zmalloc_size(ptr); #else - unsigned char *prefix = (unsigned char *)ptr - PREFIX_SIZE; - size_t data_size = *((size_t *)prefix); - if (zmallocIsAlignedAllocation(data_size)) { - size_t size = zmallocAlignedDataSize(data_size) + PREFIX_SIZE; - void *raw = *((void **)(prefix - sizeof(void *))); - zfree_internal(raw, size); - return; - } - ptr = prefix; + ptr = (char *)ptr - PREFIX_SIZE; + size_t data_size = *((size_t *)ptr); size_t size = data_size + PREFIX_SIZE; #endif @@ -508,12 +444,7 @@ void zfree_with_size(void *ptr, size_t size) { if (ptr == NULL) return; #ifndef HAVE_MALLOC_SIZE - unsigned char *prefix = (unsigned char *)ptr - PREFIX_SIZE; - if (zmallocIsAlignedAllocation(*((size_t *)prefix))) { - ptr = *((void **)(prefix - sizeof(void *))); - } else { - ptr = prefix; - } + ptr = (char *)ptr - PREFIX_SIZE; size += PREFIX_SIZE; #endif diff --git a/src/zmalloc.h b/src/zmalloc.h index da6c79cf18f..7406ea9f758 100644 --- a/src/zmalloc.h +++ b/src/zmalloc.h @@ -108,15 +108,15 @@ #define zcalloc valkey_calloc #define zrealloc valkey_realloc #define zfree valkey_free -#define zmalloc_aligned valkey_malloc_aligned +#define zmalloc_cache_aligned valkey_malloc_cache_aligned /* 'noinline' attribute is intended to prevent the `-Wstringop-overread` warning * when using gcc-12 later with LTO enabled. It may be removed once the * bug[https://gcc.gnu.org/bugzilla/show_bug.cgi?id=96503] is fixed. */ __attribute__((malloc, alloc_size(1), noinline)) void *zmalloc(size_t size); -/* Allocate memory with the requested alignment. The returned pointer can be freed - * with zfree(). */ -__attribute__((malloc, alloc_size(2), noinline)) void *zmalloc_aligned(size_t alignment, size_t size); +/* Allocate CACHE_LINE_SIZE-aligned memory. The returned pointer can be freed + * with zfree(). Only supported on platforms with HAVE_MALLOC_SIZE. */ +__attribute__((malloc, alloc_size(1), noinline)) void *zmalloc_cache_aligned(size_t size); __attribute__((malloc, alloc_size(1), noinline)) void *zcalloc(size_t size); __attribute__((malloc, alloc_size(1, 2), noinline)) void *zcalloc_num(size_t num, size_t size); __attribute__((alloc_size(2), noinline)) void *zrealloc(void *ptr, size_t size); From 7c486f09790c5f5734154831bb85c62d1d8d5c52 Mon Sep 17 00:00:00 2001 From: Sarthak Aggarwal Date: Fri, 17 Apr 2026 03:08:41 +0000 Subject: [PATCH 3/4] Support zmalloc_cache_aligned without malloc_usable_size Signed-off-by: Sarthak Aggarwal --- src/zmalloc.c | 79 ++++++++++++++++++++++++++++++++++++++++++++------- src/zmalloc.h | 2 +- 2 files changed, 69 insertions(+), 12 deletions(-) diff --git a/src/zmalloc.c b/src/zmalloc.c index 1ea6042d0b4..04de73e2f2d 100644 --- a/src/zmalloc.c +++ b/src/zmalloc.c @@ -72,6 +72,11 @@ void zlibc_free(void *ptr) { */ #define MALLOC_MIN_SIZE(x) ((x) > 0 ? (x) : sizeof(long)) +#ifndef HAVE_MALLOC_SIZE +#define ZMALLOC_CACHE_ALIGNED_FLAG ((size_t)1 << (sizeof(size_t) * 8 - 1)) +#define ZMALLOC_CACHE_ALIGNED_SIZE_MASK (~ZMALLOC_CACHE_ALIGNED_FLAG) +#endif + /* Explicitly override malloc/free etc when using tcmalloc. */ #if defined(USE_TCMALLOC) #define malloc(size) tc_malloc(size) @@ -141,6 +146,16 @@ static void zmalloc_default_oom(size_t size) { static void (*zmalloc_oom_handler)(size_t) = zmalloc_default_oom; +#ifndef HAVE_MALLOC_SIZE +static inline int zmallocIsCacheAlignedAllocation(size_t size) { + return (size & ZMALLOC_CACHE_ALIGNED_FLAG) != 0; +} + +static inline size_t zmallocCacheAlignedDataSize(size_t size) { + return size & ZMALLOC_CACHE_ALIGNED_SIZE_MASK; +} +#endif + #ifdef HAVE_MALLOC_SIZE void *extend_to_usable(void *ptr, size_t size) { @@ -191,19 +206,48 @@ void *zmalloc(size_t size) { /* Allocate CACHE_LINE_SIZE-aligned memory or panic. * The returned pointer can be freed with zfree(). */ void *zmalloc_cache_aligned(size_t size) { + size_t alloc_size = MALLOC_MIN_SIZE(size); + + if (alloc_size >= SIZE_MAX / 2) { + zmalloc_oom_handler(size); + return NULL; + } + #ifndef HAVE_MALLOC_SIZE - /* The PREFIX_SIZE header that zmalloc uses on !HAVE_MALLOC_SIZE platforms - * would break the alignment guarantee. Every platform Valkey ships on - * (Linux/glibc, macOS, FreeBSD) defines HAVE_MALLOC_SIZE. */ - panic("zmalloc_cache_aligned requires HAVE_MALLOC_SIZE"); - return NULL; + if (alloc_size & ZMALLOC_CACHE_ALIGNED_FLAG) { + zmalloc_oom_handler(size); + return NULL; + } + + size_t extra = CACHE_LINE_SIZE - 1; + if (extra > SIZE_MAX - PREFIX_SIZE - sizeof(void *)) { + zmalloc_oom_handler(size); + return NULL; + } + extra += PREFIX_SIZE + sizeof(void *); + if (alloc_size > SIZE_MAX - extra) { + zmalloc_oom_handler(size); + return NULL; + } + + unsigned char *raw = malloc(alloc_size + extra); + if (!raw) zmalloc_oom_handler(size); + + uintptr_t aligned = + ((uintptr_t)(raw + sizeof(void *) + PREFIX_SIZE + CACHE_LINE_SIZE - 1)) & ~((uintptr_t)CACHE_LINE_SIZE - 1); + void *ptr = (void *)aligned; + + *((void **)((unsigned char *)ptr - PREFIX_SIZE - sizeof(void *))) = raw; + *((size_t *)((unsigned char *)ptr - PREFIX_SIZE)) = alloc_size | ZMALLOC_CACHE_ALIGNED_FLAG; + + update_zmalloc_stat_alloc(alloc_size + PREFIX_SIZE); + return ptr; #else void *ptr = NULL; - size = MALLOC_MIN_SIZE(size); #ifdef USE_JEMALLOC - int ret = je_posix_memalign(&ptr, CACHE_LINE_SIZE, size); + int ret = je_posix_memalign(&ptr, CACHE_LINE_SIZE, alloc_size); #else - int ret = posix_memalign(&ptr, CACHE_LINE_SIZE, size); + int ret = posix_memalign(&ptr, CACHE_LINE_SIZE, alloc_size); #endif if (ret != 0 || !ptr) zmalloc_oom_handler(size); update_zmalloc_stat_alloc(zmalloc_size(ptr)); @@ -400,6 +444,7 @@ void *zrealloc_usable(void *ptr, size_t size, size_t *usable) { size_t zmalloc_size(void *ptr) { void *realptr = (char *)ptr - PREFIX_SIZE; size_t size = *((size_t *)realptr); + size = zmallocCacheAlignedDataSize(size); return size + PREFIX_SIZE; } size_t zmalloc_usable_size(void *ptr) { @@ -431,8 +476,15 @@ void zfree(void *ptr) { #ifdef HAVE_MALLOC_SIZE size_t size = zmalloc_size(ptr); #else - ptr = (char *)ptr - PREFIX_SIZE; - size_t data_size = *((size_t *)ptr); + unsigned char *prefix = (unsigned char *)ptr - PREFIX_SIZE; + size_t data_size = *((size_t *)prefix); + if (zmallocIsCacheAlignedAllocation(data_size)) { + size_t size = zmallocCacheAlignedDataSize(data_size) + PREFIX_SIZE; + void *raw = *((void **)(prefix - sizeof(void *))); + zfree_internal(raw, size); + return; + } + ptr = prefix; size_t size = data_size + PREFIX_SIZE; #endif @@ -444,7 +496,12 @@ void zfree_with_size(void *ptr, size_t size) { if (ptr == NULL) return; #ifndef HAVE_MALLOC_SIZE - ptr = (char *)ptr - PREFIX_SIZE; + unsigned char *prefix = (unsigned char *)ptr - PREFIX_SIZE; + if (zmallocIsCacheAlignedAllocation(*((size_t *)prefix))) { + ptr = *((void **)(prefix - sizeof(void *))); + } else { + ptr = prefix; + } size += PREFIX_SIZE; #endif diff --git a/src/zmalloc.h b/src/zmalloc.h index 7406ea9f758..a4f9eecbb49 100644 --- a/src/zmalloc.h +++ b/src/zmalloc.h @@ -115,7 +115,7 @@ * bug[https://gcc.gnu.org/bugzilla/show_bug.cgi?id=96503] is fixed. */ __attribute__((malloc, alloc_size(1), noinline)) void *zmalloc(size_t size); /* Allocate CACHE_LINE_SIZE-aligned memory. The returned pointer can be freed - * with zfree(). Only supported on platforms with HAVE_MALLOC_SIZE. */ + * with zfree(). */ __attribute__((malloc, alloc_size(1), noinline)) void *zmalloc_cache_aligned(size_t size); __attribute__((malloc, alloc_size(1), noinline)) void *zcalloc(size_t size); __attribute__((malloc, alloc_size(1, 2), noinline)) void *zcalloc_num(size_t num, size_t size); From cbd4fb267b0e979f062dfde35386fea68f4f8527 Mon Sep 17 00:00:00 2001 From: Sarthak Aggarwal Date: Fri, 17 Apr 2026 17:21:01 -0700 Subject: [PATCH 4/4] tests: address cache-aligned zmalloc review feedback Signed-off-by: Sarthak Aggarwal --- src/unit/test_zmalloc.cpp | 8 ++++++-- src/zmalloc.h | 2 -- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/unit/test_zmalloc.cpp b/src/unit/test_zmalloc.cpp index ed5655530c9..8b987098ba8 100644 --- a/src/unit/test_zmalloc.cpp +++ b/src/unit/test_zmalloc.cpp @@ -53,12 +53,16 @@ TEST_F(ZmallocTest, TestZmallocAllocZeroByteAndFree) { TEST_F(ZmallocTest, TestZmallocCacheAlignedAllocAndFree) { size_t used_memory_before = zmalloc_used_memory(); void *ptr = zmalloc_cache_aligned(123); + uintptr_t alignment; + size_t usable_size; ASSERT_NE(ptr, nullptr); - EXPECT_EQ(reinterpret_cast(ptr) % CACHE_LINE_SIZE, 0u); - EXPECT_GE(zmalloc_usable_size(ptr), 123u); + alignment = (uintptr_t)ptr % CACHE_LINE_SIZE; + usable_size = zmalloc_usable_size(ptr); zfree(ptr); + ASSERT_EQ(alignment, 0u); + ASSERT_GE(usable_size, 123u); ASSERT_EQ(zmalloc_used_memory(), used_memory_before); } diff --git a/src/zmalloc.h b/src/zmalloc.h index a4f9eecbb49..c1da025bc43 100644 --- a/src/zmalloc.h +++ b/src/zmalloc.h @@ -114,8 +114,6 @@ * when using gcc-12 later with LTO enabled. It may be removed once the * bug[https://gcc.gnu.org/bugzilla/show_bug.cgi?id=96503] is fixed. */ __attribute__((malloc, alloc_size(1), noinline)) void *zmalloc(size_t size); -/* Allocate CACHE_LINE_SIZE-aligned memory. The returned pointer can be freed - * with zfree(). */ __attribute__((malloc, alloc_size(1), noinline)) void *zmalloc_cache_aligned(size_t size); __attribute__((malloc, alloc_size(1), noinline)) void *zcalloc(size_t size); __attribute__((malloc, alloc_size(1, 2), noinline)) void *zcalloc_num(size_t num, size_t size);