[tracing] Implement composable memory usage estimators.

base/test/scoped_memory_usage.cc

+// Copyright 2016 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "base/test/scoped_memory_usage.h"
+
+#include <memory>
+#include <unordered_map>
+
+#include "base/allocator/allocator_shim.h"
+#include "base/logging.h"
+#include "base/memory/ptr_util.h"
+#include "base/threading/thread_local_storage.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace base {
+namespace test {
+
+namespace {
+
+// Unshimmed API
+
+void* UnshimmedAlloc(size_t size) {
+  auto& dispatch = allocator::AllocatorDispatch::default_dispatch;
+  return dispatch.alloc_function(&dispatch, size);
+}
+
+void UnshimmedFree(void* ptr) {
+  auto& dispatch = allocator::AllocatorDispatch::default_dispatch;
+  dispatch.free_function(&dispatch, ptr);
+}
+
+#define USING_UNSHIMMED_NEW()                                        \
+  void* operator new(size_t size) { return UnshimmedAlloc(size); }   \
+  void* operator new[](size_t size) { return UnshimmedAlloc(size); } \
+  void operator delete(void* ptr) { UnshimmedFree(ptr); }            \
+  void operator delete[](void* ptr) { UnshimmedFree(ptr); }
+
+template <class T>
+struct unshimmed_allocator : std::allocator<T> {
+  typedef std::allocator<T> base;
+
+  using typename base::size_type;
+  using typename base::difference_type;
+  using typename base::pointer;
+  using typename base::const_pointer;
+  using typename base::reference;
+  using typename base::const_reference;
+  using typename base::value_type;
+
+  using base::construct;
+  using base::destroy;
+
+  template <class U>
+  struct rebind {
+    typedef unshimmed_allocator<U> other;
+  };
+
+  unshimmed_allocator() {}
+
+  template <class U>
+  unshimmed_allocator(const unshimmed_allocator<U>&) {}
+
+  pointer allocate(size_type count, const void* hint = 0) {
+    return static_cast<pointer>(UnshimmedAlloc(count * sizeof(T)));
+  }
+
+  void deallocate(pointer ptr, size_type) { UnshimmedFree(ptr); }
+};
+
+template <class Key,
+          class T,
+          class Hash = std::hash<Key>,
+          class KeyEqual = std::equal_to<Key>>
+using unshimmed_unordered_map =
+    std::unordered_map<Key,
+                       T,
+                       Hash,
+                       KeyEqual,
+                       unshimmed_allocator<std::pair<const Key, T>>>;
+
+// ThreadAllocationRegister
+
+class ThreadAllocationRegister {
+ public:
+  USING_UNSHIMMED_NEW();
+
+  ThreadAllocationRegister() : total_size_(0) {}
+
+  size_t total_size() const { return total_size_; }
+
+  void Register(void* ptr, size_t size) {
+    EXPECT_TRUE(allocations_.insert({ptr, size}).second);
+    total_size_ += size;
+  }
+
+  void Remove(void* ptr) {
+    auto allocation = allocations_.find(ptr);
+    if (allocation != allocations_.end()) {
+      total_size_ -= allocation->second;
+      allocations_.erase(allocation);
+    }
+  }
+
+ private:
+  unshimmed_unordered_map<void*, size_t> allocations_;
+  size_t total_size_;
+};
+
+ThreadLocalStorage::StaticSlot g_thread_allocation_register = TLS_INITIALIZER;
+
+void InstallThreadAllocationRegister() {
+  CHECK(g_thread_allocation_register.initialized());
+  EXPECT_EQ(nullptr, g_thread_allocation_register.Get());
+  g_thread_allocation_register.Set(new ThreadAllocationRegister());
+}
+
+void UninstallThreadAllocationRegister() {
+  CHECK(g_thread_allocation_register.initialized());
+  EXPECT_NE(nullptr, g_thread_allocation_register.Get());
+  auto* allocation_register = static_cast<ThreadAllocationRegister*>(
+      g_thread_allocation_register.Get());
+  g_thread_allocation_register.Set(nullptr);
+  delete allocation_register;
+}
+
+ThreadAllocationRegister* GetThreadAllocationRegister() {
+  CHECK(g_thread_allocation_register.initialized());
+  return static_cast<ThreadAllocationRegister*>(
+      g_thread_allocation_register.Get());
+}
+
+// AllocatorDispatch
+
+void RegisterAllocation(void* ptr, size_t size) {
+  if (!ptr) {
+    return;
+  }
+  if (auto* allocation_register = GetThreadAllocationRegister()) {
+    allocation_register->Register(ptr, size);
+  }
+}
+
+void RemoveAllocation(void* ptr) {
+  if (!ptr) {
+    return;
+  }
+  if (auto* allocation_register = GetThreadAllocationRegister()) {
+    allocation_register->Remove(ptr);
+  }
+}
+
+void* HookAlloc(const allocator::AllocatorDispatch* self, size_t size) {
+  auto* next = self->next;
+  void* ptr = next->alloc_function(next, size);
+  RegisterAllocation(ptr, size);
+  return ptr;
+}
+
+void* HookAllocZeroInitialized(const allocator::AllocatorDispatch* self,
+                               size_t n,
+                               size_t size) {
+  const allocator::AllocatorDispatch* const next = self->next;
+  void* ptr = next->alloc_zero_initialized_function(next, n, size);
+  RegisterAllocation(ptr, n * size);
+  return ptr;
+}
+
+void* HookAllocAligned(const allocator::AllocatorDispatch* self,
+                       size_t alignment,
+                       size_t size) {
+  const allocator::AllocatorDispatch* const next = self->next;
+  void* ptr = next->alloc_aligned_function(next, alignment, size);
+  RegisterAllocation(ptr, size);
+  return ptr;
+}
+
+void* HookRealloc(const allocator::AllocatorDispatch* self,
+                  void* ptr,
+                  size_t new_size) {
+  const allocator::AllocatorDispatch* const next = self->next;
+  void* new_ptr = next->realloc_function(next, ptr, new_size);
+  RemoveAllocation(ptr);
+  if (new_size != 0) {
+    RegisterAllocation(new_ptr, new_size);
+  }
+  return new_ptr;
+}
+
+void HookFree(const allocator::AllocatorDispatch* self, void* ptr) {
+  RemoveAllocation(ptr);
+  const allocator::AllocatorDispatch* const next = self->next;
+  next->free_function(next, ptr);
+}
+
+size_t HookGetSizeEstimate(const allocator::AllocatorDispatch* self,
+                           void* ptr) {
+  const allocator::AllocatorDispatch* const next = self->next;
+  return next->get_size_estimate_function(next, ptr);
+}
+
+allocator::AllocatorDispatch g_allocator_hooks = {
+    &HookAlloc,                /* alloc_function */
+    &HookAllocZeroInitialized, /* alloc_zero_initialized_function */
+    &HookAllocAligned,         /* alloc_aligned_function */
+    &HookRealloc,              /* realloc_function */
+    &HookFree,                 /* free_function */
+    &HookGetSizeEstimate,      /* get_size_estimate_function */
+    nullptr,                   /* next */
+};
+
+}  // namespace
+
+ScopedMemoryUsage::ScopedMemoryUsage() {
+  InstallThreadAllocationRegister();
+}
+
+ScopedMemoryUsage::~ScopedMemoryUsage() {
+  UninstallThreadAllocationRegister();
+}
+
+size_t ScopedMemoryUsage::Usage() const {
+  return GetThreadAllocationRegister()->total_size();
+}
+
+void ScopedMemoryUsage::Initialize() {
+  if (g_thread_allocation_register.initialized()) {
+    return;
+  }
+
+  g_thread_allocation_register.Initialize([](void* allocation_register) {
+    delete static_cast<ThreadAllocationRegister*>(allocation_register);
+  });
+
+  allocator::InsertAllocatorDispatch(&g_allocator_hooks);
+}
+
+}  // namespace test
+}  // namespace base