[Tracing] Add hook to PartitionAlloc for heap profiling

third_party/WebKit/Source/wtf/PartitionAlloc.h

 /*
  * Copyright (C) 2013 Google Inc. All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are
  * met:
  *
  *     * Redistributions of source code must retain the above copyright
  * notice, this list of conditions and the following disclaimer.
  *     * Redistributions in binary form must reproduce the above
@@ skipping 393 matching lines @@
 WTF_EXPORT void partitionPurgeMemory(PartitionRoot*, int);
 WTF_EXPORT void partitionPurgeMemoryGeneric(PartitionRootGeneric*, int);
 
 WTF_EXPORT NEVER_INLINE void* partitionAllocSlowPath(PartitionRootBase*, int, size_t, PartitionBucket*);
 WTF_EXPORT NEVER_INLINE void partitionFreeSlowPath(PartitionPage*);
 WTF_EXPORT NEVER_INLINE void* partitionReallocGeneric(PartitionRootGeneric*, void*, size_t);
 
 WTF_EXPORT void partitionDumpStats(PartitionRoot*, const char* partitionName, bool isLightDump, PartitionStatsDumper*);
 WTF_EXPORT void partitionDumpStatsGeneric(PartitionRootGeneric*, const char* partitionName, bool isLightDump, PartitionStatsDumper*);
 
+class WTF_EXPORT PartitionAllocHooks {
+public:
+    typedef void AllocationHook(void* address, size_t);
+    typedef void FreeHook(void* address);
+
+    static void setAllocationHook(AllocationHook* hook) { m_allocationHook = hook; }
+    static void setFreeHook(FreeHook* hook) { m_freeHook = hook; }
+
+    static void allocationHookIfEnabled(void* address, size_t size)
+    {
+        AllocationHook* allocationHook = m_allocationHook;
+        if (UNLIKELY(allocationHook != nullptr))
+            allocationHook(address, size);
+    }
+
+    static void freeHookIfEnabled(void* address)
+    {
+        FreeHook* freeHook = m_freeHook;
+        if (UNLIKELY(freeHook != nullptr))
+            freeHook(address);
+    }
+
+    static void reallocHookIfEnabled(void* oldAddress, void* newAddress, size_t size)
+    {
+        // Report a reallocation as a free followed by an allocation.
+        AllocationHook* allocationHook = m_allocationHook;
+        FreeHook* freeHook = m_freeHook;
+        if (UNLIKELY(allocationHook && freeHook)) {
+            freeHook(oldAddress);
+            allocationHook(newAddress, size);
+        }
+    }
+
+private:
+    // Pointers to hook functions that PartitionAlloc will call on allocation and
+    // free if the pointers are non-null.
+    static AllocationHook* m_allocationHook;
+    static FreeHook* m_freeHook;
+};
+
 ALWAYS_INLINE PartitionFreelistEntry* partitionFreelistMask(PartitionFreelistEntry* ptr)
 {
     // We use bswap on little endian as a fast mask for two reasons:
     // 1) If an object is freed and its vtable used where the attacker doesn't
     // get the chance to run allocations between the free and use, the vtable
     // dereference is likely to fault.
     // 2) If the attacker has a linear buffer overflow and elects to try and
     // corrupt a freelist pointer, partial pointer overwrite attacks are
     // thwarted.
     // For big endian, similar guarantees are arrived at with a negation.
@@ skipping 192 matching lines @@
     return ret;
 }
 
 ALWAYS_INLINE void* partitionAlloc(PartitionRoot* root, size_t size)
 {
 #if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
     void* result = malloc(size);
     RELEASE_ASSERT(result);
     return result;
 #else
+    size_t requestedSize = size;
     size = partitionCookieSizeAdjustAdd(size);
     ASSERT(root->initialized);
     size_t index = size >> kBucketShift;
     ASSERT(index < root->numBuckets);
     ASSERT(size == index << kBucketShift);
     PartitionBucket* bucket = &root->buckets()[index];
-    return partitionBucketAlloc(root, 0, size, bucket);
+    void* result = partitionBucketAlloc(root, 0, size, bucket);
+    PartitionAllocHooks::allocationHookIfEnabled(result, requestedSize);
+    return result;
 #endif // defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
 }
 
 ALWAYS_INLINE void partitionFreeWithPage(void* ptr, PartitionPage* page)
 {
     // If these asserts fire, you probably corrupted memory.
 #if ENABLE(ASSERT)
     size_t slotSize = page->bucket->slotSize;
     size_t rawSize = partitionPageGetRawSize(page);
     if (rawSize)
@@ skipping 18 matching lines @@
         // correctly update the size metadata.
         ASSERT(partitionPageGetRawSize(page) == 0);
     }
 }
 
 ALWAYS_INLINE void partitionFree(void* ptr)
 {
 #if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
     free(ptr);
 #else
+    PartitionAllocHooks::freeHookIfEnabled(ptr);
     ptr = partitionCookieFreePointerAdjust(ptr);
     ASSERT(partitionPointerIsValid(ptr));
     PartitionPage* page = partitionPointerToPage(ptr);
     partitionFreeWithPage(ptr, page);
 #endif
 }
 
 ALWAYS_INLINE PartitionBucket* partitionGenericSizeToBucket(PartitionRootGeneric* root, size_t size)
 {
     size_t order = kBitsPerSizet - countLeadingZerosSizet(size);
     // The order index is simply the next few bits after the most significant bit.
     size_t orderIndex = (size >> root->orderIndexShifts[order]) & (kGenericNumBucketsPerOrder - 1);
     // And if the remaining bits are non-zero we must bump the bucket up.
     size_t subOrderIndex = size & root->orderSubIndexMasks[order];
     PartitionBucket* bucket = root->bucketLookups[(order << kGenericNumBucketsPerOrderBits) + orderIndex + !!subOrderIndex];
     ASSERT(!bucket->slotSize || bucket->slotSize >= size);
     ASSERT(!(bucket->slotSize % kGenericSmallestBucket));
     return bucket;
 }
 
 ALWAYS_INLINE void* partitionAllocGenericFlags(PartitionRootGeneric* root, int flags, size_t size)
 {
 #if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
     void* result = malloc(size);
     RELEASE_ASSERT(result);
     return result;
 #else
     ASSERT(root->initialized);
+    size_t requestedSize = size;
     size = partitionCookieSizeAdjustAdd(size);
     PartitionBucket* bucket = partitionGenericSizeToBucket(root, size);
     spinLockLock(&root->lock);
     void* ret = partitionBucketAlloc(root, flags, size, bucket);
     spinLockUnlock(&root->lock);
+    PartitionAllocHooks::allocationHookIfEnabled(ret, requestedSize);
     return ret;
 #endif
 }
 
 ALWAYS_INLINE void* partitionAllocGeneric(PartitionRootGeneric* root, size_t size)
 {
     return partitionAllocGenericFlags(root, 0, size);
 }
 
 ALWAYS_INLINE void partitionFreeGeneric(PartitionRootGeneric* root, void* ptr)
 {
 #if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
     free(ptr);
 #else
     ASSERT(root->initialized);
 
     if (UNLIKELY(!ptr))
         return;
 
+    PartitionAllocHooks::freeHookIfEnabled(ptr);
     ptr = partitionCookieFreePointerAdjust(ptr);
     ASSERT(partitionPointerIsValid(ptr));
     PartitionPage* page = partitionPointerToPage(ptr);
     spinLockLock(&root->lock);
     partitionFreeWithPage(ptr, page);
     spinLockUnlock(&root->lock);
 #endif
 }
 
 ALWAYS_INLINE size_t partitionDirectMapSize(size_t size)
@@ skipping 81 matching lines @@
 using WTF::partitionAlloc;
 using WTF::partitionFree;
 using WTF::partitionAllocGeneric;
 using WTF::partitionFreeGeneric;
 using WTF::partitionReallocGeneric;
 using WTF::partitionAllocActualSize;
 using WTF::partitionAllocSupportsGetSize;
 using WTF::partitionAllocGetSize;
 
 #endif // WTF_PartitionAlloc_h