Reintroducing MallocHooks changes with fix for infinite loop in MallocHooks on Platform::Exit.

runtime/vm/malloc_hooks.cc

Index: runtime/vm/malloc_hooks.cc
diff --git a/runtime/vm/malloc_hooks.cc b/runtime/vm/malloc_hooks.cc
index 65edca7bff776d8d8c571b97c1e3e95c0e199a2c..dd471d977e72612fa686537c2d56a43283c7414d 100644
--- a/runtime/vm/malloc_hooks.cc
+++ b/runtime/vm/malloc_hooks.cc
@@ -2,16 +2,267 @@
 // for details. All rights reserved. Use of this source code is governed by a
 // BSD-style license that can be found in the LICENSE file.
 
-#if defined(DART_USE_TCMALLOC)
+#include "platform/globals.h"
+
+#if defined(DART_USE_TCMALLOC) && !defined(PRODUCT)
 
 #include "vm/malloc_hooks.h"
 
+#include "gperftools/malloc_hook.h"
+
+#include "platform/assert.h"
+#include "vm/hash_map.h"
+#include "vm/lockers.h"
+
 namespace dart {
 
-void MallocHooks::Init() {
-  // TODO(bkonyi): Implement
+// A locker-type class to automatically grab and release the
+// in_malloc_hook_flag_.
+class MallocHookScope {
+ public:
+  static void InitMallocHookFlag() {
+    ASSERT(in_malloc_hook_flag_ == kUnsetThreadLocalKey);
+    in_malloc_hook_flag_ = OSThread::CreateThreadLocal();
+    OSThread::SetThreadLocal(in_malloc_hook_flag_, 0);
+  }
+
+  static void DestroyMallocHookFlag() {
+    ASSERT(in_malloc_hook_flag_ != kUnsetThreadLocalKey);
+    OSThread::DeleteThreadLocal(in_malloc_hook_flag_);
+    in_malloc_hook_flag_ = kUnsetThreadLocalKey;
+  }
+
+  MallocHookScope() {
+    ASSERT(in_malloc_hook_flag_ != kUnsetThreadLocalKey);
+    OSThread::SetThreadLocal(in_malloc_hook_flag_, 1);
+  }
+
+  ~MallocHookScope() {
+    ASSERT(in_malloc_hook_flag_ != kUnsetThreadLocalKey);
+    OSThread::SetThreadLocal(in_malloc_hook_flag_, 0);
+  }
+
+  static bool IsInHook() {
+    ASSERT(in_malloc_hook_flag_ != kUnsetThreadLocalKey);
+    return OSThread::GetThreadLocal(in_malloc_hook_flag_);
+  }
+
+ private:
+  static ThreadLocalKey in_malloc_hook_flag_;
+
+  DISALLOW_ALLOCATION();
+  DISALLOW_COPY_AND_ASSIGN(MallocHookScope);
+};
+
+
+// Custom key/value trait specifically for address/size pairs. Unlike
+// RawPointerKeyValueTrait, the default value is -1 as 0 can be a valid entry.
+class AddressKeyValueTrait {
+ public:
+  typedef const void* Key;
+  typedef intptr_t Value;
+
+  struct Pair {
+    Key key;
+    Value value;
+    Pair() : key(NULL), value(-1) {}
+    Pair(const Key key, const Value& value) : key(key), value(value) {}
+    Pair(const Pair& other) : key(other.key), value(other.value) {}
+  };
+
+  static Key KeyOf(Pair kv) { return kv.key; }
+  static Value ValueOf(Pair kv) { return kv.value; }
+  static intptr_t Hashcode(Key key) { return reinterpret_cast<intptr_t>(key); }
+  static bool IsKeyEqual(Pair kv, Key key) { return kv.key == key; }
+};
+
+
+// Map class that will be used to store mappings between ptr -> allocation size.
+class AddressMap : public MallocDirectChainedHashMap<AddressKeyValueTrait> {
+ public:
+  typedef AddressKeyValueTrait::Key Key;
+  typedef AddressKeyValueTrait::Value Value;
+  typedef AddressKeyValueTrait::Pair Pair;
+
+  inline void Insert(const Key& key, const Value& value) {
+    Pair pair(key, value);
+    MallocDirectChainedHashMap<AddressKeyValueTrait>::Insert(pair);
+  }
+
+  inline bool Lookup(const Key& key, Value* value) {
+    ASSERT(value != NULL);
+    Pair* pair = MallocDirectChainedHashMap<AddressKeyValueTrait>::Lookup(key);
+    if (pair == NULL) {
+      return false;
+    } else {
+      *value = pair->value;
+      return true;
+    }
+  }
+};
+
+
+class MallocHooksState {
+ public:
+  static void RecordAllocHook(const void* ptr, size_t size);
+  static void RecordFreeHook(const void* ptr);
+
+  static bool initialized() { return initialized_; }
+  static void Init() {
+    address_map_ = new AddressMap();
+    initialized_ = true;
+  }
+
+  static Mutex* malloc_hook_mutex() { return malloc_hook_mutex_; }
+
+  static intptr_t allocation_count() { return allocation_count_; }
+
+  static intptr_t heap_allocated_memory_in_bytes() {
+    return heap_allocated_memory_in_bytes_;
+  }
+
+  static void IncrementHeapAllocatedMemoryInBytes(intptr_t size) {
+    ASSERT(size >= 0);
+    heap_allocated_memory_in_bytes_ += size;
+    ++allocation_count_;
+  }
+
+  static void DecrementHeapAllocatedMemoryInBytes(intptr_t size) {
+    ASSERT(size >= 0);
+    ASSERT(heap_allocated_memory_in_bytes_ >= size);
+    heap_allocated_memory_in_bytes_ -= size;
+    --allocation_count_;
+    ASSERT(allocation_count_ >= 0);
+  }
+
+  static AddressMap* address_map() { return address_map_; }
+
+  static void ResetStats() {
+    allocation_count_ = 0;
+    heap_allocated_memory_in_bytes_ = 0;
+    address_map_->Clear();
+  }
+
+  static void TearDown() {
+    initialized_ = false;
+    ResetStats();
+    delete address_map_;
+  }
+
+ private:
+  static bool initialized_;
+  static Mutex* malloc_hook_mutex_;
+  static intptr_t allocation_count_;
+  static intptr_t heap_allocated_memory_in_bytes_;
+  static AddressMap* address_map_;
+
+ private:
+  DISALLOW_ALLOCATION();
+  DISALLOW_COPY_AND_ASSIGN(MallocHooksState);
+};
+
+
+// MallocHooks state / locks.
+ThreadLocalKey MallocHookScope::in_malloc_hook_flag_ = kUnsetThreadLocalKey;
+bool MallocHooksState::initialized_ = false;
+Mutex* MallocHooksState::malloc_hook_mutex_ = new Mutex();
+
+// Memory allocation state information.
+intptr_t MallocHooksState::allocation_count_ = 0;
+intptr_t MallocHooksState::heap_allocated_memory_in_bytes_ = 0;
+AddressMap* MallocHooksState::address_map_ = NULL;
+
+
+void MallocHooks::InitOnce() {
+  MutexLocker ml(MallocHooksState::malloc_hook_mutex());
+  ASSERT(!MallocHooksState::initialized());
+
+  MallocHookScope::InitMallocHookFlag();
+  MallocHooksState::Init();
+
+  // Register malloc hooks.
+  bool success = false;
+  success = MallocHook::AddNewHook(&MallocHooksState::RecordAllocHook);
+  ASSERT(success);
+  success = MallocHook::AddDeleteHook(&MallocHooksState::RecordFreeHook);
+  ASSERT(success);
+}
+
+
+void MallocHooks::TearDown() {
+  MutexLocker ml(MallocHooksState::malloc_hook_mutex());
+  ASSERT(MallocHooksState::initialized());
+
+  // Remove malloc hooks.
+  bool success = false;
+  success = MallocHook::RemoveNewHook(&MallocHooksState::RecordAllocHook);
+  ASSERT(success);
+  success = MallocHook::RemoveDeleteHook(&MallocHooksState::RecordFreeHook);
+  ASSERT(success);
+
+  MallocHooksState::TearDown();
+  MallocHookScope::DestroyMallocHookFlag();
+}
+
+
+void MallocHooks::ResetStats() {
+  MutexLocker ml(MallocHooksState::malloc_hook_mutex());
+  ASSERT(MallocHooksState::initialized());
+
+  MallocHooksState::ResetStats();
+}
+
+
+intptr_t MallocHooks::allocation_count() {
+  MutexLocker ml(MallocHooksState::malloc_hook_mutex());
+  return MallocHooksState::allocation_count();
+}
+
+
+intptr_t MallocHooks::heap_allocated_memory_in_bytes() {
+  MutexLocker ml(MallocHooksState::malloc_hook_mutex());
+  return MallocHooksState::heap_allocated_memory_in_bytes();
+}
+
+
+void MallocHooksState::RecordAllocHook(const void* ptr, size_t size) {
+  if (MallocHookScope::IsInHook()) {
+    return;
+  }
+
+  // Set the malloc hook flag before grabbing the mutex to avoid calling hooks
+  // again.
+  MallocHookScope mhs;
+  MutexLocker ml(MallocHooksState::malloc_hook_mutex());
+  ASSERT(MallocHooksState::initialized());
+
+  if (ptr != NULL) {
+    MallocHooksState::IncrementHeapAllocatedMemoryInBytes(size);
+    MallocHooksState::address_map()->Insert(ptr, size);
+  }
+}
+
+
+void MallocHooksState::RecordFreeHook(const void* ptr) {
+  if (MallocHookScope::IsInHook()) {
+    return;
+  }
+
+  // Set the malloc hook flag before grabbing the mutex to avoid calling hooks
+  // again.
+  MallocHookScope mhs;
+  MutexLocker ml(MallocHooksState::malloc_hook_mutex());
+  ASSERT(MallocHooksState::initialized());
+
+  if (ptr != NULL) {
+    intptr_t size = 0;
+    if (MallocHooksState::address_map()->Lookup(ptr, &size)) {
+      MallocHooksState::DecrementHeapAllocatedMemoryInBytes(size);
+      MallocHooksState::address_map()->Remove(ptr);
+    }
+  }
 }
 
 }  // namespace dart
 
-#endif  // defined(DART_USE_TCMALLOC)
+#endif  // defined(DART_USE_TCMALLOC) && !defined(PRODUCT)