[Store] Support 512MB hugepage size for client buffer allocation

docs/source/deployment/mooncake-store-deployment-guide.md

@@ -1202,7 +1202,7 @@ Do not run binaries from before and after checksum support was introduced in the
 | Variable | Default | Description |
 |----------|---------|-------------|
 | `MC_STORE_USE_HUGEPAGE` | unset | Set `1` to request HugeTLB-backed `mmap()` |
-| `MC_STORE_HUGEPAGE_SIZE` | `2MB` | Supported: `2MB`, `1GB` |
+| `MC_STORE_HUGEPAGE_SIZE` | `2MB` | Supported: `2MB`, `512MB`, `1GB` |
 | `MC_MMAP_ARENA_POOL_SIZE` | unset | Pre-allocated arena pool size (e.g., `8gb`). Explicitly set to enable the arena |
 | `MC_DISABLE_MMAP_ARENA` | unset | Disable arena, fall back to per-call `mmap()`. Accepts `1`/`true`/`yes`/`on` (or `0`/`false`/`no`/`off`) |
 

mooncake-store/include/utils.h

@@ -262,9 +262,19 @@ std::string expected_to_str(const tl::expected<T, ErrorCode>& expected) {
 // Buffer allocator functions
 
 constexpr size_t SZ_2MB = 2 * 1024 * 1024;
+constexpr size_t SZ_512MB = 512 * 1024 * 1024;
 constexpr size_t SZ_1GB = 1024 * 1024 * 1024;
 constexpr double BYTES_PER_GIB = static_cast<double>(SZ_1GB);
 
+// 512MiB hugepages (PMD size on arm64 kernels with 64K base pages) are not
+// defined by older glibc/kernel headers; provide fallbacks.
+#ifndef MAP_HUGE_512MB
+#define MAP_HUGE_512MB (29 << 26)  // MAP_HUGE_SHIFT = 26
+#endif
+#ifndef MFD_HUGE_512MB
+#define MFD_HUGE_512MB (29 << 26)  // MFD_HUGE_SHIFT = 26
+#endif
+
 /**
  * @brief Allocates memory for the `BufferAllocator` class.
  * @param total_size The total size of the memory to allocate.
@@ -286,7 +296,7 @@ inline size_t align_up(size_t size, size_t alignment) {
  * @brief Get hugepage size from env and optionally set the corresponding memfd
  * flags.
  * * @param out_flags Optional pointer to an int. If provided,
- * MAP_HUGETLB and MAP_HUGE_2MB/1GB will be OR-ed into it.
+ * MAP_HUGETLB and MAP_HUGE_2MB/512MB/1GB will be OR-ed into it.
  * @return size_t Hugepage size in bytes, or 0 if disabled.
  */
 [[nodiscard]] inline size_t get_hugepage_size_from_env(
@@ -302,11 +312,12 @@ inline size_t align_up(size_t size, size_t alignment) {
     if (size_env != nullptr) {
         size_t parsed_size = string_to_byte_size(size_env);
 
-        if (parsed_size == SZ_2MB || parsed_size == SZ_1GB) {
+        if (parsed_size == SZ_2MB || parsed_size == SZ_512MB ||
+            parsed_size == SZ_1GB) {
             size = parsed_size;
         } else {
             LOG(WARNING) << "Invalid MC_STORE_HUGEPAGE_SIZE='" << size_env
-                         << "'. Supported: 2MB, 1GB. Fallback to 2MB.";
+                         << "'. Supported: 2MB, 512MB, 1GB. Fallback to 2MB.";
             size = SZ_2MB;
         }
     }
@@ -325,6 +336,12 @@ inline size_t align_up(size_t size, size_t alignment) {
             } else {
                 *out_flags |= MAP_HUGE_2MB;
             }
+        } else if (size == SZ_512MB) {
+            if (use_memfd) {
+                *out_flags |= MFD_HUGE_512MB;
+            } else {
+                *out_flags |= MAP_HUGE_512MB;
+            }
         } else if (size == SZ_1GB) {
             if (use_memfd) {
                 *out_flags |= MFD_HUGE_1GB;
@@ -333,7 +350,9 @@ inline size_t align_up(size_t size, size_t alignment) {
             }
         }
         LOG(INFO) << "Using hugepage size: "
-                  << (size == SZ_2MB ? "2MB" : "1GB");
+                  << (size == SZ_2MB     ? "2MB"
+                      : size == SZ_512MB ? "512MB"
+                                         : "1GB");
     }
 
     return size;

mooncake-store/src/mmap_arena.cpp

@@ -76,9 +76,13 @@ bool MmapArena::initialize(size_t pool_size, size_t alignment,
 
     const size_t actual_alignment = std::max(alignment, kMinAlignment);
 
-    // Align pool size to 2MB for huge pages with overflow protection
+    // Align pool size to the configured huge page size (2MB when hugepages
+    // are not explicitly requested) with overflow protection
+    unsigned int hugepage_flags = 0;
+    const size_t hugepage_size = get_hugepage_size_from_env(&hugepage_flags);
+    const size_t page_align = hugepage_size > 0 ? hugepage_size : SZ_2MB;
     size_t aligned_pool_size;
-    if (!safe_align_up(pool_size, SZ_2MB, &aligned_pool_size)) {
+    if (!safe_align_up(pool_size, page_align, &aligned_pool_size)) {
         LOG(ERROR) << "Arena pool size overflow: requested=" << pool_size;
         return false;
     }
@@ -91,7 +95,12 @@ bool MmapArena::initialize(size_t pool_size, size_t alignment,
 
 // Try huge pages for better TLB performance
 #ifdef MAP_HUGETLB
-    flags |= MAP_HUGETLB;
+    if (hugepage_size > 0) {
+        // MC_STORE_USE_HUGEPAGE set: honor the configured page size.
+        flags |= static_cast<int>(hugepage_flags);
+    } else {
+        flags |= MAP_HUGETLB;
+    }
 #endif
 
     void* pool_base =
@@ -111,7 +120,7 @@ bool MmapArena::initialize(size_t pool_size, size_t alignment,
         }
 
         // Retry without huge pages
-        flags &= ~MAP_HUGETLB;
+        flags &= ~(MAP_HUGETLB | (0x3f << 26));  // strip MAP_HUGE_* size bits
         LOG(WARNING) << "Arena hugepage mmap failed for pool_size="
                      << aligned_pool_size << " bytes"
                      << ", errno=" << mmap_errno << " (" << strerror(mmap_errno)

mooncake-store/src/utils.cpp

@@ -425,6 +425,8 @@ void *allocate_buffer_numa_segments(size_t total_size,
     unsigned int flags = MAP_PRIVATE | MAP_ANONYMOUS;
     if (page_size == SZ_2MB) {
         flags |= MAP_HUGETLB | MAP_HUGE_2MB;
+    } else if (page_size == SZ_512MB) {
+        flags |= MAP_HUGETLB | MAP_HUGE_512MB;
     } else if (page_size == SZ_1GB) {
         flags |= MAP_HUGETLB | MAP_HUGE_1GB;
     } else if (page_size != static_cast<size_t>(getpagesize())) {

mooncake-store/tests/utils_test.cpp

@@ -2,6 +2,7 @@
 #include "random.h"
 
 #include <cstdint>
+#include <cstdlib>
 #include <future>
 #include <gtest/gtest.h>
 #include <netinet/in.h>
@@ -268,3 +269,33 @@ TEST(UtilsTest, AutoPortBinderCustomRange) {
     EXPECT_GE(port, 50000);
     EXPECT_LE(port, 50100);
 }
+
+TEST(HugepageSizeEnvTest, Accepts512Mb) {
+    setenv("MC_STORE_USE_HUGEPAGE", "1", 1);
+    setenv("MC_STORE_HUGEPAGE_SIZE", "512MB", 1);
+
+    unsigned int flags = 0;
+    EXPECT_EQ(get_hugepage_size_from_env(&flags), SZ_512MB);
+    EXPECT_TRUE(flags & MAP_HUGETLB);
+    EXPECT_TRUE(flags & MAP_HUGE_512MB);
+
+    flags = 0;
+    EXPECT_EQ(get_hugepage_size_from_env(&flags, /*use_memfd=*/true), SZ_512MB);
+    EXPECT_TRUE(flags & MFD_HUGETLB);
+    EXPECT_TRUE(flags & MFD_HUGE_512MB);
+
+    unsetenv("MC_STORE_HUGEPAGE_SIZE");
+    unsetenv("MC_STORE_USE_HUGEPAGE");
+}
+
+TEST(HugepageSizeEnvTest, FallsBackTo2MbOnInvalidSize) {
+    setenv("MC_STORE_USE_HUGEPAGE", "1", 1);
+    setenv("MC_STORE_HUGEPAGE_SIZE", "256MB", 1);
+
+    unsigned int flags = 0;
+    EXPECT_EQ(get_hugepage_size_from_env(&flags), SZ_2MB);
+    EXPECT_TRUE(flags & MAP_HUGE_2MB);
+
+    unsetenv("MC_STORE_HUGEPAGE_SIZE");
+    unsetenv("MC_STORE_USE_HUGEPAGE");
+}