Add zmalloc_aligned() and fix SPMC queue buffer alignment

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_cache_aligned(sizeof(spmcCell) * SPMC_QUEUE_SIZE);
     atomic_init(&q->head, 0);
     q->tail = 0;
     q->head_cache = 0;

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<uintptr_t>(q.buffer) % CACHE_LINE_SIZE, 0u);
     }
 
     void TearDown() override {

src/unit/test_zmalloc.cpp

@@ -6,6 +6,7 @@
 
 #include "generated_wrappers.hpp"
 
+#include <cstdint>
 #include <cstdio>
 
 extern "C" {
@@ -48,3 +49,20 @@ TEST_F(ZmallocTest, TestZmallocAllocZeroByteAndFree) {
 
     ASSERT_EQ(zmalloc_used_memory(), used_memory_before);
 }
+
+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);
+    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);
+}

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,17 @@ 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) {
     UNUSED(size);
@@ -187,6 +203,58 @@ void *zmalloc(size_t size) {
     return ptr;
 }
 
+/* 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
+    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;
+#ifdef USE_JEMALLOC
+    int ret = je_posix_memalign(&ptr, CACHE_LINE_SIZE, alloc_size);
+#else
+    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));
+    return ptr;
+#endif
+}
+
 /* Try allocating memory, and return NULL if failed. */
 void *ztrymalloc(size_t size) {
     void *ptr = ztrymalloc_usable_internal(size, NULL);
@@ -376,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) {
@@ -407,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
 
@@ -420,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
 

src/zmalloc.h

@@ -108,11 +108,13 @@
 #define zcalloc valkey_calloc
 #define zrealloc valkey_realloc
 #define zfree valkey_free
+#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);
+__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);