NUMA bindings support for private memory

cachelib/allocator/CMakeLists.txt

@@ -55,6 +55,7 @@ add_library (cachelib_allocator
     PoolOptimizeStrategy.cpp
     PoolRebalancer.cpp
     PoolResizer.cpp
+    PrivateMemoryManager.cpp
     RebalanceStrategy.cpp
     SlabReleaseStats.cpp
     TempShmMapping.cpp

cachelib/allocator/CacheAllocator-inl.h

@@ -60,6 +60,9 @@ CacheAllocator<CacheTrait>::CacheAllocator(
       tempShm_(type == InitMemType::kNone && isOnShm_
                    ? std::make_unique<TempShmMapping>(config_.getCacheSize())
                    : nullptr),
+      privMemManager_(type == InitMemType::kNone && !isOnShm_
+                          ? std::make_unique<PrivateMemoryManager>()
+                          : nullptr),
       shmManager_(type != InitMemType::kNone
                       ? std::make_unique<ShmManager>(config_.cacheDir,
                                                      config_.isUsingPosixShm())
@@ -123,6 +126,16 @@ ShmSegmentOpts CacheAllocator<CacheTrait>::createShmCacheOpts(TierId tid) {
   return opts;
 }
 
+template <typename CacheTrait>
+PrivateSegmentOpts CacheAllocator<CacheTrait>::createPrivateSegmentOpts(TierId tid) {
+  PrivateSegmentOpts opts;
+  opts.alignment = sizeof(Slab);
+  auto memoryTierConfigs = config_.getMemoryTierConfigs();
+  opts.memBindNumaNodes = memoryTierConfigs[tid].getMemBind();
+
+  return opts;
+}
+
 template <typename CacheTrait>
 size_t CacheAllocator<CacheTrait>::memoryTierSize(TierId tid) const {
   auto partitions = std::accumulate(memoryTierConfigs.begin(), memoryTierConfigs.end(), 0UL,
@@ -134,21 +147,18 @@ size_t CacheAllocator<CacheTrait>::memoryTierSize(TierId tid) const {
 }
 
 template <typename CacheTrait>
-std::vector<std::unique_ptr<MemoryAllocator>>
-CacheAllocator<CacheTrait>::createPrivateAllocator() {
-  std::vector<std::unique_ptr<MemoryAllocator>> allocators;
-
+std::unique_ptr<MemoryAllocator>
+CacheAllocator<CacheTrait>::createPrivateAllocator(TierId tid) {
   if (isOnShm_)
-    allocators.emplace_back(std::make_unique<MemoryAllocator>(
+    return std::make_unique<MemoryAllocator>(
                             getAllocatorConfig(config_),
                             tempShm_->getAddr(),
-                            config_.getCacheSize()));
+                            memoryTierSize(tid));
   else
-    allocators.emplace_back(std::make_unique<MemoryAllocator>(
+    return std::make_unique<MemoryAllocator>(
                             getAllocatorConfig(config_),
-                            config_.getCacheSize()));
-
-  return allocators;
+                            privMemManager_->createMapping(config_.size, createPrivateSegmentOpts(tid)),
+                            memoryTierSize(tid));
 }
 
 template <typename CacheTrait>
@@ -177,6 +187,16 @@ CacheAllocator<CacheTrait>::restoreMemoryAllocator(TierId tid) {
       config_.disableFullCoredump);
 }
 
+template <typename CacheTrait>
+std::vector<std::unique_ptr<MemoryAllocator>>
+CacheAllocator<CacheTrait>::createPrivateAllocators() {
+  std::vector<std::unique_ptr<MemoryAllocator>> allocators;
+  for (int tid = 0; tid < getNumTiers(); tid++) {
+    allocators.emplace_back(createPrivateAllocator(tid));
+  }
+  return allocators;
+}
+
 template <typename CacheTrait>
 std::vector<std::unique_ptr<MemoryAllocator>>
 CacheAllocator<CacheTrait>::createAllocators() {
@@ -310,7 +330,7 @@ std::vector<std::unique_ptr<MemoryAllocator>>
 CacheAllocator<CacheTrait>::initAllocator(
     InitMemType type) {
   if (type == InitMemType::kNone) {
-    return createPrivateAllocator();
+    return createPrivateAllocators();
   } else if (type == InitMemType::kMemNew) {
     return createAllocators();
   } else if (type == InitMemType::kMemAttach) {

cachelib/allocator/CacheAllocator.h

@@ -60,6 +60,7 @@
 #include "cachelib/allocator/PoolOptimizer.h"
 #include "cachelib/allocator/PoolRebalancer.h"
 #include "cachelib/allocator/PoolResizer.h"
+#include "cachelib/allocator/PrivateMemoryManager.h"
 #include "cachelib/allocator/ReadOnlySharedCacheView.h"
 #include "cachelib/allocator/Reaper.h"
 #include "cachelib/allocator/RebalanceStrategy.h"
@@ -2155,6 +2156,8 @@ auto& mmContainer = getMMContainer(tid, pid, cid);
                   std::chrono::seconds timeout = std::chrono::seconds{0});
 
   ShmSegmentOpts createShmCacheOpts(TierId tid);
+  PrivateSegmentOpts createPrivateSegmentOpts(TierId tid);
+  std::unique_ptr<MemoryAllocator> createPrivateAllocator(TierId tid);
   std::unique_ptr<MemoryAllocator> createNewMemoryAllocator(TierId tid);
   std::unique_ptr<MemoryAllocator> restoreMemoryAllocator(TierId tid);
   std::unique_ptr<CCacheManager> restoreCCacheManager(TierId tid);
@@ -2204,7 +2207,7 @@ auto& mmContainer = getMMContainer(tid, pid, cid);
   // @throw std::runtime_error if type is invalid
   std::vector<std::unique_ptr<MemoryAllocator>> initAllocator(InitMemType type);
 
-  std::vector<std::unique_ptr<MemoryAllocator>> createPrivateAllocator();
+  std::vector<std::unique_ptr<MemoryAllocator>> createPrivateAllocators();
   std::vector<std::unique_ptr<MemoryAllocator>> createAllocators();
   std::vector<std::unique_ptr<MemoryAllocator>> restoreAllocators();
 
@@ -2373,6 +2376,8 @@ auto& mmContainer = getMMContainer(tid, pid, cid);
   // is not persisted when cache process exits.
   std::unique_ptr<TempShmMapping> tempShm_;
 
+  std::unique_ptr<PrivateMemoryManager> privMemManager_;
+
   std::unique_ptr<ShmManager> shmManager_;
 
   // Deserialize data to restore cache allocator. Used only while attaching to

cachelib/allocator/MemoryTierCacheConfig.h

@@ -16,11 +16,14 @@
 
 #pragma once
 
+#include "cachelib/common/Utils.h"
 #include "cachelib/shm/ShmCommon.h"
 
 namespace facebook {
 namespace cachelib {
 class MemoryTierCacheConfig {
+  using bitmask_type = util::NumaBitMask;
+
  public:
   // Creates instance of MemoryTierCacheConfig for Posix/SysV Shared memory.
   static MemoryTierCacheConfig fromShm() { return MemoryTierCacheConfig(); }
@@ -39,12 +42,12 @@ class MemoryTierCacheConfig {
   size_t getRatio() const noexcept { return ratio; }
 
   // Allocate memory only from specified NUMA nodes
-  MemoryTierCacheConfig& setMemBind(const NumaBitMask& _numaNodes) {
+  MemoryTierCacheConfig& setMemBind(const bitmask_type& _numaNodes) {
     numaNodes = _numaNodes;
     return *this;
   }
 
-  const NumaBitMask& getMemBind() const noexcept { return numaNodes; }
+  const bitmask_type& getMemBind() const noexcept { return numaNodes; }
 
   size_t calculateTierSize(size_t totalCacheSize, size_t partitionNum) const {
     // TODO: Call this method when tiers are enabled in allocator
@@ -71,7 +74,7 @@ class MemoryTierCacheConfig {
   size_t ratio{1};
 
   // Numa node(s) to bind the tier
-  NumaBitMask numaNodes;
+  bitmask_type numaNodes;
 
   // TODO: introduce a container for tier settings when adding support for
   // file-mapped memory

cachelib/allocator/PrivateMemoryManager.cpp

@@ -0,0 +1,50 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "cachelib/allocator/PrivateMemoryManager.h"
+
+#include <folly/ScopeGuard.h>
+
+namespace facebook {
+namespace cachelib {
+
+PrivateMemoryManager::~PrivateMemoryManager() {
+  for (auto& entry : mappings) {
+    util::munmapMemory(entry.first, entry.second);
+  }
+}
+
+void* PrivateMemoryManager::createMapping(size_t size,
+                                          PrivateSegmentOpts opts) {
+  void* addr = util::mmapAlignedZeroedMemory(opts.alignment, size);
+  auto guard = folly::makeGuard([&]() {
+    util::munmapMemory(addr, size);
+    mappings.erase(addr);
+  });
+
+  XDCHECK_EQ(reinterpret_cast<uint64_t>(addr) & (opts.alignment - 1), 0ULL);
+
+  if (!opts.memBindNumaNodes.empty()) {
+    util::mbindMemory(addr, size, MPOL_BIND, opts.memBindNumaNodes, 0);
+  }
+
+  mappings.emplace(addr, size);
+
+  guard.dismiss();
+  return addr;
+}
+} // namespace cachelib
+} // namespace facebook

cachelib/allocator/PrivateMemoryManager.h

@@ -0,0 +1,44 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <cstddef>
+#include <unordered_map>
+
+#include "cachelib/common/Utils.h"
+
+namespace facebook {
+namespace cachelib {
+
+struct PrivateSegmentOpts {
+  size_t alignment{1}; // alignment for mapping.
+  util::NumaBitMask memBindNumaNodes;
+};
+
+class PrivateMemoryManager {
+ public:
+  PrivateMemoryManager() {}
+  ~PrivateMemoryManager();
+
+  void* createMapping(size_t size, PrivateSegmentOpts opts);
+
+ private:
+  std::unordered_map<void*, size_t> mappings;
+};
+
+} // namespace cachelib
+} // namespace facebook

cachelib/cachebench/util/CacheConfig.h

@@ -52,7 +52,7 @@ struct MemoryTierConfig : public JSONConfig {
   MemoryTierCacheConfig getMemoryTierCacheConfig() {
     MemoryTierCacheConfig config = MemoryTierCacheConfig::fromShm();
     config.setRatio(ratio);
-    config.setMemBind(NumaBitMask(memBindNodes));
+    config.setMemBind(util::NumaBitMask(memBindNodes));
     return config;
   }
 

cachelib/common/CMakeLists.txt

@@ -39,6 +39,7 @@ target_link_libraries(cachelib_common PUBLIC
   Folly::folly_exception_tracer
   Folly::folly_exception_tracer_base
   Folly::folly_exception_counter
+  numa
 )
 
 install(TARGETS cachelib_common

cachelib/common/Utils.cpp

@@ -16,6 +16,7 @@
 
 #include <dirent.h>
 #include <folly/experimental/exception_tracer/ExceptionTracer.h>
+#include <numaif.h>
 #include <sys/mman.h>
 #include <sys/resource.h>
 #include <sys/shm.h>
@@ -181,6 +182,22 @@ void* mmapAlignedZeroedMemory(size_t alignment,
   throw std::system_error(errno, std::system_category(), "Cannot mmap");
 }
 
+void munmapMemory(void* addr, size_t size) { munmap(addr, size); }
+
+void mbindMemory(void* addr,
+                 unsigned long len,
+                 int mode,
+                 const NumaBitMask& mask,
+                 unsigned int flags) {
+  auto nodesMask = mask.getNativeBitmask();
+
+  long ret = mbind(addr, len, mode, nodesMask->maskp, nodesMask->size, flags);
+  if (ret != 0) {
+    util::throwSystemError(
+        errno, folly::sformat("mbind() failed: {}", std::strerror(errno)));
+  }
+}
+
 void setMaxLockMemory(uint64_t bytes) {
   struct rlimit rlim {
     bytes, bytes