#include <cstdio>
#ifdef MEMKIND
#include "memkind_kmem_allocator.h"
#include "rocksdb/cache.h"
#include "rocksdb/db.h"
#include "rocksdb/options.h"
#include "table/block_based/block_based_table_factory.h"
#include "test_util/testharness.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) {
try {
MemkindKmemAllocator allocator;
allocator.Allocate(1024);
} catch (const std::bad_alloc& e) {
return; }
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);
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()));
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);
s = db->Close();
ASSERT_OK(s);
ASSERT_OK(DestroyDB(dbname, options));
}
}
int main(int argc, char** argv) {
::testing::InitGoogleTest(&argc, argv);
return RUN_ALL_TESTS();
}
#else
int main(int , char** ) {
printf(
"Skip memkind_kmem_allocator_test as the required library memkind is "
"missing.");
}
#endif