[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 198 matching lines @@
     void* result = malloc(size);
     RELEASE_ASSERT(result);
     return result;
 #else
     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, size);
+    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 22 matching lines @@
 
 ALWAYS_INLINE void partitionFree(void* ptr)
 {
 #if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
     free(ptr);
 #else
     ptr = partitionCookieFreePointerAdjust(ptr);
     ASSERT(partitionPointerIsValid(ptr));
     PartitionPage* page = partitionPointerToPage(ptr);
     partitionFreeWithPage(ptr, page);
+    PartitionAllocHooks::freeHookIfEnabled(ptr);

[Jens Widell, 2015/10/21 13:56:51]

In debug, |ptr| here doesn't match what we called the allocation hook with. The
value passed to the allocation hook is what was returned to the caller (which is
what was actually allocated plus kCookieSize) whereas the value we pass to the
free hook here is what was actually allocated (sans that kCookieSize
adjustment.)

I don't know if this is important or not.

[Ruud van Asseldonk, 2015/10/21 15:00:27]

The pointers in the allocation hook and free hook definitely need to match. It
should be fixed now, can you please verify?

 #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 = partitionCookieSizeAdjustAdd(size);
     PartitionBucket* bucket = partitionGenericSizeToBucket(root, size);
     spinLockLock(&root->lock);
     void* ret = partitionBucketAlloc(root, flags, size, bucket);
     spinLockUnlock(&root->lock);
+    PartitionAllocHooks::allocationHookIfEnabled(ret, size);

[Jens Widell, 2015/10/21 13:56:51]

Note that |size| here isn't necessarily the size of the returned memory block,
which may be larger. Don't know if we consider this important to adjust for.

If we intend to record the actual size the caller requested, we might want undo
the adjustment made by partitionCookieSizeAdjustAdd() above. Again, I don't know
if this is something that's important enough to care about.

[Ruud van Asseldonk, 2015/10/21 15:00:27]

Good point. I want to record the requested size. If the bucket size happens to
be larger, the extra memory would be allocator overhead (along with the metadata
pages), and the cookie is overhead as well. Fixed it now to report the requested
size.

On a side note, the way in which |partitionDumpPageStats| treats the
requested/allocated distinction is not ideal if we want to combine that
information with the numbers reported by the heap profiler. See
https://crbug.com/527121.

     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;
 
     ptr = partitionCookieFreePointerAdjust(ptr);
     ASSERT(partitionPointerIsValid(ptr));
     PartitionPage* page = partitionPointerToPage(ptr);
     spinLockLock(&root->lock);
     partitionFreeWithPage(ptr, page);
     spinLockUnlock(&root->lock);
+    PartitionAllocHooks::freeHookIfEnabled(ptr);

[Jens Widell, 2015/10/21 13:56:51]

Same kCookieSize adjustment issue here.

[Ruud van Asseldonk, 2015/10/21 15:00:27]

Fixed.

 #endif
 }
 
 ALWAYS_INLINE size_t partitionDirectMapSize(size_t size)
 {
     // Caller must check that the size is not above the kGenericMaxDirectMapped
     // limit before calling. This also guards against integer overflow in the
     // calculation here.
     ASSERT(size <= kGenericMaxDirectMapped);
     return (size + kSystemPageOffsetMask) & kSystemPageBaseMask;
@@ skipping 75 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