NUMA bindings support for Shm Segments

cachelib/allocator/CacheAllocator-inl.h

@@ -105,28 +105,37 @@ CacheAllocator<CacheTrait>::~CacheAllocator() {
 }
 
 template <typename CacheTrait>
-std::unique_ptr<MemoryAllocator>
-CacheAllocator<CacheTrait>::createNewMemoryAllocator() {
+ShmSegmentOpts CacheAllocator<CacheTrait>::createShmCacheOpts() {
   ShmSegmentOpts opts;
   opts.alignment = sizeof(Slab);
+  auto memoryTierConfigs = config_.getMemoryTierConfigs();
+  // TODO: we support single tier so far
+  XDCHECK_EQ(memoryTierConfigs.size(), 1ul);
+  opts.memBindNumaNodes = memoryTierConfigs[0].getMemBind();
+
+  return opts;
+}
+
+template <typename CacheTrait>
+std::unique_ptr<MemoryAllocator>
+CacheAllocator<CacheTrait>::createNewMemoryAllocator() {
   return std::make_unique<MemoryAllocator>(
       getAllocatorConfig(config_),
       shmManager_
           ->createShm(detail::kShmCacheName, config_.size,
-                      config_.slabMemoryBaseAddr, opts)
+                      config_.slabMemoryBaseAddr, createShmCacheOpts())
           .addr,
       config_.size);
 }
 
 template <typename CacheTrait>
 std::unique_ptr<MemoryAllocator>
 CacheAllocator<CacheTrait>::restoreMemoryAllocator() {
-  ShmSegmentOpts opts;
-  opts.alignment = sizeof(Slab);
   return std::make_unique<MemoryAllocator>(
       deserializer_->deserialize<MemoryAllocator::SerializationType>(),
       shmManager_
-          ->attachShm(detail::kShmCacheName, config_.slabMemoryBaseAddr, opts)
+          ->attachShm(detail::kShmCacheName, config_.slabMemoryBaseAddr,
+                      createShmCacheOpts())
           .addr,
       config_.size,
       config_.disableFullCoredump);

cachelib/allocator/CacheAllocator.h

@@ -1868,6 +1868,7 @@ class CacheAllocator : public CacheBase {
                   std::unique_ptr<T>& worker,
                   std::chrono::seconds timeout = std::chrono::seconds{0});
 
+  ShmSegmentOpts createShmCacheOpts();
   std::unique_ptr<MemoryAllocator> createNewMemoryAllocator();
   std::unique_ptr<MemoryAllocator> restoreMemoryAllocator();
   std::unique_ptr<CCacheManager> restoreCCacheManager();

cachelib/allocator/CacheAllocatorConfig.h

@@ -210,7 +210,7 @@ class CacheAllocatorConfig {
   CacheAllocatorConfig& configureMemoryTiers(const MemoryTierConfigs& configs);
 
   // Return reference to MemoryTierCacheConfigs.
-  const MemoryTierConfigs& getMemoryTierConfigs();
+  const MemoryTierConfigs& getMemoryTierConfigs() const noexcept;
 
   // This turns on a background worker that periodically scans through the
   // access container and look for expired items and remove them.
@@ -877,7 +877,7 @@ CacheAllocatorConfig<T>& CacheAllocatorConfig<T>::configureMemoryTiers(
 
 template <typename T>
 const typename CacheAllocatorConfig<T>::MemoryTierConfigs&
-CacheAllocatorConfig<T>::getMemoryTierConfigs() {
+CacheAllocatorConfig<T>::getMemoryTierConfigs() const noexcept {
   return memoryTierConfigs;
 }
 

cachelib/allocator/MemoryTierCacheConfig.h

@@ -16,8 +16,6 @@
 
 #pragma once
 
-#include <string>
-
 #include "cachelib/shm/ShmCommon.h"
 
 namespace facebook {
@@ -43,6 +41,16 @@ class MemoryTierCacheConfig {
 
   size_t getRatio() const noexcept { return ratio; }
 
+  // Allocate memory only from specified NUMA nodes
+  MemoryTierCacheConfig& setMemBind(const NumaBitMask& _numaNodes) {
+    numaNodes = _numaNodes;
+    return *this;
+  }
+
+  const NumaBitMask& getMemBind() const noexcept {
+    return numaNodes;
+  }
+
   size_t calculateTierSize(size_t totalCacheSize, size_t partitionNum) {
     // TODO: Call this method when tiers are enabled in allocator
     // to calculate tier sizes in bytes.
@@ -59,14 +67,17 @@ class MemoryTierCacheConfig {
     return getRatio() * (totalCacheSize / partitionNum);
   }
 
+ private:
   // Ratio is a number of parts of the total cache size to be allocated for this
   // tier. E.g. if X is a total cache size, Yi are ratios specified for memory
   // tiers, and Y is the sum of all Yi, then size of the i-th tier
   // Xi = (X / Y) * Yi. For examle, to configure 2-tier cache where each
   // tier is a half of the total cache size, set both tiers' ratios to 1.
   size_t ratio{1};
 
- private:
+  // Numa node(s) to bind the tier
+  NumaBitMask numaNodes;
+
   // TODO: introduce a container for tier settings when adding support for
   // file-mapped memory
   MemoryTierCacheConfig() = default;

cachelib/cachebench/cache/Cache-inl.h

@@ -80,6 +80,18 @@ Cache<Allocator>::Cache(const CacheConfig& config,
 
   allocatorConfig_.setCacheSize(config_.cacheSizeMB * (MB));
 
+  if (!cacheDir.empty()) {
+    allocatorConfig_.cacheDir = cacheDir;
+  }
+
+  if (config_.usePosixShm) {
+    allocatorConfig_.usePosixForShm();
+  }
+
+  if (!config_.memoryTierConfigs.empty()) {
+    allocatorConfig_.configureMemoryTiers(config_.memoryTierConfigs);
+  }
+
   auto cleanupGuard = folly::makeGuard([&] {
     if (!nvmCacheFilePath_.empty()) {
       util::removePath(nvmCacheFilePath_);
@@ -222,8 +234,7 @@ Cache<Allocator>::Cache(const CacheConfig& config,
   allocatorConfig_.cacheName = "cachebench";
 
   bool isRecovered = false;
-  if (!cacheDir.empty()) {
-    allocatorConfig_.cacheDir = cacheDir;
+  if (!allocatorConfig_.cacheDir.empty()) {
     try {
       cache_ = std::make_unique<Allocator>(Allocator::SharedMemAttach,
                                            allocatorConfig_);

cachelib/cachebench/test_configs/simple_tiers_test.json

@@ -0,0 +1,36 @@
+// @nolint instantiates a small cache and runs a quick run of basic operations.
+{
+    "cache_config" : {
+      "cacheSizeMB" : 512,
+      "usePosixShm" : false,
+      "cacheDir" : "/tmp/mem-tiers",
+      "memoryTiers" : [
+        {
+          "ratio": 1,
+          "memBindNodes": "0"
+        }
+      ],
+      "poolRebalanceIntervalSec" : 1,
+      "moveOnSlabRelease" : false,
+
+      "numPools" : 2,
+      "poolSizes" : [0.3, 0.7]
+    },
+    "test_config" : {
+        "numOps" : 100000,
+        "numThreads" : 32,
+        "numKeys" : 1000000,
+
+        "keySizeRange" : [1, 8, 64],
+        "keySizeRangeProbability" : [0.3, 0.7],
+
+        "valSizeRange" : [1, 32, 10240, 409200],
+        "valSizeRangeProbability" : [0.1, 0.2, 0.7],
+
+        "getRatio" : 0.15,
+        "setRatio" : 0.8,
+        "delRatio" : 0.05,
+        "keyPoolDistribution": [0.4, 0.6],
+        "opPoolDistribution" : [0.5, 0.5]
+    }
+  }

cachelib/cachebench/util/CacheConfig.cpp

@@ -84,6 +84,13 @@ CacheConfig::CacheConfig(const folly::dynamic& configJson) {
 
   JSONSetVal(configJson, memoryOnlyTTL);
 
+  JSONSetVal(configJson, usePosixShm);
+  if (configJson.count("memoryTiers")) {
+    for (auto& it : configJson["memoryTiers"]) {
+      memoryTierConfigs.push_back(MemoryTierConfig(it).getMemoryTierCacheConfig());
+    }
+  }
+
   JSONSetVal(configJson, useTraceTimeStamp);
   JSONSetVal(configJson, printNvmCounters);
   JSONSetVal(configJson, tickerSynchingSeconds);
@@ -95,7 +102,7 @@ CacheConfig::CacheConfig(const folly::dynamic& configJson) {
   // if you added new fields to the configuration, update the JSONSetVal
   // to make them available for the json configs and increment the size
   // below
-  checkCorrectSize<CacheConfig, 696>();
+  checkCorrectSize<CacheConfig, 728>();
 
   if (numPools != poolSizes.size()) {
     throw std::invalid_argument(folly::sformat(
@@ -124,6 +131,13 @@ std::shared_ptr<RebalanceStrategy> CacheConfig::getRebalanceStrategy() const {
         RandomStrategy::Config{static_cast<unsigned int>(rebalanceMinSlabs)});
   }
 }
+
+MemoryTierConfig::MemoryTierConfig(const folly::dynamic& configJson) {
+  JSONSetVal(configJson, ratio);
+  JSONSetVal(configJson, memBindNodes);
+
+  checkCorrectSize<MemoryTierConfig, 40>();
+}
 } // namespace cachebench
 } // namespace cachelib
 } // namespace facebook

cachelib/cachebench/util/CacheConfig.h

@@ -41,6 +41,26 @@ class CacheMonitorFactory {
   virtual std::unique_ptr<CacheMonitor> create(Lru2QAllocator& cache) = 0;
 };
 
+// Parse memory tiers configuration from JSON config
+struct MemoryTierConfig : public JSONConfig {
+  MemoryTierConfig() {}
+
+  explicit MemoryTierConfig(const folly::dynamic& configJson);
+
+  // Returns MemoryTierCacheConfig parsed from JSON config
+  MemoryTierCacheConfig getMemoryTierCacheConfig() {
+    MemoryTierCacheConfig config = MemoryTierCacheConfig::fromShm();
+    config.setRatio(ratio);
+    config.setMemBind(NumaBitMask(memBindNodes));
+    return config;
+  }
+
+  // Specifies ratio of this memory tier to other tiers
+  size_t ratio{0};
+  // Allocate memory only from specified NUMA nodes
+  std::string memBindNodes{""};
+};
+
 struct CacheConfig : public JSONConfig {
   // by defaullt, lru allocator. can be set to LRU-2Q.
   std::string allocator{"LRU"};
@@ -194,6 +214,12 @@ struct CacheConfig : public JSONConfig {
   // Not used when its value is 0.  In seconds.
   uint32_t memoryOnlyTTL{0};
 
+  // Use Posix Shm instead of SysVShm
+  bool usePosixShm{false};
+
+  // Memory tiers configs
+  std::vector<MemoryTierCacheConfig> memoryTierConfigs{};
+
   // If enabled, we will use the timestamps from the trace file in the ticker
   // so that the cachebench will observe time based on timestamps from the trace
   // instead of the system time.

cachelib/shm/CMakeLists.txt

@@ -24,6 +24,7 @@ add_library (cachelib_shm
 add_dependencies(cachelib_shm thrift_generated_files)
 target_link_libraries(cachelib_shm PUBLIC
   cachelib_common
+  numa
 )
 
 install(TARGETS cachelib_shm

cachelib/shm/PosixShmSegment.cpp

@@ -16,11 +16,15 @@
 
 #include "cachelib/shm/PosixShmSegment.h"
 
+#include <cstring>
+
 #include <fcntl.h>
 #include <folly/logging/xlog.h>
 #include <sys/mman.h>
 #include <sys/stat.h>
 #include <sys/types.h>
+#include <numa.h>
+#include <numaif.h>
 
 #include "cachelib/common/Utils.h"
 
@@ -166,6 +170,29 @@ void munmapImpl(void* addr, size_t length) {
   }
 }
 
+void getMempolicyImpl(int &oldMode, NumaBitMask &memBindNumaNodes) {
+  auto nodeMask = memBindNumaNodes.getNativeBitmask();
+
+  long ret = get_mempolicy(&oldMode, nodeMask->maskp, nodeMask->size,
+                           nullptr, 0);
+
+  if (ret != 0) {
+    util::throwSystemError(errno, folly::sformat("get_mempolicy() failed: {}",
+                                                 std::strerror(errno)));
+  }
+}
+
+void setMempolicyImpl(int oldMode, const NumaBitMask &memBindNumaNodes) {
+  auto nodeMask = memBindNumaNodes.getNativeBitmask();
+
+  long ret = set_mempolicy(oldMode, nodeMask->maskp, nodeMask->size);
+
+  if (ret != 0) {
+    util::throwSystemError(errno, folly::sformat("set_mempolicy() failed: {}",
+                                                 std::strerror(errno)));
+  }
+}
+
 } // namespace detail
 
 PosixShmSegment::PosixShmSegment(ShmAttachT,
@@ -312,13 +339,50 @@ void* PosixShmSegment::mapAddress(void* addr) const {
     util::throwSystemError(EINVAL, "Address already mapped");
   }
   XDCHECK(retAddr == addr || addr == nullptr);
+  memBind(addr);
   return retAddr;
 }
 
 void PosixShmSegment::unMap(void* addr) const {
   detail::munmapImpl(addr, getSize());
 }
 
+static void forcePageAllocation(void* addr, size_t size, size_t pageSize) {
+  char* startAddr = reinterpret_cast<char*>(addr);
+  char* endAddr = startAddr + size;
+  for (volatile char* curAddr = startAddr; curAddr < endAddr; curAddr += pageSize) {
+    *curAddr = *curAddr;
+  }
+}
+
+void PosixShmSegment::memBind(void* addr) const {
+  if (opts_.memBindNumaNodes.empty()) {
+    return;
+  }
+
+  NumaBitMask oldMemBindNumaNodes;
+  int oldMode = 0;
+
+  // mbind() cannot be used because mmap was called with MAP_SHARED flag
+  // But we can set memory policy for current thread and force page allocation.
+  // The following logic is used:
+  // 1. Remember current memory policy for the current thread
+  // 2. Set new memory policy as specified by config
+  // 3. Force page allocation by touching every page in the segment
+  // 4. Restore memory policy
+
+  // Remember current memory policy
+  detail::getMempolicyImpl(oldMode, oldMemBindNumaNodes);
+
+  // Set memory bindings
+  detail::setMempolicyImpl(MPOL_BIND, opts_.memBindNumaNodes);
+
+  forcePageAllocation(addr, getSize(), detail::getPageSize(opts_.pageSize));
+
+  // Restore memory policy for the thread
+  detail::setMempolicyImpl(oldMode, oldMemBindNumaNodes);
+}
+
 std::string PosixShmSegment::createKeyForName(
     const std::string& name) noexcept {
   // ensure that the slash is always there in the head. repetitive

cachelib/shm/PosixShmSegment.h

@@ -110,6 +110,8 @@ class PosixShmSegment : public ShmBase {
   void createReferenceMapping();
   void deleteReferenceMapping() const;
 
+  void memBind(void* addr) const;
+
   // file descriptor associated with the shm. This has FD_CLOEXEC set
   // and once opened, we close this only on destruction of this object
   int fd_{kInvalidFD};