OLD | NEW |
(Empty) | |
| 1 // Copyright 2016 The Chromium Authors. All rights reserved. |
| 2 // Use of this source code is governed by a BSD-style license that can be |
| 3 // found in the LICENSE file. |
| 4 |
| 5 #include "base/test/scoped_memory_usage.h" |
| 6 |
| 7 #include <memory> |
| 8 #include <unordered_map> |
| 9 |
| 10 #include "base/allocator/allocator_shim.h" |
| 11 #include "base/logging.h" |
| 12 #include "base/memory/ptr_util.h" |
| 13 #include "base/threading/thread_local_storage.h" |
| 14 #include "testing/gtest/include/gtest/gtest.h" |
| 15 |
| 16 namespace base { |
| 17 namespace test { |
| 18 |
| 19 namespace { |
| 20 |
| 21 // Unshimmed API |
| 22 |
| 23 void* UnshimmedAlloc(size_t size) { |
| 24 auto& dispatch = allocator::AllocatorDispatch::default_dispatch; |
| 25 return dispatch.alloc_function(&dispatch, size); |
| 26 } |
| 27 |
| 28 void UnshimmedFree(void* ptr) { |
| 29 auto& dispatch = allocator::AllocatorDispatch::default_dispatch; |
| 30 dispatch.free_function(&dispatch, ptr); |
| 31 } |
| 32 |
| 33 #define USING_UNSHIMMED_NEW() \ |
| 34 void* operator new(size_t size) { return UnshimmedAlloc(size); } \ |
| 35 void* operator new[](size_t size) { return UnshimmedAlloc(size); } \ |
| 36 void operator delete(void* ptr) { UnshimmedFree(ptr); } \ |
| 37 void operator delete[](void* ptr) { UnshimmedFree(ptr); } |
| 38 |
| 39 template <class T> |
| 40 struct unshimmed_allocator : std::allocator<T> { |
| 41 typedef std::allocator<T> base; |
| 42 |
| 43 using typename base::size_type; |
| 44 using typename base::difference_type; |
| 45 using typename base::pointer; |
| 46 using typename base::const_pointer; |
| 47 using typename base::reference; |
| 48 using typename base::const_reference; |
| 49 using typename base::value_type; |
| 50 |
| 51 using base::construct; |
| 52 using base::destroy; |
| 53 |
| 54 template <class U> |
| 55 struct rebind { |
| 56 typedef unshimmed_allocator<U> other; |
| 57 }; |
| 58 |
| 59 unshimmed_allocator() {} |
| 60 |
| 61 template <class U> |
| 62 unshimmed_allocator(const unshimmed_allocator<U>&) {} |
| 63 |
| 64 pointer allocate(size_type count, const void* hint = 0) { |
| 65 return static_cast<pointer>(UnshimmedAlloc(count * sizeof(T))); |
| 66 } |
| 67 |
| 68 void deallocate(pointer ptr, size_type) { UnshimmedFree(ptr); } |
| 69 }; |
| 70 |
| 71 template <class Key, |
| 72 class T, |
| 73 class Hash = std::hash<Key>, |
| 74 class KeyEqual = std::equal_to<Key>> |
| 75 using unshimmed_unordered_map = |
| 76 std::unordered_map<Key, |
| 77 T, |
| 78 Hash, |
| 79 KeyEqual, |
| 80 unshimmed_allocator<std::pair<const Key, T>>>; |
| 81 |
| 82 // ThreadAllocationRegister |
| 83 |
| 84 class ThreadAllocationRegister { |
| 85 public: |
| 86 USING_UNSHIMMED_NEW(); |
| 87 |
| 88 ThreadAllocationRegister() : total_size_(0) {} |
| 89 |
| 90 size_t total_size() const { return total_size_; } |
| 91 |
| 92 void Register(void* ptr, size_t size) { |
| 93 EXPECT_TRUE(allocations_.insert({ptr, size}).second); |
| 94 total_size_ += size; |
| 95 } |
| 96 |
| 97 void Remove(void* ptr) { |
| 98 auto allocation = allocations_.find(ptr); |
| 99 if (allocation != allocations_.end()) { |
| 100 total_size_ -= allocation->second; |
| 101 allocations_.erase(allocation); |
| 102 } |
| 103 } |
| 104 |
| 105 private: |
| 106 unshimmed_unordered_map<void*, size_t> allocations_; |
| 107 size_t total_size_; |
| 108 }; |
| 109 |
| 110 ThreadLocalStorage::StaticSlot g_thread_allocation_register = TLS_INITIALIZER; |
| 111 |
| 112 void InstallThreadAllocationRegister() { |
| 113 CHECK(g_thread_allocation_register.initialized()); |
| 114 EXPECT_EQ(nullptr, g_thread_allocation_register.Get()); |
| 115 g_thread_allocation_register.Set(new ThreadAllocationRegister()); |
| 116 } |
| 117 |
| 118 void UninstallThreadAllocationRegister() { |
| 119 CHECK(g_thread_allocation_register.initialized()); |
| 120 EXPECT_NE(nullptr, g_thread_allocation_register.Get()); |
| 121 auto* allocation_register = static_cast<ThreadAllocationRegister*>( |
| 122 g_thread_allocation_register.Get()); |
| 123 g_thread_allocation_register.Set(nullptr); |
| 124 delete allocation_register; |
| 125 } |
| 126 |
| 127 ThreadAllocationRegister* GetThreadAllocationRegister() { |
| 128 CHECK(g_thread_allocation_register.initialized()); |
| 129 return static_cast<ThreadAllocationRegister*>( |
| 130 g_thread_allocation_register.Get()); |
| 131 } |
| 132 |
| 133 // AllocatorDispatch |
| 134 |
| 135 void RegisterAllocation(void* ptr, size_t size) { |
| 136 if (!ptr) { |
| 137 return; |
| 138 } |
| 139 if (auto* allocation_register = GetThreadAllocationRegister()) { |
| 140 allocation_register->Register(ptr, size); |
| 141 } |
| 142 } |
| 143 |
| 144 void RemoveAllocation(void* ptr) { |
| 145 if (!ptr) { |
| 146 return; |
| 147 } |
| 148 if (auto* allocation_register = GetThreadAllocationRegister()) { |
| 149 allocation_register->Remove(ptr); |
| 150 } |
| 151 } |
| 152 |
| 153 void* HookAlloc(const allocator::AllocatorDispatch* self, size_t size) { |
| 154 auto* next = self->next; |
| 155 void* ptr = next->alloc_function(next, size); |
| 156 RegisterAllocation(ptr, size); |
| 157 return ptr; |
| 158 } |
| 159 |
| 160 void* HookAllocZeroInitialized(const allocator::AllocatorDispatch* self, |
| 161 size_t n, |
| 162 size_t size) { |
| 163 const allocator::AllocatorDispatch* const next = self->next; |
| 164 void* ptr = next->alloc_zero_initialized_function(next, n, size); |
| 165 RegisterAllocation(ptr, n * size); |
| 166 return ptr; |
| 167 } |
| 168 |
| 169 void* HookAllocAligned(const allocator::AllocatorDispatch* self, |
| 170 size_t alignment, |
| 171 size_t size) { |
| 172 const allocator::AllocatorDispatch* const next = self->next; |
| 173 void* ptr = next->alloc_aligned_function(next, alignment, size); |
| 174 RegisterAllocation(ptr, size); |
| 175 return ptr; |
| 176 } |
| 177 |
| 178 void* HookRealloc(const allocator::AllocatorDispatch* self, |
| 179 void* ptr, |
| 180 size_t new_size) { |
| 181 const allocator::AllocatorDispatch* const next = self->next; |
| 182 void* new_ptr = next->realloc_function(next, ptr, new_size); |
| 183 RemoveAllocation(ptr); |
| 184 if (new_size != 0) { |
| 185 RegisterAllocation(new_ptr, new_size); |
| 186 } |
| 187 return new_ptr; |
| 188 } |
| 189 |
| 190 void HookFree(const allocator::AllocatorDispatch* self, void* ptr) { |
| 191 RemoveAllocation(ptr); |
| 192 const allocator::AllocatorDispatch* const next = self->next; |
| 193 next->free_function(next, ptr); |
| 194 } |
| 195 |
| 196 size_t HookGetSizeEstimate(const allocator::AllocatorDispatch* self, |
| 197 void* ptr) { |
| 198 const allocator::AllocatorDispatch* const next = self->next; |
| 199 return next->get_size_estimate_function(next, ptr); |
| 200 } |
| 201 |
| 202 allocator::AllocatorDispatch g_allocator_hooks = { |
| 203 &HookAlloc, /* alloc_function */ |
| 204 &HookAllocZeroInitialized, /* alloc_zero_initialized_function */ |
| 205 &HookAllocAligned, /* alloc_aligned_function */ |
| 206 &HookRealloc, /* realloc_function */ |
| 207 &HookFree, /* free_function */ |
| 208 &HookGetSizeEstimate, /* get_size_estimate_function */ |
| 209 nullptr, /* next */ |
| 210 }; |
| 211 |
| 212 } // namespace |
| 213 |
| 214 ScopedMemoryUsage::ScopedMemoryUsage() { |
| 215 InstallThreadAllocationRegister(); |
| 216 } |
| 217 |
| 218 ScopedMemoryUsage::~ScopedMemoryUsage() { |
| 219 UninstallThreadAllocationRegister(); |
| 220 } |
| 221 |
| 222 size_t ScopedMemoryUsage::Usage() const { |
| 223 return GetThreadAllocationRegister()->total_size(); |
| 224 } |
| 225 |
| 226 void ScopedMemoryUsage::Initialize() { |
| 227 if (g_thread_allocation_register.initialized()) { |
| 228 return; |
| 229 } |
| 230 |
| 231 g_thread_allocation_register.Initialize([](void* allocation_register) { |
| 232 delete static_cast<ThreadAllocationRegister*>(allocation_register); |
| 233 }); |
| 234 |
| 235 allocator::InsertAllocatorDispatch(&g_allocator_hooks); |
| 236 } |
| 237 |
| 238 } // namespace test |
| 239 } // namespace base |
OLD | NEW |