Provide an allocator for new memory type to be used with RocksDB block cache

CMakeLists.txt

@@ -595,6 +595,7 @@ set(SOURCES
         memory/arena.cc
         memory/concurrent_arena.cc
         memory/jemalloc_nodump_allocator.cc
+        memory/memkind_kmem_allocator.cc
         memtable/alloc_tracker.cc
         memtable/hash_linklist_rep.cc
         memtable/hash_skiplist_rep.cc
@@ -1029,6 +1030,7 @@ if(WITH_TESTS)
         logging/env_logger_test.cc
         logging/event_logger_test.cc
         memory/arena_test.cc
+        memory/memkind_kmem_allocator_test.cc
         memtable/inlineskiplist_test.cc
         memtable/skiplist_test.cc
         memtable/write_buffer_manager_test.cc

Makefile

@@ -505,6 +505,7 @@ TESTS = \
 	column_family_test \
 	table_properties_collector_test \
 	arena_test \
+	memkind_kmem_allocator_test \
 	block_test \
 	data_block_hash_index_test \
 	cache_test \
@@ -1239,6 +1240,9 @@ db_repl_stress: tools/db_repl_stress.o $(LIBOBJECTS) $(TESTUTIL)
 
 arena_test: memory/arena_test.o $(LIBOBJECTS) $(TESTHARNESS)
 	$(AM_LINK)
+
+memkind_kmem_allocator_test: memory/memkind_kmem_allocator_test.o $(LIBOBJECTS) $(TESTHARNESS)
+	$(AM_LINK)	
 
 autovector_test: util/autovector_test.o $(LIBOBJECTS) $(TESTHARNESS)
 	$(AM_LINK)

build_tools/build_detect_platform

@@ -28,6 +28,7 @@
 #       -DZSTD                      if the ZSTD library is present
 #       -DNUMA                      if the NUMA library is present
 #       -DTBB                       if the TBB library is present
+#       -DMEMKIND                   if the memkind library is present
 #
 # Using gflags in rocksdb:
 # Our project depends on gflags, which requires users to take some extra steps
@@ -425,7 +426,23 @@ EOF
             COMMON_FLAGS="$COMMON_FLAGS -DROCKSDB_MALLOC_USABLE_SIZE"
         fi
     fi
-
+
+    if ! test $ROCKSDB_DISABLE_MEMKIND; then
+        # Test whether memkind library is installed
+        $CXX $CFLAGS $COMMON_FLAGS -lmemkind -x c++ - -o /dev/null 2>/dev/null  <<EOF
+          #include <memkind.h>
+          int main() {
+            memkind_malloc(MEMKIND_DAX_KMEM, 1024);
+            return 0;
+          }
+EOF
+        if [ "$?" = 0 ]; then
+            COMMON_FLAGS="$COMMON_FLAGS -DMEMKIND"
+            PLATFORM_LDFLAGS="$PLATFORM_LDFLAGS -lmemkind"
+            JAVA_LDFLAGS="$JAVA_LDFLAGS -lmemkind"
+        fi
+    fi
+
     if ! test $ROCKSDB_DISABLE_PTHREAD_MUTEX_ADAPTIVE_NP; then
         # Test whether PTHREAD_MUTEX_ADAPTIVE_NP mutex type is available
         $CXX $CFLAGS -x c++ - -o /dev/null 2>/dev/null  <<EOF

memory/memkind_kmem_allocator.cc

@@ -0,0 +1,32 @@
+//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+//  Copyright (c) 2019 Intel Corporation
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#ifdef MEMKIND
+
+#include "memkind_kmem_allocator.h"
+
+namespace rocksdb {
+
+void* MemkindKmemAllocator::Allocate(size_t size) {
+  void* p = memkind_malloc(MEMKIND_DAX_KMEM, size);
+  if (p == NULL) {
+	throw std::bad_alloc();
+  }
+  return p;
+}
+
+void MemkindKmemAllocator::Deallocate(void* p) {
+  memkind_free(MEMKIND_DAX_KMEM, p);
+}
+
+#ifdef ROCKSDB_MALLOC_USABLE_SIZE
+size_t MemkindKmemAllocator::UsableSize(void* p, size_t /*allocation_size*/) const {
+  return memkind_malloc_usable_size(MEMKIND_DAX_KMEM, p);
+}
+#endif  // ROCKSDB_MALLOC_USABLE_SIZE
+
+}  // namespace rocksdb
+#endif  // MEMKIND

memory/memkind_kmem_allocator.h

@@ -0,0 +1,28 @@
+//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+//  Copyright (c) 2019 Intel Corporation
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#pragma once
+
+#ifdef MEMKIND
+
+#include <memkind.h>
+#include "rocksdb/memory_allocator.h"
+
+namespace rocksdb {
+
+class MemkindKmemAllocator : public MemoryAllocator {
+ public:
+  const char* Name() const override { return "MemkindKmemAllocator"; };
+  void* Allocate(size_t size) override;
+  void Deallocate(void* p) override;
+#ifdef ROCKSDB_MALLOC_USABLE_SIZE
+  size_t UsableSize(void* p, size_t /*allocation_size*/) const override;
+#endif
+};
+
+}  // namespace rocksdb
+#endif  // MEMKIND
+

memory/memkind_kmem_allocator_test.cc

@@ -0,0 +1,99 @@
+//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+//  Copyright (c) 2019 Intel Corporation
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#include <cstdio>
+
+#ifdef MEMKIND
+#include "memkind_kmem_allocator.h"
+#include "test_util/testharness.h"
+#include "rocksdb/cache.h"
+#include "rocksdb/db.h"
+#include "rocksdb/options.h"
+#include "table/block_based/block_based_table_factory.h"
+
+namespace rocksdb {
+TEST(MemkindKmemAllocatorTest, Allocate) {
+  MemkindKmemAllocator allocator;
+  void* p;
+  try {
+    p = allocator.Allocate(1024);
+  } catch (const std::bad_alloc& e) {
+    return;
+  }
+  ASSERT_NE(p, nullptr);
+  size_t size = allocator.UsableSize(p, 1024);
+  ASSERT_GE(size, 1024);
+  allocator.Deallocate(p);
+}
+
+TEST(MemkindKmemAllocatorTest, DatabaseBlockCache) {
+  // Check if a memory node is available for allocation
+  try {
+    MemkindKmemAllocator allocator;
+    allocator.Allocate(1024);
+  } catch (const std::bad_alloc& e) {
+    return;  // if no node available, skip the test
+  }
+
+  // Create database with block cache using MemkindKmemAllocator
+  Options options;
+  std::string dbname = test::PerThreadDBPath("memkind_kmem_allocator_test");
+  ASSERT_OK(DestroyDB(dbname, options));
+
+  options.create_if_missing = true;
+  std::shared_ptr<Cache> cache = NewLRUCache(1024 * 1024, 6, false, false,
+      std::make_shared<MemkindKmemAllocator>());
+  BlockBasedTableOptions table_options;
+  table_options.block_cache = cache;
+  options.table_factory.reset(NewBlockBasedTableFactory(table_options));
+
+  DB* db = nullptr;
+  Status s = DB::Open(options, dbname, &db);
+  ASSERT_OK(s);
+  ASSERT_NE(db, nullptr);
+  ASSERT_EQ(cache->GetUsage(), 0);
+
+  // Write 2kB (200 values, each 10 bytes)
+  int num_keys = 200;
+  WriteOptions wo;
+  std::string val = "0123456789";
+  for (int i = 0; i < num_keys; i++) {
+    std::string key = std::to_string(i);
+    s = db->Put(wo, Slice(key), Slice(val));
+    ASSERT_OK(s);
+  }
+  ASSERT_OK(db->Flush(FlushOptions()));  // Flush all data from memtable so that reads are from block cache
+
+  // Read and check block cache usage
+  ReadOptions ro;
+  std::string result;
+  for (int i = 0; i < num_keys; i++) {
+    std::string key = std::to_string(i);
+    s = db->Get(ro, key, &result);
+    ASSERT_OK(s);
+    ASSERT_EQ(result, val);
+  }
+  ASSERT_GT(cache->GetUsage(), 2000);
+
+  // Close database
+  s = db->Close();
+  ASSERT_OK(s);
+  ASSERT_OK(DestroyDB(dbname, options));
+}
+}  // namespace rocksdb
+
+int main(int argc, char** argv) {
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}
+
+#else
+
+int main(int /*argc*/, char** /*argv*/) {
+  printf("Skip memkind_kmem_allocator_test as the required library memkind is missing.");
+}
+
+#endif  // MEMKIND

src.mk

@@ -92,6 +92,7 @@ LIB_SOURCES =                                                   \
   memory/arena.cc                                               \
   memory/concurrent_arena.cc                                    \
   memory/jemalloc_nodump_allocator.cc                           \
+  memory/memkind_kmem_allocator.cc                              \
   memtable/alloc_tracker.cc                                     \
   memtable/hash_linklist_rep.cc                                 \
   memtable/hash_skiplist_rep.cc                                 \
@@ -402,6 +403,7 @@ MAIN_SOURCES =                                                          \
   logging/env_logger_test.cc                                            \
   logging/event_logger_test.cc                                          \
   memory/arena_test.cc                                                  \
+  memory/memkind_kmem_allocator_test.cc                                 \
   memtable/inlineskiplist_test.cc                                       \
   memtable/memtablerep_bench.cc                                         \
   memtable/skiplist_test.cc                                             \

tools/db_bench_tool.cc

@@ -73,6 +73,10 @@
 #include "utilities/merge_operators/sortlist.h"
 #include "utilities/persistent_cache/block_cache_tier.h"
 
+#ifdef MEMKIND
+#include "memory/memkind_kmem_allocator.h"
+#endif
+
 #ifdef OS_WIN
 #include <io.h>  // open/close
 #endif
@@ -452,6 +456,9 @@ DEFINE_int64(simcache_size, -1,
 DEFINE_bool(cache_index_and_filter_blocks, false,
             "Cache index/filter blocks in block cache.");
 
+DEFINE_bool(use_cache_memkind_kmem_allocator, false,
+            "Use memkind kmem allocator for block cache.");
+
 DEFINE_bool(partition_index_and_filters, false,
             "Partition index and filter blocks.");
 
@@ -2631,9 +2638,22 @@ class Benchmark {
       }
       return cache;
     } else {
-      return NewLRUCache(
-          static_cast<size_t>(capacity), FLAGS_cache_numshardbits,
-          false /*strict_capacity_limit*/, FLAGS_cache_high_pri_pool_ratio);
+      if(FLAGS_use_cache_memkind_kmem_allocator) {
+#ifdef MEMKIND
+        return NewLRUCache(
+            static_cast<size_t>(capacity), FLAGS_cache_numshardbits,
+            false /*strict_capacity_limit*/, FLAGS_cache_high_pri_pool_ratio,
+            std::make_shared<MemkindKmemAllocator>());
+
+#else
+        fprintf(stderr, "Memkind library is not linked with the binary.");
+        exit(1);
+#endif
+      } else {
+        return NewLRUCache(
+            static_cast<size_t>(capacity), FLAGS_cache_numshardbits,
+            false /*strict_capacity_limit*/, FLAGS_cache_high_pri_pool_ratio);
+      }
     }
   }