[Store] add HugePage support (#1300)

* [Store] add HugePage support

Co-authored-by: Teng Ma <sima.mt@alibaba-inc.com>
Signed-off-by: Xingrui Yi <yixingrui@linux.alibaba.com>

* doc: add huge page env introduction

Signed-off-by: Xingrui Yi <yixingrui@linux.alibaba.com>

---------

Signed-off-by: Xingrui Yi <yixingrui@linux.alibaba.com>
Co-authored-by: Teng Ma <sima.mt@alibaba-inc.com>

docs/source/design/mooncake-store.md

@@ -637,6 +637,8 @@ The HTTP metadata server can be configured using the following parameters:
 - MC_STORE_MEMCPY: Enables or disables local memcpy optimization, set to 1/true to enable, 0/false to disable.
 - MC_STORE_CLIENT_METRIC: Enables client metric reporting, enabled by default; set to 0/false to disable.
 - MC_STORE_CLIENT_METRIC_INTERVAL: Reporting interval in seconds, default 0 (collects but does not report).
+- MC_STORE_USE_HUGEPAGE: Enables huge page support, disabled by default.
+- MC_STORE_HUGEPAGE_SIZE: Specifies the page size of the huge page to use, default 2M.
 #### Usage Example
 To start the master service with the HTTP metadata server enabled:
 ```bash

docs/source/zh_archive/mooncake-store.md

@@ -647,6 +647,8 @@ HTTP 元数据服务器可通过以下参数进行配置：
 - **MC_STORE_MEMCPY**: 控制是否启用本地 memcpy 优化, 1/true 启用, 0/false 禁用
 - **MC_STORE_CLIENT_METRIC**: 启用客户端指标上报, 默认启用；设为 0/false 禁用
 - **MC_STORE_CLIENT_METRIC_INTERVAL**: 指标上报间隔(秒), 默认 0(仅收集不上报)
+- **MC_STORE_USE_HUGEPAGE**: 启用 hugepage 优化, 默认禁用, 设置为 1/true 启用
+- **MC_STORE_HUGEPAGE_SIZE**: hugepage 页大小, 默认 2M
 
 #### 使用示例
 

mooncake-store/include/client_buffer.hpp

@@ -28,7 +28,8 @@ class ClientBufferAllocator
    public:
     // Create for heap-allocated memory
     static std::shared_ptr<ClientBufferAllocator> create(
-        size_t size, const std::string& protocol = "");
+        size_t size, const std::string& protocol = "",
+        bool use_hugepage = false);
 
     // Create for shared memory
     static std::shared_ptr<ClientBufferAllocator> create(
@@ -52,7 +53,8 @@ class ClientBufferAllocator
 
    private:
     // Private constructors for different memory types
-    ClientBufferAllocator(size_t size, const std::string& protocol);
+    ClientBufferAllocator(size_t size, const std::string& protocol,
+                          bool use_hugepage);
     ClientBufferAllocator(void* addr, size_t size, const std::string& protocol);
 
     std::shared_ptr<offset_allocator::OffsetAllocator> allocator_;
@@ -61,6 +63,7 @@ class ClientBufferAllocator
     void* buffer_;
     size_t buffer_size_;
     bool is_external_memory_ = false;
+    bool use_hugepage_ = false;
 };
 
 /**

mooncake-store/include/dummy_client.h

@@ -34,6 +34,8 @@ class ShmHelper {
         return shms_;
     }
 
+    bool is_hugepage() const { return use_hugepage_; }
+
     ShmHelper(const ShmHelper &) = delete;
     ShmHelper &operator=(const ShmHelper &) = delete;
 
@@ -43,6 +45,7 @@ class ShmHelper {
 
     std::vector<std::shared_ptr<ShmSegment>> shms_;
     static std::mutex shm_mutex_;
+    bool use_hugepage_ = false;
 };
 
 class DummyClient : public PyClient {

mooncake-store/include/real_client.h

@@ -428,6 +428,15 @@ class RealClient : public PyClient {
         }
     };
 
+    struct HugepageSegmentDeleter {
+        size_t size = 0;
+        void operator()(void *ptr) const {
+            if (ptr && size > 0) {
+                free_buffer_mmap_memory(ptr, size);
+            }
+        }
+    };
+
     struct AscendSegmentDeleter {
         void operator()(void *ptr) {
             if (ptr) {
@@ -436,12 +445,15 @@ class RealClient : public PyClient {
         }
     };
 
+    std::vector<std::unique_ptr<void, HugepageSegmentDeleter>>
+        hugepage_segment_ptrs_;
     std::vector<std::unique_ptr<void, SegmentDeleter>> segment_ptrs_;
     std::vector<std::unique_ptr<void, AscendSegmentDeleter>>
         ascend_segment_ptrs_;
     std::string protocol;
     std::string device_name;
     std::string local_hostname;
+    bool use_hugepage_ = false;
 
     struct MappedShm {
         std::string shm_name;

mooncake-store/include/utils.h

@@ -2,10 +2,13 @@
 
 #include <cstddef>
 #include <cstdlib>
+#include <linux/memfd.h>
+#include <linux/mman.h>
 #include <string>
 #include <limits>
 #include <ylt/util/tl/expected.hpp>
 
+#include "rpc_types.h"
 #include "types.h"
 
 namespace mooncake {
@@ -119,17 +122,13 @@ std::string expected_to_str(const tl::expected<T, ErrorCode>& expected) {
     return oss.str();
 }
 
-/*
-    @brief Allocates memory for the `BufferAllocator` class.
-    @param total_size The total size of the memory to allocate.
-    @return A pointer to the allocated memory.
-*/
-void* allocate_buffer_allocator_memory(
-    size_t total_size, const std::string& protocol = "",
-    size_t alignment = facebook::cachelib::Slab::kSize);
-
-void free_memory(const std::string& protocol, void* ptr);
+// String utility functions
 
+/**
+ * @brief Convert a byte size to a human-readable string
+ * @param bytes Number of bytes
+ * @return std::string Human-readable string representation of size
+ */
 [[nodiscard]] inline std::string byte_size_to_string(uint64_t bytes) {
     const double KB = 1024.0;
     const double MB = KB * 1024.0;
@@ -230,8 +229,6 @@ void free_memory(const std::string& protocol, void* ptr);
     }
 }
 
-// String utility functions
-
 /**
  * @brief Split a string by delimiter into a vector of strings
  * @param str The string to split
@@ -245,6 +242,93 @@ std::vector<std::string> splitString(const std::string& str,
                                      bool trim_spaces = true,
                                      bool keep_empty = false);
 
+// Buffer allocator functions
+
+constexpr size_t SZ_2MB = 2 * 1024 * 1024;
+constexpr size_t SZ_1GB = 1024 * 1024 * 1024;
+
+/**
+ * @brief Allocates memory for the `BufferAllocator` class.
+ * @param total_size The total size of the memory to allocate.
+ * @return A pointer to the allocated memory.
+ */
+void* allocate_buffer_allocator_memory(
+    size_t total_size, const std::string& protocol = "",
+    size_t alignment = facebook::cachelib::Slab::kSize);
+
+inline size_t align_up(size_t size, size_t alignment) {
+    if (alignment == 0) {
+        return size;
+    }
+    return ((size + alignment - 1) / alignment) * 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.
+ * @return size_t Hugepage size in bytes, or 0 if disabled.
+ */
+[[nodiscard]] inline size_t get_hugepage_size_from_env(
+    unsigned int* out_flags = nullptr, bool use_memfd = false) {
+    const char* use_hp_env = std::getenv("MC_STORE_USE_HUGEPAGE");
+    if (use_hp_env == nullptr) {
+        return 0;
+    }
+
+    constexpr size_t SZ_2MB = 2 * 1024 * 1024;
+    constexpr size_t SZ_1GB = 1024 * 1024 * 1024;
+
+    size_t size = SZ_2MB;  // Default to 2MB
+
+    const char* size_env = std::getenv("MC_STORE_HUGEPAGE_SIZE");
+    if (size_env != nullptr) {
+        size_t parsed_size = string_to_byte_size(size_env);
+
+        if (parsed_size == SZ_2MB || parsed_size == SZ_1GB) {
+            size = parsed_size;
+        } else {
+            LOG(WARNING) << "Invalid MC_STORE_HUGEPAGE_SIZE='" << size_env
+                         << "'. Supported: 2MB, 1GB. Fallback to 2MB.";
+            size = SZ_2MB;
+        }
+    }
+
+    if (out_flags != nullptr) {
+        if (use_memfd) {
+            *out_flags |= MFD_HUGETLB;
+        } else {
+            *out_flags |= MAP_HUGETLB;
+        }
+
+        // Add size specific flag
+        if (size == SZ_2MB) {
+            if (use_memfd) {
+                *out_flags |= MFD_HUGE_2MB;
+            } else {
+                *out_flags |= MAP_HUGE_2MB;
+            }
+        } else if (size == SZ_1GB) {
+            if (use_memfd) {
+                *out_flags |= MFD_HUGE_1GB;
+            } else {
+                *out_flags |= MAP_HUGE_1GB;
+            }
+        }
+        LOG(INFO) << "Using hugepage size: "
+                  << (size == SZ_2MB ? "2MB" : "1GB");
+    }
+
+    return size;
+}
+
+// Hugepage-backed allocation helpers (MAP_HUGETLB + MADV_HUGEPAGE)
+void* allocate_buffer_mmap_memory(size_t total_size, size_t alignment);
+void free_buffer_mmap_memory(void* ptr, size_t total_size);
+
+void free_memory(const std::string& protocol, void* ptr);
+
 // Network utility functions
 
 /**

mooncake-store/src/client_buffer.cpp

@@ -13,9 +13,9 @@
 namespace mooncake {
 
 std::shared_ptr<ClientBufferAllocator> ClientBufferAllocator::create(
-    size_t size, const std::string& protocol) {
+    size_t size, const std::string& protocol, bool use_hugepage) {
     return std::shared_ptr<ClientBufferAllocator>(
-        new ClientBufferAllocator(size, protocol));
+        new ClientBufferAllocator(size, protocol, use_hugepage));
 }
 
 std::shared_ptr<ClientBufferAllocator> ClientBufferAllocator::create(
@@ -25,8 +25,9 @@ std::shared_ptr<ClientBufferAllocator> ClientBufferAllocator::create(
 }
 
 ClientBufferAllocator::ClientBufferAllocator(size_t size,
-                                             const std::string& protocol)
-    : protocol(protocol), buffer_size_(size) {
+                                             const std::string& protocol,
+                                             bool use_hugepage)
+    : protocol(protocol), buffer_size_(size), use_hugepage_(use_hugepage) {
     if (size == 0) {
         buffer_ = nullptr;
         allocator_ = nullptr;
@@ -34,7 +35,11 @@ ClientBufferAllocator::ClientBufferAllocator(size_t size,
     }
     // Align to 64 bytes(cache line size) for better cache performance
     constexpr size_t alignment = 64;
-    buffer_ = allocate_buffer_allocator_memory(size, protocol, alignment);
+    if (use_hugepage_) {
+        buffer_ = allocate_buffer_mmap_memory(size, alignment);
+    } else {
+        buffer_ = allocate_buffer_allocator_memory(size, protocol, alignment);
+    }
     if (!buffer_) {
         throw std::bad_alloc();
     }
@@ -55,7 +60,11 @@ ClientBufferAllocator::ClientBufferAllocator(void* addr, size_t size,
 ClientBufferAllocator::~ClientBufferAllocator() {
     // Free the aligned allocated memory or unmap shared memory
     if (!is_external_memory_ && buffer_) {
-        free_memory(protocol, buffer_);
+        if (use_hugepage_) {
+            free_buffer_mmap_memory(buffer_, buffer_size_);
+        } else {
+            free_memory(protocol, buffer_);
+        }
     }
 }
 

mooncake-store/src/dummy_client.cpp

@@ -6,9 +6,11 @@
 #include <sys/stat.h>  // For S_IRUSR, S_IWUSR
 #include <fcntl.h>     // For O_CREAT, O_RDWR
 #include <unistd.h>    // For ftruncate, close, shm_unlink
+#include <cstdlib>
 
 #include "real_client.h"
 #include "dummy_client.h"
+#include "utils.h"
 #include "utils/scoped_vlog_timer.h"
 #include "rpc_types.h"
 #include "types.h"
@@ -31,7 +33,10 @@ ShmHelper* ShmHelper::getInstance() {
     return &instance;
 }
 
-ShmHelper::ShmHelper() {}
+ShmHelper::ShmHelper() {
+    const char* hp = std::getenv("MC_STORE_USE_HUGEPAGE");
+    use_hugepage_ = (hp != nullptr);
+}
 
 ShmHelper::~ShmHelper() { cleanup(); }
 
@@ -59,10 +64,20 @@ bool ShmHelper::cleanup() {
 void* ShmHelper::allocate(size_t size) {
     std::lock_guard<std::mutex> lock(shm_mutex_);
 
+    unsigned int flags = MFD_CLOEXEC;
+    if (use_hugepage_) {
+        bool use_memfd = true;
+        size = align_up(size, get_hugepage_size_from_env(&flags, use_memfd));
+        LOG(INFO) << "Using huge pages for shared memory, size: " << size;
+    }
+
     // Create memfd
-    int fd = memfd_create_wrapper(MOONCAKE_SHM_NAME, MFD_CLOEXEC);
+    int fd = memfd_create_wrapper(MOONCAKE_SHM_NAME, flags);
     if (fd == -1) {
-        throw std::runtime_error("Failed to create anonymous shared memory: " +
+        std::string extra_msg =
+            use_hugepage_ ? " (Check /proc/sys/vm/nr_hugepages?)" : "";
+        throw std::runtime_error("Failed to create anonymous shared memory" +
+                                 extra_msg + ": " +
                                  std::string(strerror(errno)));
     }
 
@@ -74,8 +89,8 @@ void* ShmHelper::allocate(size_t size) {
     }
 
     // Map memory
-    void* base_addr =
-        mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
+    void* base_addr = mmap(nullptr, size, PROT_READ | PROT_WRITE,
+                           MAP_SHARED | MAP_POPULATE, fd, 0);
     if (base_addr == MAP_FAILED) {
         close(fd);
         throw std::runtime_error("Failed to map shared memory: " +
@@ -413,14 +428,17 @@ int DummyClient::register_buffer(void* buffer, size_t size) {
         LOG(ERROR) << "Buffer is not in any registered shared memory";
         return -1;
     }
+    if (shm_helper_->is_hugepage()) {
+        size = align_up(size, get_hugepage_size_from_env());
+    }
     // Check bounds
     if (reinterpret_cast<uint8_t*>(buffer) !=
             reinterpret_cast<uint8_t*>(shm->base_addr) ||
         size != shm->size) {
-        LOG(ERROR)
-            << "Invalid buffer address or size for registration: Buffer addr: "
-            << buffer << ", need addr: " << shm->base_addr
-            << ", buffer size: " << size << ", need size: " << shm->size;
+        LOG(ERROR) << "Invalid buffer address or size for registration: "
+                      "Buffer addr: "
+                   << buffer << ", need addr: " << shm->base_addr
+                   << ", buffer size: " << size << ", need size: " << shm->size;
         return -1;
     }
 

mooncake-store/src/real_client.cpp

@@ -9,6 +9,8 @@
 #include <stop_token>
 
 #include <cstdlib>  // for atexit
+#include <algorithm>
+#include <cctype>
 #include <optional>
 #include <vector>
 
@@ -143,6 +145,8 @@ void ResourceTracker::startSignalThread() {
 RealClient::RealClient() {
     // Initialize logging severity (leave as before)
     mooncake::init_ylt_log_level();
+    const char *hp = std::getenv("MC_STORE_USE_HUGEPAGE");
+    use_hugepage_ = (hp != nullptr);
 }
 
 RealClient::~RealClient() {
@@ -165,6 +169,8 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
     const std::string &ipc_socket_path, bool enable_offload) {
     this->protocol = protocol;
     this->ipc_socket_path_ = ipc_socket_path;
+    const bool should_use_hugepage =
+        use_hugepage_ && this->protocol != "ascend";
 
     // Remove port if hostname already contains one
     std::string hostname = local_hostname;
@@ -200,8 +206,8 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
     // fail in some rdma implementations.
     // Dummy Client can create shm and share it with Real Client, so Real Client
     // can create client buffer allocator on the shared memory later.
-    client_buffer_allocator_ =
-        ClientBufferAllocator::create(local_buffer_size, this->protocol);
+    client_buffer_allocator_ = ClientBufferAllocator::create(
+        local_buffer_size, this->protocol, should_use_hugepage);
     if (local_buffer_size > 0) {
         LOG(INFO) << "Registering local memory: " << local_buffer_size
                   << " bytes";
@@ -229,18 +235,30 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
         current_glbseg_size += segment_size;
         LOG(INFO) << "Mounting segment: " << segment_size << " bytes, "
                   << current_glbseg_size << " of " << total_glbseg_size;
-        void *ptr =
-            allocate_buffer_allocator_memory(segment_size, this->protocol);
+        size_t mapped_size = segment_size;
+        void *ptr = nullptr;
+        if (should_use_hugepage) {
+            mapped_size = align_up(segment_size, get_hugepage_size_from_env());
+            ptr = allocate_buffer_mmap_memory(mapped_size,
+                                              get_hugepage_size_from_env());
+        } else {
+            ptr =
+                allocate_buffer_allocator_memory(segment_size, this->protocol);
+        }
+
         if (!ptr) {
             LOG(ERROR) << "Failed to allocate segment memory";
             return tl::unexpected(ErrorCode::INVALID_PARAMS);
         }
         if (this->protocol == "ascend") {
             ascend_segment_ptrs_.emplace_back(ptr);
+        } else if (should_use_hugepage) {
+            hugepage_segment_ptrs_.emplace_back(
+                ptr, HugepageSegmentDeleter{mapped_size});
         } else {
             segment_ptrs_.emplace_back(ptr);
         }
-        auto mount_result = client_->MountSegment(ptr, segment_size);
+        auto mount_result = client_->MountSegment(ptr, mapped_size);
         if (!mount_result.has_value()) {
             LOG(ERROR) << "Failed to mount segment: "
                        << toString(mount_result.error());
@@ -324,6 +342,7 @@ tl::expected<void, ErrorCode> RealClient::tearDownAll_internal() {
     client_.reset();
     client_buffer_allocator_.reset();
     port_binder_.reset();
+    hugepage_segment_ptrs_.clear();
     segment_ptrs_.clear();
     local_hostname = "";
     device_name = "";

mooncake-store/src/utils.cpp

@@ -7,7 +7,11 @@
 #include <unistd.h>
 #include <boost/algorithm/string.hpp>
 
+#include <algorithm>
 #include <random>
+#include <cerrno>
+#include <cstring>
+#include <sys/mman.h>
 #ifdef USE_ASCEND_DIRECT
 #include "acl/acl.h"
 #include "config.h"
@@ -136,6 +140,39 @@ void *allocate_buffer_allocator_memory(size_t total_size,
     return aligned_alloc(alignment, total_size);
 }
 
+void *allocate_buffer_mmap_memory(size_t total_size, size_t alignment) {
+    if (total_size == 0) {
+        LOG(ERROR) << "Total size must be greater than 0 for hugepage mmap";
+        return nullptr;
+    }
+
+    unsigned int flags = MAP_PRIVATE | MAP_ANONYMOUS | MAP_POPULATE;
+    const size_t effective_alignment =
+        std::max(alignment, get_hugepage_size_from_env(&flags));
+    const size_t map_size = align_up(total_size, effective_alignment);
+
+    void *ptr = mmap(nullptr, map_size, PROT_READ | PROT_WRITE, flags, -1, 0);
+    if (ptr == MAP_FAILED) {
+        LOG(ERROR) << "Hugepage mmap failed, size=" << map_size
+                   << ", errno=" << errno << " (" << strerror(errno) << ")";
+        return nullptr;
+    }
+
+    return ptr;
+}
+
+void free_buffer_mmap_memory(void *ptr, size_t total_size) {
+    if (!ptr || total_size == 0) {
+        return;
+    }
+
+    const size_t map_size = align_up(total_size, get_hugepage_size_from_env());
+    if (munmap(ptr, map_size) != 0) {
+        LOG(ERROR) << "munmap hugepage failed, size=" << map_size
+                   << ", errno=" << errno << " (" << strerror(errno) << ")";
+    }
+}
+
 void free_memory(const std::string &protocol, void *ptr) {
 #ifdef USE_ASCEND_DIRECT
     if (protocol == "ascend") {