diff --git a/util/arena.cc b/util/arena.cc index 6efe687c6..3f00f0845 100644 --- a/util/arena.cc +++ b/util/arena.cc @@ -32,13 +32,20 @@ size_t OptimizeBlockSize(size_t block_size) { return block_size; } -Arena::Arena(size_t block_size) : kBlockSize(OptimizeBlockSize(block_size)) { +Arena::Arena(size_t block_size, size_t huge_page_size) + : kBlockSize(OptimizeBlockSize(block_size)) { assert(kBlockSize >= kMinBlockSize && kBlockSize <= kMaxBlockSize && kBlockSize % kAlignUnit == 0); alloc_bytes_remaining_ = sizeof(inline_block_); blocks_memory_ += alloc_bytes_remaining_; aligned_alloc_ptr_ = inline_block_; unaligned_alloc_ptr_ = inline_block_ + alloc_bytes_remaining_; +#ifdef MAP_HUGETLB + hugetlb_size_ = huge_page_size; + if (hugetlb_size_ && kBlockSize > hugetlb_size_) { + hugetlb_size_ = ((kBlockSize - 1U) / hugetlb_size_ + 1U) * hugetlb_size_; + } +#endif } Arena::~Arena() { @@ -62,20 +69,49 @@ char* Arena::AllocateFallback(size_t bytes, bool aligned) { } // We waste the remaining space in the current block. - auto block_head = AllocateNewBlock(kBlockSize); - alloc_bytes_remaining_ = kBlockSize - bytes; + size_t size; + char* block_head = nullptr; + if (hugetlb_size_) { + size = hugetlb_size_; + block_head = AllocateFromHugePage(size); + } + if (!block_head) { + size = kBlockSize; + block_head = AllocateNewBlock(size); + } + alloc_bytes_remaining_ = size - bytes; if (aligned) { aligned_alloc_ptr_ = block_head + bytes; - unaligned_alloc_ptr_ = block_head + kBlockSize; + unaligned_alloc_ptr_ = block_head + size; return block_head; } else { aligned_alloc_ptr_ = block_head; - unaligned_alloc_ptr_ = block_head + kBlockSize - bytes; + unaligned_alloc_ptr_ = block_head + size - bytes; return unaligned_alloc_ptr_; } } +char* Arena::AllocateFromHugePage(size_t bytes) { +#ifdef MAP_HUGETLB + if (hugetlb_size_ == 0) { + return nullptr; + } + + void* addr = mmap(nullptr, bytes, (PROT_READ | PROT_WRITE), + (MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB), 0, 0); + + if (addr == MAP_FAILED) { + return nullptr; + } + huge_blocks_.push_back(MmapInfo(addr, bytes)); + blocks_memory_ += bytes; + return reinterpret_cast(addr); +#else + return nullptr; +#endif +} + char* Arena::AllocateAligned(size_t bytes, size_t huge_page_size, Logger* logger) { assert((kAlignUnit & (kAlignUnit - 1)) == @@ -88,17 +124,14 @@ char* Arena::AllocateAligned(size_t bytes, size_t huge_page_size, size_t reserved_size = ((bytes - 1U) / huge_page_size + 1U) * huge_page_size; assert(reserved_size >= bytes); - void* addr = mmap(nullptr, reserved_size, (PROT_READ | PROT_WRITE), - (MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB), 0, 0); - if (addr == MAP_FAILED) { + char* addr = AllocateFromHugePage(reserved_size); + if (addr == nullptr) { Warn(logger, "AllocateAligned fail to allocate huge TLB pages: %s", strerror(errno)); // fail back to malloc } else { - blocks_memory_ += reserved_size; - huge_blocks_.push_back(MmapInfo(addr, reserved_size)); - return reinterpret_cast(addr); + return addr; } } #endif diff --git a/util/arena.h b/util/arena.h index dfd8e2b24..4764c1568 100644 --- a/util/arena.h +++ b/util/arena.h @@ -35,7 +35,10 @@ class Arena { static const size_t kMinBlockSize; static const size_t kMaxBlockSize; - explicit Arena(size_t block_size = kMinBlockSize); + // huge_page_size: if 0, don't use huge page TLB. If > 0 (should set to the + // supported hugepage size of the system), block allocation will try huge + // page TLB first. If allocation fails, will fall back to normal case. + explicit Arena(size_t block_size = kMinBlockSize, size_t huge_page_size = 0); ~Arena(); char* Allocate(size_t bytes); @@ -100,6 +103,8 @@ class Arena { // How many bytes left in currently active block? size_t alloc_bytes_remaining_ = 0; + size_t hugetlb_size_ = 0; + char* AllocateFromHugePage(size_t bytes); char* AllocateFallback(size_t bytes, bool aligned); char* AllocateNewBlock(size_t block_bytes); diff --git a/util/arena_test.cc b/util/arena_test.cc index 7b6cfd0af..7f55a7e53 100644 --- a/util/arena_test.cc +++ b/util/arena_test.cc @@ -13,17 +13,21 @@ namespace rocksdb { +namespace { +const size_t kHugePageSize = 2 * 1024 * 1024; +} // namespace class ArenaTest {}; TEST(ArenaTest, Empty) { Arena arena0; } -TEST(ArenaTest, MemoryAllocatedBytes) { +namespace { +void MemoryAllocatedBytesTest(size_t huge_page_size) { const int N = 17; size_t req_sz; // requested size size_t bsz = 8192; // block size size_t expected_memory_allocated; - Arena arena(bsz); + Arena arena(bsz, huge_page_size); // requested size > quarter of a block: // allocate requested size separately @@ -44,8 +48,15 @@ TEST(ArenaTest, MemoryAllocatedBytes) { for (int i = 0; i < N; i++) { arena.Allocate(req_sz); } - expected_memory_allocated += bsz; - ASSERT_EQ(arena.MemoryAllocatedBytes(), expected_memory_allocated); + if (huge_page_size) { + ASSERT_TRUE(arena.MemoryAllocatedBytes() == + expected_memory_allocated + bsz || + arena.MemoryAllocatedBytes() == + expected_memory_allocated + huge_page_size); + } else { + expected_memory_allocated += bsz; + ASSERT_EQ(arena.MemoryAllocatedBytes(), expected_memory_allocated); + } // requested size > quarter of a block: // allocate requested size separately @@ -54,16 +65,23 @@ TEST(ArenaTest, MemoryAllocatedBytes) { arena.Allocate(req_sz); } expected_memory_allocated += req_sz * N; - ASSERT_EQ(arena.MemoryAllocatedBytes(), expected_memory_allocated); + if (huge_page_size) { + ASSERT_TRUE(arena.MemoryAllocatedBytes() == + expected_memory_allocated + bsz || + arena.MemoryAllocatedBytes() == + expected_memory_allocated + huge_page_size); + } else { + ASSERT_EQ(arena.MemoryAllocatedBytes(), expected_memory_allocated); + } } // Make sure we didn't count the allocate but not used memory space in // Arena::ApproximateMemoryUsage() -TEST(ArenaTest, ApproximateMemoryUsageTest) { +static void ApproximateMemoryUsageTest(size_t huge_page_size) { const size_t kBlockSize = 4096; const size_t kEntrySize = kBlockSize / 8; const size_t kZero = 0; - Arena arena(kBlockSize); + Arena arena(kBlockSize, huge_page_size); ASSERT_EQ(kZero, arena.ApproximateMemoryUsage()); // allocate inline bytes @@ -78,7 +96,12 @@ TEST(ArenaTest, ApproximateMemoryUsageTest) { // first allocation arena.AllocateAligned(kEntrySize); auto mem_usage = arena.MemoryAllocatedBytes(); - ASSERT_EQ(mem_usage, kBlockSize + Arena::kInlineSize); + if (huge_page_size) { + ASSERT_TRUE(mem_usage == kBlockSize + Arena::kInlineSize || + mem_usage == huge_page_size + Arena::kInlineSize); + } else { + ASSERT_EQ(mem_usage, kBlockSize + Arena::kInlineSize); + } auto usage = arena.ApproximateMemoryUsage(); ASSERT_LT(usage, mem_usage); for (size_t i = 1; i < num_blocks; ++i) { @@ -87,12 +110,17 @@ TEST(ArenaTest, ApproximateMemoryUsageTest) { ASSERT_EQ(arena.ApproximateMemoryUsage(), usage + kEntrySize); usage = arena.ApproximateMemoryUsage(); } - ASSERT_GT(usage, mem_usage); + if (huge_page_size) { + ASSERT_TRUE(usage > mem_usage || + usage + huge_page_size - kBlockSize == mem_usage); + } else { + ASSERT_GT(usage, mem_usage); + } } -TEST(ArenaTest, Simple) { +static void SimpleTest(size_t huge_page_size) { std::vector> allocated; - Arena arena; + Arena arena(Arena::kMinBlockSize, huge_page_size); const int N = 100000; size_t bytes = 0; Random rnd(301); @@ -136,7 +164,22 @@ TEST(ArenaTest, Simple) { } } } +} // namespace +TEST(ArenaTest, MemoryAllocatedBytes) { + MemoryAllocatedBytesTest(0); + MemoryAllocatedBytesTest(kHugePageSize); +} + +TEST(ArenaTest, ApproximateMemoryUsage) { + ApproximateMemoryUsageTest(0); + ApproximateMemoryUsageTest(kHugePageSize); +} + +TEST(ArenaTest, Simple) { + SimpleTest(0); + SimpleTest(kHugePageSize); +} } // namespace rocksdb int main(int argc, char** argv) { return rocksdb::test::RunAllTests(); }