| // Licensed to the Apache Software Foundation (ASF) under one |
| // or more contributor license agreements. See the NOTICE file |
| // distributed with this work for additional information |
| // regarding copyright ownership. The ASF licenses this file |
| // to you 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. |
| |
| #include "kudu/util/memory/arena.h" |
| |
| #include <cstdint> |
| #include <cstring> |
| #include <memory> |
| #include <string> |
| #include <thread> |
| #include <type_traits> |
| #include <vector> |
| |
| #include <gflags/gflags.h> |
| #include <glog/logging.h> |
| #include <gtest/gtest.h> |
| |
| #include "kudu/gutil/stringprintf.h" |
| #include "kudu/util/mem_tracker.h" |
| #include "kudu/util/memory/memory.h" |
| #include "kudu/util/random.h" |
| #include "kudu/util/slice.h" |
| #include "kudu/util/test_util.h" |
| |
| DEFINE_int32(num_threads, 16, "Number of threads to test"); |
| DEFINE_int32(allocs_per_thread, 10000, "Number of allocations each thread should do"); |
| DEFINE_int32(alloc_size, 4, "number of bytes in each allocation"); |
| |
| namespace kudu { |
| |
| using std::shared_ptr; |
| using std::string; |
| using std::thread; |
| using std::vector; |
| |
| // From the "arena" allocate number of bytes required to copy the "to_write" buffer |
| // and add the allocated buffer to the output "ptrs" vector. |
| template<class ArenaType> |
| static void AllocateBytesAndWrite(ArenaType* arena, const Slice& to_write, vector<void*>* ptrs) { |
| void *allocated_bytes = arena->AllocateBytes(to_write.size()); |
| ASSERT_NE(nullptr, allocated_bytes); |
| memcpy(allocated_bytes, to_write.data(), to_write.size()); |
| ptrs->push_back(allocated_bytes); |
| } |
| |
| // From the "arena" allocate aligned bytes as specified by "alignment". Number of bytes |
| // must be at least the size required to copy the "to_write" buffer. |
| // Add the allocated aligned buffer to the output "ptrs" vector. |
| template<class ArenaType, class RNG> |
| static void AllocateAlignedBytesAndWrite(ArenaType* arena, const size_t alignment, |
| const Slice& to_write, RNG* r, vector<void*>* ptrs) { |
| // To test alignment we allocate random number of bytes within bounds |
| // but write to fixed number of bytes "to_write". |
| size_t num_bytes = FLAGS_alloc_size + r->Uniform32(128 - FLAGS_alloc_size + 1); |
| ASSERT_LE(to_write.size(), num_bytes); |
| void* allocated_bytes = arena->AllocateBytesAligned(num_bytes, alignment); |
| ASSERT_NE(nullptr, allocated_bytes); |
| ASSERT_EQ(0, reinterpret_cast<uintptr_t>(allocated_bytes) % alignment) << |
| "failed to align on " << alignment << "b boundary: " << allocated_bytes; |
| memcpy(allocated_bytes, to_write.data(), to_write.size()); |
| ptrs->push_back(allocated_bytes); |
| } |
| |
| // Thread callback function used by bunch of test cases below. |
| template<class ArenaType, class RNG, bool InvokeAligned = false> |
| static void AllocateAndTestThreadFunc(ArenaType *arena, uint8_t thread_index, RNG* r) { |
| vector<void *> ptrs; |
| ptrs.reserve(FLAGS_allocs_per_thread); |
| |
| uint8_t buf[FLAGS_alloc_size]; |
| memset(buf, thread_index, FLAGS_alloc_size); |
| Slice data(buf, FLAGS_alloc_size); |
| |
| for (int i = 0; i < FLAGS_allocs_per_thread; i++) { |
| if (InvokeAligned) { |
| // Test alignment up to 64 bytes. |
| const size_t alignment = 1 << (i % 7); |
| AllocateAlignedBytesAndWrite(arena, alignment, data, r, &ptrs); |
| } else { |
| AllocateBytesAndWrite(arena, data, &ptrs); |
| } |
| } |
| |
| for (void *p : ptrs) { |
| if (memcmp(buf, p, FLAGS_alloc_size) != 0) { |
| FAIL() << StringPrintf("overwritten pointer at %p", p); |
| } |
| } |
| } |
| |
| // Non-templated function to forward to above -- simplifies thread creation |
| static void AllocateAndTestTSArenaFunc(ThreadSafeArena *arena, uint8_t thread_index, |
| ThreadSafeRandom* r) { |
| AllocateAndTestThreadFunc(arena, thread_index, r); |
| } |
| |
| template<typename ArenaType, typename RNG> |
| static void TestArenaAlignmentHelper() { |
| RNG r(SeedRandom()); |
| |
| for (size_t initial_size = 16; initial_size <= 2048; initial_size <<= 1) { |
| ArenaType arena(initial_size); |
| static constexpr bool kIsMultiThreaded = std::is_same<ThreadSafeArena, ArenaType>::value; |
| if (kIsMultiThreaded) { |
| vector<thread> threads; |
| threads.reserve(FLAGS_num_threads); |
| for (auto i = 0; i < FLAGS_num_threads; i++) { |
| threads.emplace_back( |
| AllocateAndTestThreadFunc<ArenaType, RNG, true /* InvokeAligned */>, &arena, i, &r); |
| } |
| for (thread& thr : threads) { |
| thr.join(); |
| } |
| } else { |
| // Invoke the helper method on the same thread avoiding separate single |
| // thread creation/join. |
| AllocateAndTestThreadFunc<ArenaType, RNG, true /* InvokedAligned */>(&arena, 0, &r); |
| } |
| } |
| } |
| |
| TEST(TestArena, TestSingleThreaded) { |
| Arena arena(128); |
| Random r(SeedRandom()); |
| AllocateAndTestThreadFunc(&arena, 0, &r); |
| } |
| |
| TEST(TestArena, TestMultiThreaded) { |
| CHECK(FLAGS_num_threads < 256); |
| ThreadSafeRandom r(SeedRandom()); |
| ThreadSafeArena arena(1024); |
| |
| vector<thread> threads; |
| threads.reserve(FLAGS_num_threads); |
| for (auto i = 0; i < FLAGS_num_threads; i++) { |
| threads.emplace_back(AllocateAndTestTSArenaFunc, &arena, i, &r); |
| } |
| |
| for (thread& thr : threads) { |
| thr.join(); |
| } |
| } |
| |
| TEST(TestArena, TestAlignmentThreadSafe) { |
| TestArenaAlignmentHelper<ThreadSafeArena, ThreadSafeRandom>(); |
| } |
| |
| TEST(TestArena, TestAlignmentNotThreadSafe) { |
| TestArenaAlignmentHelper<Arena, Random>(); |
| } |
| |
| TEST(TestArena, TestAlignmentSmallArena) { |
| // Start with small initial size and allocate bytes more than the size of the current |
| // component to trigger fallback code path in Arena. Moreover allocate number of bytes |
| // with alignment such that "aligned_size" exceeds "next_component_size". |
| Arena arena(16); |
| constexpr size_t alignment = 32; |
| void *ret = arena.AllocateBytesAligned(33, alignment); |
| ASSERT_NE(nullptr, ret); |
| ASSERT_EQ(0, reinterpret_cast<uintptr_t>(ret) % alignment) << |
| "failed to align on " << alignment << "b boundary: " << ret; |
| } |
| |
| TEST(TestArena, TestObjectAlignment) { |
| struct MyStruct { |
| int64_t v; |
| }; |
| Arena a(256); |
| // Allocate a junk byte to ensure that the next allocation isn't "accidentally" aligned. |
| a.AllocateBytes(1); |
| void* v = a.NewObject<MyStruct>(); |
| ASSERT_EQ(reinterpret_cast<uintptr_t>(v) % alignof(MyStruct), 0); |
| } |
| |
| |
| // MemTrackers update their ancestors when consuming and releasing memory to compute |
| // usage totals. However, the lifetimes of parent and child trackers can be different. |
| // Validate that child trackers can still correctly update their parent stats even when |
| // the parents go out of scope. |
| TEST(TestArena, TestMemoryTrackerParentReferences) { |
| // Set up a parent and child MemTracker. |
| const string parent_id = "parent-id"; |
| const string child_id = "child-id"; |
| shared_ptr<MemTracker> child_tracker; |
| { |
| shared_ptr<MemTracker> parent_tracker = MemTracker::CreateTracker(1024, parent_id); |
| child_tracker = MemTracker::CreateTracker(-1, child_id, parent_tracker); |
| // Parent falls out of scope here. Should still be owned by the child. |
| } |
| shared_ptr<MemoryTrackingBufferAllocator> allocator( |
| new MemoryTrackingBufferAllocator(HeapBufferAllocator::Get(), child_tracker)); |
| MemoryTrackingArena arena(256, allocator); |
| |
| // Try some child operations. |
| ASSERT_EQ(256, child_tracker->consumption()); |
| void *allocated = arena.AllocateBytes(256); |
| ASSERT_TRUE(allocated); |
| ASSERT_EQ(256, child_tracker->consumption()); |
| allocated = arena.AllocateBytes(256); |
| ASSERT_TRUE(allocated); |
| ASSERT_EQ(768, child_tracker->consumption()); |
| } |
| |
| TEST(TestArena, TestMemoryTrackingDontEnforce) { |
| shared_ptr<MemTracker> mem_tracker = MemTracker::CreateTracker(1024, "arena-test-tracker"); |
| shared_ptr<MemoryTrackingBufferAllocator> allocator( |
| new MemoryTrackingBufferAllocator(HeapBufferAllocator::Get(), mem_tracker)); |
| MemoryTrackingArena arena(256, allocator); |
| ASSERT_EQ(256, mem_tracker->consumption()); |
| void *allocated = arena.AllocateBytes(256); |
| ASSERT_TRUE(allocated); |
| ASSERT_EQ(256, mem_tracker->consumption()); |
| allocated = arena.AllocateBytes(256); |
| ASSERT_TRUE(allocated); |
| ASSERT_EQ(768, mem_tracker->consumption()); |
| |
| // In DEBUG mode after Reset() the last component of an arena is |
| // cleared, but is then created again; in release mode, the last |
| // component is not cleared. In either case, after Reset() |
| // consumption() should equal the size of the last component which |
| // is 512 bytes. |
| arena.Reset(); |
| ASSERT_EQ(512, mem_tracker->consumption()); |
| |
| // Allocate beyond allowed consumption. This should still go |
| // through, since enforce_limit is false. |
| allocated = arena.AllocateBytes(1024); |
| ASSERT_TRUE(allocated); |
| |
| ASSERT_EQ(1536, mem_tracker->consumption()); |
| } |
| |
| TEST(TestArena, TestMemoryTrackingEnforced) { |
| shared_ptr<MemTracker> mem_tracker = MemTracker::CreateTracker(1024, "arena-test-tracker"); |
| shared_ptr<MemoryTrackingBufferAllocator> allocator( |
| new MemoryTrackingBufferAllocator(HeapBufferAllocator::Get(), mem_tracker, |
| // enforce limit |
| true)); |
| MemoryTrackingArena arena(256, allocator); |
| ASSERT_EQ(256, mem_tracker->consumption()); |
| void *allocated = arena.AllocateBytes(256); |
| ASSERT_TRUE(allocated); |
| ASSERT_EQ(256, mem_tracker->consumption()); |
| allocated = arena.AllocateBytes(1024); |
| ASSERT_FALSE(allocated); |
| ASSERT_EQ(256, mem_tracker->consumption()); |
| } |
| |
| TEST(TestArena, TestSTLAllocator) { |
| Arena a(256); |
| typedef vector<int, ArenaAllocator<int, false> > ArenaVector; |
| ArenaAllocator<int, false> alloc(&a); |
| ArenaVector v(alloc); |
| for (int i = 0; i < 10000; i++) { |
| v.push_back(i); |
| } |
| for (int i = 0; i < 10000; i++) { |
| ASSERT_EQ(i, v[i]); |
| } |
| } |
| |
| } // namespace kudu |