tools/memory_watcher/memory_hook.h

   1 // Copyright (c) 2006-2008 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 // Static class for hooking Win32 API routines.  For now,
   6 // we only add one watcher at a time.
   7 //
   8 // TODO(mbelshe):  Support multiple watchers.
   9
  10 #ifndef MEMORY_WATCHER_MEMORY_HOOK_
  11 #define MEMORY_WATCHER_MEMORY_HOOK_
  12
  13 #include "base/logging.h"
  14
  15 // When allocating memory for internal use with the MemoryHook,
  16 // we must always use the MemoryHook's heap; otherwise, the memory
  17 // gets tracked, and it becomes an infinite loop (allocation() calls
  18 // MemoryHook() which calls allocation(), etc).
  19 //
  20 // PrivateHookAllocator is an STL-friendly Allocator so that STL lists,
  21 // maps, etc can be used on the global MemoryHook's heap.
  22 template <class T>
  23 class PrivateHookAllocator {
  24  public:
  25   // These type definitions are needed for stl allocators.
  26   typedef size_t    size_type;
  27   typedef ptrdiff_t difference_type;
  28   typedef T*        pointer;
  29   typedef const T*  const_pointer;
  30   typedef T&        reference;
  31   typedef const T&  const_reference;
  32   typedef T         value_type;
  33
  34   PrivateHookAllocator() {}
  35
  36   // Allocate memory for STL.
  37   pointer allocate(size_type n, const void * = 0) {
  38     return reinterpret_cast<T*>(MemoryHook::Alloc(n * sizeof(T)));
  39   }
  40
  41   // Deallocate memory for STL.
  42   void deallocate(void* p, size_type) {
  43     if (p)
  44       MemoryHook::Free(p);
  45   }
  46
  47   // Construct the object
  48   void construct(pointer p, const T& val) {
  49     new (reinterpret_cast<T*>(p))T(val);
  50   }
  51
  52   // Destruct an object
  53   void destroy(pointer p) { p->~T(); }
  54
  55   size_type max_size() const { return size_t(-1); }
  56
  57   template <class U>
  58   struct rebind { typedef PrivateHookAllocator<U> other; };
  59
  60   template <class U>
  61   PrivateHookAllocator(const PrivateHookAllocator<U>&) {}
  62 };
  63
  64 template<class T, class U> inline
  65 bool operator==(const PrivateHookAllocator<T>&,
  66                 const PrivateHookAllocator<U>&) {
  67   return (true);
  68 }
  69
  70 template<class T, class U> inline
  71 bool operator!=(const PrivateHookAllocator<T>& left,
  72                 const PrivateHookAllocator<U>& right) {
  73   return (!(left == right));
  74 }
  75
  76
  77 // Classes which monitor memory from these hooks implement
  78 // the MemoryObserver interface.
  79 class MemoryObserver {
  80  public:
  81   virtual ~MemoryObserver() {}
  82
  83   // Track a pointer.  Will capture the current StackTrace.
  84   virtual void OnTrack(HANDLE heap, int32 id, int32 size) = 0;
  85
  86   // Untrack a pointer, removing it from our list.
  87   virtual void OnUntrack(HANDLE heap, int32 id, int32 size) = 0;
  88 };
  89
  90 class MemoryHook : MemoryObserver {
  91  public:
  92   // Initialize the MemoryHook.  Must be called before
  93   // registering watchers.  This can be called repeatedly,
  94   // but is not thread safe.
  95   static bool Initialize();
  96
  97   // Returns true is memory allocations and deallocations
  98   // are being traced.
  99   static bool hooked() { return hooked_ != NULL; }
 100
 101   // Register a class to receive memory allocation & deallocation
 102   // callbacks.  If we haven't hooked memory yet, this call will
 103   // force memory hooking to start.
 104   static bool RegisterWatcher(MemoryObserver* watcher);
 105
 106   // Register a class to stop receiving callbacks.  If there are
 107   // no more watchers, this call will unhook memory.
 108   static bool UnregisterWatcher(MemoryObserver* watcher);
 109
 110   // MemoryHook provides a private heap for allocating
 111   // unwatched memory.
 112   static void* Alloc(size_t size) {
 113     DCHECK(global_hook_ && global_hook_->heap_);
 114     return HeapAlloc(global_hook_->heap_, 0, size);
 115   }
 116   static void Free(void* ptr) {
 117     DCHECK(global_hook_ && global_hook_->heap_);
 118     HeapFree(global_hook_->heap_, 0, ptr);
 119   }
 120
 121   // Access the global hook.  For internal use only from static "C"
 122   // hooks.
 123   static MemoryHook* hook() { return global_hook_; }
 124
 125   // MemoryObserver interface.
 126   virtual void OnTrack(HANDLE hHeap, int32 id, int32 size);
 127   virtual void OnUntrack(HANDLE hHeap, int32 id, int32 size);
 128
 129  private:
 130   MemoryHook();
 131   ~MemoryHook();
 132
 133   // Enable memory tracing.  When memory is 'hooked',
 134   // MemoryWatchers which have registered will be called
 135   // as memory is allocated and deallocated.
 136   static bool Hook();
 137
 138   // Disables memory tracing.
 139   static bool Unhook();
 140
 141   // Create our private heap
 142   bool CreateHeap();
 143
 144   // Close our private heap.
 145   bool CloseHeap();
 146
 147   MemoryObserver* watcher_;
 148   HANDLE heap_;   // An internal accounting heap.
 149   static bool hooked_;
 150   static MemoryHook* global_hook_;
 151 };
 152
 153 #endif  // MEMORY_WATCHER_MEMORY_HOOK_