Initial support of NUMA bindings

cachelib/allocator/CacheAllocator-inl.h

@@ -125,6 +125,7 @@ ShmSegmentOpts CacheAllocator<CacheTrait>::createShmCacheOpts(TierId tid) {
   ShmSegmentOpts opts;
   opts.alignment = sizeof(Slab);
   opts.typeOpts = memoryTierConfigs[tid].getShmTypeOpts();
+  opts.memBindNumaNodes = memoryTierConfigs[tid].getMemBind();
   if (auto *v = std::get_if<PosixSysVSegmentOpts>(&opts.typeOpts)) {
     v->usePosix = config_.usePosixShm;
   }

cachelib/allocator/MemoryTierCacheConfig.h

@@ -53,6 +53,16 @@ class MemoryTierCacheConfig {
 
   size_t getRatio() const noexcept { return ratio; }
 
+  // Allocate memory only from specified NUMA nodes
+  MemoryTierCacheConfig& setMemBind(const std::vector<size_t>& _numaNodes) {
+    numaNodes = _numaNodes;
+    return *this;
+  }
+
+  std::vector<size_t> getMemBind() const {
+    return numaNodes;
+  }
+
   size_t calculateTierSize(size_t totalCacheSize, size_t partitionNum) const {
     // TODO: Call this method when tiers are enabled in allocator
     // to calculate tier sizes in bytes.
@@ -82,6 +92,9 @@ class MemoryTierCacheConfig {
   // Options specific to shm type
   ShmTypeOpts shmOpts;
 
+  // Numa node(s) to bind the tier
+  std::vector<size_t> numaNodes;
+
   MemoryTierCacheConfig() = default;
 };
 } // namespace cachelib

cachelib/allocator/tests/AllocatorMemoryTiersTest.cpp

@@ -23,9 +23,11 @@ namespace tests {
 using LruAllocatorMemoryTiersTest = AllocatorMemoryTiersTest<LruAllocator>;
 
 // TODO(MEMORY_TIER): add more tests with different eviction policies
-TEST_F(LruAllocatorMemoryTiersTest, MultiTiersInvalid) { this->testMultiTiersInvalid(); }
-TEST_F(LruAllocatorMemoryTiersTest, MultiTiersValid) { this->testMultiTiersValid(); }
+TEST_F(LruAllocatorMemoryTiersTest, MultiTiersFromFileInvalid) { this->testMultiTiersFormFileInvalid(); }
+TEST_F(LruAllocatorMemoryTiersTest, MultiTiersFromFileValid) { this->testMultiTiersFromFileValid(); }
 TEST_F(LruAllocatorMemoryTiersTest, MultiTiersValidMixed) { this->testMultiTiersValidMixed(); }
+TEST_F(LruAllocatorMemoryTiersTest, MultiTiersNumaBindingsSysVValid) { this->testMultiTiersNumaBindingsSysVValid(); }
+TEST_F(LruAllocatorMemoryTiersTest, MultiTiersNumaBindingsPosixValid) { this->testMultiTiersNumaBindingsPosixValid(); }
 
 } // end of namespace tests
 } // end of namespace cachelib

cachelib/allocator/tests/AllocatorMemoryTiersTest.h

@@ -27,7 +27,7 @@ namespace tests {
 template <typename AllocatorT>
 class AllocatorMemoryTiersTest : public AllocatorTest<AllocatorT> {
  public:
-  void testMultiTiersInvalid() {
+  void testMultiTiersFormFileInvalid() {
     typename AllocatorT::Config config;
     config.setCacheSize(100 * Slab::kSize);
     config.configureMemoryTiers({
@@ -42,7 +42,7 @@ class AllocatorMemoryTiersTest : public AllocatorTest<AllocatorT> {
                  std::invalid_argument);
   }
 
-  void testMultiTiersValid() {
+  void testMultiTiersFromFileValid() {
     typename AllocatorT::Config config;
     config.setCacheSize(100 * Slab::kSize);
     config.enableCachePersistence("/tmp");
@@ -83,6 +83,47 @@ class AllocatorMemoryTiersTest : public AllocatorTest<AllocatorT> {
     ASSERT(handle != nullptr);
     ASSERT_NO_THROW(alloc->insertOrReplace(handle));
   }
+
+  void testMultiTiersNumaBindingsSysVValid() {
+    typename AllocatorT::Config config;
+    config.setCacheSize(100 * Slab::kSize);
+    config.enableCachePersistence("/tmp");
+    config.configureMemoryTiers({
+        MemoryTierCacheConfig::fromShm()
+            .setRatio(1).setMemBind({0}),
+        MemoryTierCacheConfig::fromShm()
+            .setRatio(1).setMemBind({0})
+    });
+
+    auto alloc = std::make_unique<AllocatorT>(AllocatorT::SharedMemNew, config);
+    ASSERT(alloc != nullptr);
+
+    auto pool = alloc->addPool("default", alloc->getCacheMemoryStats().cacheSize);
+    auto handle = alloc->allocate(pool, "key", std::string("value").size());
+    ASSERT(handle != nullptr);
+    ASSERT_NO_THROW(alloc->insertOrReplace(handle));
+  }
+
+  void testMultiTiersNumaBindingsPosixValid() {
+    typename AllocatorT::Config config;
+    config.setCacheSize(100 * Slab::kSize);
+    config.enableCachePersistence("/tmp");
+    config.usePosixForShm();
+    config.configureMemoryTiers({
+        MemoryTierCacheConfig::fromShm()
+            .setRatio(1).setMemBind({0}),
+        MemoryTierCacheConfig::fromShm()
+            .setRatio(1).setMemBind({0})
+    });
+
+    auto alloc = std::make_unique<AllocatorT>(AllocatorT::SharedMemNew, config);
+    ASSERT(alloc != nullptr);
+
+    auto pool = alloc->addPool("default", alloc->getCacheMemoryStats().cacheSize);
+    auto handle = alloc->allocate(pool, "key", std::string("value").size());
+    ASSERT(handle != nullptr);
+    ASSERT_NO_THROW(alloc->insertOrReplace(handle));
+  }
 };
 } // namespace tests
 } // namespace cachelib

cachelib/cachebench/CMakeLists.txt

@@ -89,5 +89,6 @@ if (BUILD_TESTS)
   add_test (consistency/tests/ValueHistoryTest.cpp)
   add_test (consistency/tests/ValueTrackerTest.cpp)
   add_test (util/tests/NandWritesTest.cpp)
+  add_test (util/tests/MemoryTierConfigTest.cpp)
   add_test (cache/tests/TimeStampTickerTest.cpp)
 endif()

cachelib/cachebench/util/CacheConfig.cpp

@@ -137,8 +137,53 @@ std::shared_ptr<RebalanceStrategy> CacheConfig::getRebalanceStrategy() const {
 MemoryTierConfig::MemoryTierConfig(const folly::dynamic& configJson) {
   JSONSetVal(configJson, file);
   JSONSetVal(configJson, ratio);
+  JSONSetVal(configJson, memBindNodes);
 
-  checkCorrectSize<MemoryTierConfig, 40>();
+  checkCorrectSize<MemoryTierConfig, 72>();
+}
+
+static bool starts_with() {return true;}
+
+std::vector<size_t> MemoryTierConfig::parseNumaNodes() {
+  std::vector<size_t> numaNodes;
+
+  std::vector<folly::StringPiece> tokens;
+  folly::split(",", memBindNodes, tokens, true /*ignore empty*/);
+  for(const auto &token : tokens) {
+    if(token.startsWith("!")) {
+      throw std::invalid_argument(folly::sformat(
+        "invalid NUMA nodes binding in memory tier config: {} "
+        "inverse !N or !N-N is not supported "
+        "nodes may be specified as N,N,N or N-N or N,N-N or N-N,N-N and so forth.",
+        token));
+    }
+    else if(token.startsWith("+")) {
+      throw std::invalid_argument(folly::sformat(
+        "invalid NUMA nodes binding in memory tier config: {} "
+        "relative nodes are not supported. "
+        "nodes may be specified as N,N,N or N-N or N,N-N or N-N,N-N and so forth.",
+        token));
+    }
+    else if (token.contains("-")) {
+      size_t begin, end;
+      if(folly::split("-", token, begin, end) && begin < end) {
+        while(begin <=end) {
+          numaNodes.push_back(begin++);
+        }
+      } else {
+        throw std::invalid_argument(folly::sformat(
+        "invalid NUMA nodes binding in memory tier config: {} "
+        "Invalid range format. "
+        "nodes may be specified as N,N,N or N-N or N,N-N or N-N,N-N and so forth.",
+        token));
+      }
+    }
+    else {
+      numaNodes.push_back(folly::to<size_t>(token));
+    }
+  }
+
+  return numaNodes;
 }
 
 } // namespace cachebench

cachelib/cachebench/util/CacheConfig.h

@@ -48,11 +48,13 @@ struct MemoryTierConfig : public JSONConfig {
   MemoryTierCacheConfig getMemoryTierCacheConfig() {
     MemoryTierCacheConfig config = memoryTierCacheConfigFromSource();
     config.setRatio(ratio);
+    config.setMemBind(parseNumaNodes());
     return config;
   }
 
   std::string file{""};
   size_t ratio{0};
+  std::string memBindNodes{""};
 
 private:
   MemoryTierCacheConfig memoryTierCacheConfigFromSource() {
@@ -62,6 +64,8 @@ struct MemoryTierConfig : public JSONConfig {
       return MemoryTierCacheConfig::fromFile(file);
     }
   }
+
+  std::vector<size_t> parseNumaNodes();
 };
 
 struct CacheConfig : public JSONConfig {

cachelib/cachebench/util/tests/MemoryTierConfigTest.cpp

@@ -0,0 +1,86 @@
+/*
+ * Copyright (c) Facebook, Inc. and its 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.
+ */
+
+// Copyright 2022-present Facebook. All Rights Reserved.
+
+#include <algorithm>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+#include "cachelib/cachebench/util/CacheConfig.h"
+
+namespace facebook {
+namespace cachelib {
+namespace cachebench {
+
+TEST(MemoryTierConfigTest, MemBind_SingleNumaNode) {
+  const std::string configString =
+    "{"
+    "  \"ratio\": 1,"
+    "  \"memBindNodes\": 1"
+    "}";
+
+  const std::vector<size_t> expectedNumaNodes = {1};
+
+  auto configJson = folly::parseJson(folly::json::stripComments(configString));
+
+  MemoryTierConfig memoryTierConfig(configJson);
+  MemoryTierCacheConfig tierCacheConfig = memoryTierConfig.getMemoryTierCacheConfig();
+
+  auto parsedNumaNodes = tierCacheConfig.getMemBind();
+  ASSERT_TRUE(std::equal(expectedNumaNodes.begin(), expectedNumaNodes.end(), parsedNumaNodes.begin()));
+}
+
+TEST(MemoryTierConfigTest, MemBind_RangeNumaNodes) {
+  const std::string configString =
+    "{"
+    "  \"ratio\": 1,"
+    "  \"memBindNodes\": \"0-2\""
+    "}";
+
+  const std::vector<size_t> expectedNumaNodes = {0, 1, 2};
+
+  auto configJson = folly::parseJson(folly::json::stripComments(configString));
+
+  MemoryTierConfig memoryTierConfig(configJson);
+  MemoryTierCacheConfig tierCacheConfig = memoryTierConfig.getMemoryTierCacheConfig();
+
+  auto parsedNumaNodes = tierCacheConfig.getMemBind();
+  ASSERT_TRUE(std::equal(expectedNumaNodes.begin(), expectedNumaNodes.end(), parsedNumaNodes.begin()));
+}
+
+TEST(MemoryTierConfigTest, MemBind_SingleAndRangeNumaNodes) {
+  const std::string configString =
+    "{"
+    "  \"ratio\": 1,"
+    "  \"memBindNodes\": \"0,2-5\""
+    "}";
+
+  const std::vector<size_t> expectedNumaNodes = {0, 2, 3, 4, 5};
+
+  auto configJson = folly::parseJson(folly::json::stripComments(configString));
+
+  MemoryTierConfig memoryTierConfig(configJson);
+  MemoryTierCacheConfig tierCacheConfig = memoryTierConfig.getMemoryTierCacheConfig();
+
+  auto parsedNumaNodes = tierCacheConfig.getMemBind();
+  ASSERT_TRUE(std::equal(expectedNumaNodes.begin(), expectedNumaNodes.end(), parsedNumaNodes.begin()));
+}
+
+} // namespace facebook
+} // namespace cachelib
+} // namespace cachebench

cachelib/shm/CMakeLists.txt

@@ -25,6 +25,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

@@ -21,6 +21,8 @@
 #include <sys/mman.h>
 #include <sys/stat.h>
 #include <sys/types.h>
+#include <numa.h>
+#include <numaif.h>
 
 #include "cachelib/common/Utils.h"
 
@@ -176,13 +178,52 @@ 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) {
+  for(volatile char* curAddr = (char*)addr; curAddr < (char*)addr+size; curAddr += pageSize) {
+    *curAddr = *curAddr;
+  }
+}
+
+void PosixShmSegment::memBind(void* addr) const {
+  if(opts_.memBindNumaNodes.empty()) return;
+
+  struct bitmask *oldNodeMask = numa_allocate_nodemask();
+  int oldMode = 0;
+  struct bitmask *nodesMask = numa_allocate_nodemask();
+  auto guard = folly::makeGuard([&] { numa_bitmask_free(nodesMask); numa_bitmask_free(oldNodeMask); });
+
+  for(auto node : opts_.memBindNumaNodes) {
+    numa_bitmask_setbit(nodesMask, node);
+  }
+
+  // mbind() cannot be used because mmap was called with MAP_SHARED flag
+  // But we can set memory policy for current thread and force page allcoation.
+  // The following logic is used:
+  // 1. Remember current memory policy for the current thread
+  // 2. Set new memory policy as specifiec by config
+  // 3. Force page allocation by touching every page in the segment
+  // 4. Restore memory policy
+
+  // Remember current memory policy 
+  get_mempolicy(&oldMode, oldNodeMask->maskp, oldNodeMask->size, nullptr, 0);
+
+  // Set memory bindings
+  set_mempolicy(MPOL_BIND, nodesMask->maskp, nodesMask->size);
+
+  forcePageAllocation(addr, getSize(), detail::getPageSize(opts_.pageSize));
+
+  // Restore memory policy for the thread
+  set_mempolicy(oldMode, nodesMask->maskp, nodesMask->size);
+}
+
 std::string PosixShmSegment::createKeyForName(
     const std::string& name) noexcept {
   // ensure that the slash is always there in the head. repetitive

cachelib/shm/PosixShmSegment.h

@@ -108,6 +108,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};

cachelib/shm/ShmCommon.h

@@ -93,6 +93,7 @@ struct ShmSegmentOpts {
   PageSizeT pageSize{PageSizeT::NORMAL};
   bool readOnly{false};
   size_t alignment{1}; // alignment for mapping.
+  std::vector<size_t> memBindNumaNodes;
   // opts specific to segment type
   ShmTypeOpts typeOpts{PosixSysVSegmentOpts(false)};