Allow custom profiling of JS and non-JS allocation counts

https://bugs.webkit.org/show_bug.cgi?id=262827
rdar://114841891

Reviewed by Mark Lam.

Adds macros to FastMalloc, IsoHeap, and Gigacage malloc functions to allow
for custom profiling code around different kinds of mallocs. Profiling for
different categories of mallocs can be provided by defining certain macros
(see AllocationCounts.h for more information).

* Source/WTF/wtf/Gigacage.cpp:
(Gigacage::tryAlignedMalloc):
(Gigacage::tryMalloc):
(Gigacage::tryZeroedMalloc):
(Gigacage::tryRealloc):
* Source/bmalloc/CMakeLists.txt:
* Source/bmalloc/bmalloc.xcodeproj/project.pbxproj:
* Source/bmalloc/bmalloc/AllocationCounts.cpp: Added.
* Source/bmalloc/bmalloc/AllocationCounts.h: Added.
* Source/bmalloc/bmalloc/BCompiler.h:
* Source/bmalloc/bmalloc/IsoHeap.cpp:
(bmalloc::api::isoAllocate):
(bmalloc::api::isoTryAllocate):
* Source/bmalloc/bmalloc/bmalloc.h:

Canonical link: https://commits.webkit.org/269033@main

Source/WTF/wtf/FastMalloc.cpp

@@ -535,6 +535,7 @@ void* fastMalloc(size_t size)
     if (!AvoidRecordingScope::avoidRecordingCount())
         MallocCallTracker::singleton().recordMalloc(result, size);
 #endif
+    BPROFILE_ALLOCATION(NON_JS_CELL, result, size);
     return result;
 }
 
@@ -547,14 +548,17 @@ void* fastZeroedMalloc(size_t size)
     if (!AvoidRecordingScope::avoidRecordingCount())
         MallocCallTracker::singleton().recordMalloc(result, size);
 #endif
+    BPROFILE_ALLOCATION(NON_JS_CELL, result, size);
     return result;
 }
 
 TryMallocReturnValue tryFastZeroedMalloc(size_t size)
 {
     FAIL_IF_EXCEEDS_LIMIT(size);
     ASSERT(!forbidMallocUseScopeCount || disableMallocRestrictionScopeCount);
-    return bmalloc::api::tryZeroedMalloc(size);
+    void* result = bmalloc::api::tryZeroedMalloc(size);
+    BPROFILE_TRY_ALLOCATION(NON_JS_CELL, result, size);
+    return result;
 }
 
 void* fastCalloc(size_t numElements, size_t elementSize)
@@ -565,6 +569,7 @@ void* fastCalloc(size_t numElements, size_t elementSize)
     void* result = fastZeroedMalloc(checkedSize);
     if (!result)
         CRASH();
+    BPROFILE_ALLOCATION(NON_JS_CELL, result, size);
     return result;
 }
 
@@ -577,6 +582,7 @@ void* fastRealloc(void* object, size_t size)
     if (!AvoidRecordingScope::avoidRecordingCount())
         MallocCallTracker::singleton().recordRealloc(object, result, size);
 #endif
+    BPROFILE_ALLOCATION(NON_JS_CELL, result, size);
     return result;
 }
 
@@ -620,6 +626,7 @@ void* fastAlignedMalloc(size_t alignment, size_t size)
     if (!AvoidRecordingScope::avoidRecordingCount())
         MallocCallTracker::singleton().recordMalloc(result, size);
 #endif
+    BPROFILE_ALLOCATION(NON_JS_CELL, result, size);
     return result;
 }
 
@@ -632,6 +639,7 @@ void* tryFastAlignedMalloc(size_t alignment, size_t size)
     if (!AvoidRecordingScope::avoidRecordingCount())
         MallocCallTracker::singleton().recordMalloc(result, size);
 #endif
+    BPROFILE_TRY_ALLOCATION(NON_JS_CELL, result, size);
     return result;
 }
 
@@ -644,7 +652,9 @@ TryMallocReturnValue tryFastMalloc(size_t size)
 {
     FAIL_IF_EXCEEDS_LIMIT(size);
     ASSERT(!forbidMallocUseScopeCount || disableMallocRestrictionScopeCount);
-    return bmalloc::api::tryMalloc(size);
+    void* result = bmalloc::api::tryMalloc(size);
+    BPROFILE_TRY_ALLOCATION(NON_JS_CELL, result, size);
+    return result;
 }
 
 TryMallocReturnValue tryFastCalloc(size_t numElements, size_t elementSize)
@@ -661,7 +671,9 @@ TryMallocReturnValue tryFastRealloc(void* object, size_t newSize)
 {
     FAIL_IF_EXCEEDS_LIMIT(newSize);
     ASSERT(!forbidMallocUseScopeCount || disableMallocRestrictionScopeCount);
-    return bmalloc::api::tryRealloc(object, newSize);
+    void* result = bmalloc::api::tryRealloc(object, newSize);
+    BPROFILE_TRY_ALLOCATION(NON_JS_CELL, result, size);
+    return result;
 }
 
 void releaseFastMallocFreeMemoryForThisThread()

Source/WTF/wtf/Gigacage.cpp

@@ -81,6 +81,7 @@ namespace Gigacage {
 void* tryAlignedMalloc(Kind kind, size_t alignment, size_t size)
 {
     void* result = bmalloc::api::tryMemalign(alignment, size, bmalloc::heapKind(kind));
+    BPROFILE_TRY_ALLOCATION(GIGACAGE, kind, result, size);
     WTF::compilerFence();
     return result;
 }
@@ -97,20 +98,23 @@ void alignedFree(Kind kind, void* p)
 void* tryMalloc(Kind kind, size_t size)
 {
     void* result = bmalloc::api::tryMalloc(size, bmalloc::heapKind(kind));
+    BPROFILE_TRY_ALLOCATION(GIGACAGE, kind, result, size);
     WTF::compilerFence();
     return result;
 }
 
 void* tryZeroedMalloc(Kind kind, size_t size)
 {
     void* result = bmalloc::api::tryZeroedMalloc(size, bmalloc::heapKind(kind));
+    BPROFILE_TRY_ALLOCATION(GIGACAGE, kind, result, size);
     WTF::compilerFence();
     return result;
 }
 
 void* tryRealloc(Kind kind, void* pointer, size_t size)
 {
     void* result = bmalloc::api::tryRealloc(pointer, size, bmalloc::heapKind(kind));
+    BPROFILE_TRY_ALLOCATION(GIGACAGE, kind, result, size);
     WTF::compilerFence();
     return result;
 }
@@ -127,6 +131,7 @@ void free(Kind kind, void* p)
 void* tryAllocateZeroedVirtualPages(Kind kind, size_t size)
 {
     void* result = bmalloc::api::tryLargeZeroedMemalignVirtual(WTF::pageSize(), size, bmalloc::heapKind(kind));
+    BPROFILE_TRY_ALLOCATION(GIGACAGE, kind, result, size);
     WTF::compilerFence();
     return result;
 }

Source/bmalloc/CMakeLists.txt

@@ -192,6 +192,7 @@ set(bmalloc_PUBLIC_HEADERS
     bmalloc/Algorithm.h
     bmalloc/AllIsoHeaps.h
     bmalloc/AllIsoHeapsInlines.h
+    bmalloc/AllocationCounts.h
     bmalloc/Allocator.h
     bmalloc/AvailableMemory.h
     bmalloc/BAssert.h

Source/bmalloc/bmalloc.xcodeproj/project.pbxproj

@@ -131,6 +131,8 @@
 		2B6EB7A029EE102E00F10400 /* pas_report_crash_pgm_report.h in Headers */ = {isa = PBXBuildFile; fileRef = 2B6EB79F29EE102E00F10400 /* pas_report_crash_pgm_report.h */; settings = {ATTRIBUTES = (Private, ); }; };
 		2BDF4F4C29E8B8BA0056BF50 /* pas_report_crash.c in Sources */ = {isa = PBXBuildFile; fileRef = 2BDF4F4A29E8B8BA0056BF50 /* pas_report_crash.c */; };
 		2BDF4F4D29E8B8BA0056BF50 /* pas_report_crash.h in Headers */ = {isa = PBXBuildFile; fileRef = 2BDF4F4B29E8B8BA0056BF50 /* pas_report_crash.h */; settings = {ATTRIBUTES = (Private, ); }; };
+		306D490A2ACCA05900CCFEE5 /* AllocationCounts.h in Headers */ = {isa = PBXBuildFile; fileRef = 306D49052ACCA01400CCFEE5 /* AllocationCounts.h */; settings = {ATTRIBUTES = (Private, ); }; };
+		306D490B2ACCA05A00CCFEE5 /* AllocationCounts.h in Headers */ = {isa = PBXBuildFile; fileRef = 306D49052ACCA01400CCFEE5 /* AllocationCounts.h */; settings = {ATTRIBUTES = (Private, ); }; };
 		4426E2801C838EE0008EB042 /* Logging.cpp in Sources */ = {isa = PBXBuildFile; fileRef = 4426E27E1C838EE0008EB042 /* Logging.cpp */; };
 		4426E2811C838EE0008EB042 /* Logging.h in Headers */ = {isa = PBXBuildFile; fileRef = 4426E27F1C838EE0008EB042 /* Logging.h */; settings = {ATTRIBUTES = (Private, ); }; };
 		52F47249210BA30200B730BB /* MemoryStatusSPI.h in Headers */ = {isa = PBXBuildFile; fileRef = 52F47248210BA2F500B730BB /* MemoryStatusSPI.h */; settings = {ATTRIBUTES = (Private, ); }; };
@@ -1329,6 +1331,7 @@
 		2CE2AE512769928200D02BBC /* pas_compact_tagged_void_ptr.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; name = pas_compact_tagged_void_ptr.h; path = libpas/src/libpas/pas_compact_tagged_void_ptr.h; sourceTree = "<group>"; };
 		2CE2AE572769928300D02BBC /* pas_probabilistic_guard_malloc_allocator.c */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.c; name = pas_probabilistic_guard_malloc_allocator.c; path = libpas/src/libpas/pas_probabilistic_guard_malloc_allocator.c; sourceTree = "<group>"; };
 		2CE2AE582769928300D02BBC /* pas_probabilistic_guard_malloc_allocator.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; name = pas_probabilistic_guard_malloc_allocator.h; path = libpas/src/libpas/pas_probabilistic_guard_malloc_allocator.h; sourceTree = "<group>"; };
+		306D49052ACCA01400CCFEE5 /* AllocationCounts.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; name = AllocationCounts.h; path = bmalloc/AllocationCounts.h; sourceTree = "<group>"; };
 		4426E27E1C838EE0008EB042 /* Logging.cpp */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.cpp; name = Logging.cpp; path = bmalloc/Logging.cpp; sourceTree = "<group>"; };
 		4426E27F1C838EE0008EB042 /* Logging.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; name = Logging.h; path = bmalloc/Logging.h; sourceTree = "<group>"; };
 		52F47248210BA2F500B730BB /* MemoryStatusSPI.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; name = MemoryStatusSPI.h; path = bmalloc/darwin/MemoryStatusSPI.h; sourceTree = "<group>"; };
@@ -2071,6 +2074,7 @@
 		14D9DB4E17F2866E00EAAB79 /* heap */ = {
 			isa = PBXGroup;
 			children = (
+				306D49052ACCA01400CCFEE5 /* AllocationCounts.h */,
 				7939885A2076EEB50074A2E7 /* BulkDecommit.h */,
 				140FA00219CE429C00FFD3C8 /* BumpRange.h */,
 				147DC6E21CA5B70B00724E8D /* Chunk.h */,
@@ -2163,6 +2167,7 @@
 			isa = PBXHeadersBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				306D490A2ACCA05900CCFEE5 /* AllocationCounts.h in Headers */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};
@@ -2174,6 +2179,7 @@
 				14DD78C518F48D7500950702 /* Algorithm.h in Headers */,
 				0F7EB8341F9541B000F1ABCB /* AllIsoHeaps.h in Headers */,
 				0F7EB84F1F954D4E00F1ABCB /* AllIsoHeapsInlines.h in Headers */,
+				306D490B2ACCA05A00CCFEE5 /* AllocationCounts.h in Headers */,
 				14DD789818F48D4A00950702 /* Allocator.h in Headers */,
 				6599C5CD1EC3F15900A2F7BB /* AvailableMemory.h in Headers */,
 				14DD78C718F48D7500950702 /* BAssert.h in Headers */,

Source/bmalloc/bmalloc/AllocationCounts.h

@@ -0,0 +1,97 @@
+/*
+ * Copyright (C) 2023 Apple 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:
+ * 1. Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ * 2. Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY APPLE INC. ``AS IS'' AND ANY
+ * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+ * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL APPLE INC. OR
+ * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+ * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+ * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+ * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+#pragma once
+
+#include "BExport.h"
+#include <atomic>
+
+#define BFOR_EACH_ALLOCATION_KIND(macro) \
+    macro(JS_CELL)      /* Allocation of any JSCell */ \
+    macro(NON_JS_CELL)  /* Allocation of any non-JSCell */ \
+    macro(GIGACAGE)     /* Allocation within a gigacage, takes heap kind as first parameter */
+
+#define BPROFILE_ALLOCATION(kind, ...) \
+    BPROFILE_ALLOCATION_ ## kind(__VA_ARGS__)
+
+#define BPROFILE_TRY_ALLOCATION(kind, ...) \
+    BPROFILE_TRY_ALLOCATION_ ## kind(__VA_ARGS__)
+
+// Definitions of specializations of the above macro (i.e. BPROFILE_ALLOCATION_JS_CELL(ptr, size))
+// may be provided in an AllocationCountsAdditions.h header to add custom profiling code
+// at an allocation site, taking both the pointer variable (which may be modified) and the
+// size of the allocation in bytes.
+
+#if __has_include(<WebKitAdditions/AllocationCountsAdditions.h>)
+#include <WebKitAdditions/AllocationCountsAdditions.h>
+#elif __has_include(<AllocationCountsAdditions.h>)
+#include <AllocationCountsAdditions.h>
+#endif
+
+// If allocation profiling macros weren't defined above, we define them below as no-ops.
+// Additionally, BENABLE(PROFILE_<type>_ALLOCATION) can be queried to see if we are profiling
+// allocations of a specific kind.
+
+#ifdef BPROFILE_ALLOCATION_JS_CELL
+#define BENABLE_PROFILE_JS_CELL_ALLOCATION 1
+#else
+#define BENABLE_PROFILE_JS_CELL_ALLOCATION 0
+#define BPROFILE_ALLOCATION_JS_CELL(ptr, size) do { } while (false)
+#endif
+
+#ifdef BPROFILE_ALLOCATION_NON_JS_CELL
+#define BENABLE_PROFILE_NON_JS_CELL_ALLOCATION 1
+#else
+#define BENABLE_PROFILE_NON_JS_CELL_ALLOCATION 0
+#define BPROFILE_ALLOCATION_NON_JS_CELL(ptr, size) do { } while (false)
+#endif
+
+#ifdef BPROFILE_ALLOCATION_GIGACAGE
+#define BENABLE_PROFILE_GIGACAGE_ALLOCATION 1
+#else
+#define BENABLE_PROFILE_GIGACAGE_ALLOCATION 0
+#define BPROFILE_ALLOCATION_GIGACAGE(kind, ptr, size) do { } while (false)
+#endif
+
+#ifdef BPROFILE_TRY_ALLOCATION_JS_CELL
+#define BENABLE_PROFILE_JS_CELL_TRY_ALLOCATION 1
+#else
+#define BENABLE_PROFILE_JS_CELL_TRY_ALLOCATION 0
+#define BPROFILE_TRY_ALLOCATION_JS_CELL(ptr, size) do { } while (false)
+#endif
+
+#ifdef BPROFILE_TRY_ALLOCATION_NON_JS_CELL
+#define BENABLE_PROFILE_NON_JS_CELL_TRY_ALLOCATION 1
+#else
+#define BENABLE_PROFILE_NON_JS_CELL_TRY_ALLOCATION 0
+#define BPROFILE_TRY_ALLOCATION_NON_JS_CELL(ptr, size) do { } while (false)
+#endif
+
+#ifdef BPROFILE_TRY_ALLOCATION_GIGACAGE
+#define BENABLE_PROFILE_GIGACAGE_TRY_ALLOCATION 1
+#else
+#define BENABLE_PROFILE_GIGACAGE_TRY_ALLOCATION 0
+#define BPROFILE_TRY_ALLOCATION_GIGACAGE(kind, ptr, size) do { } while (false)
+#endif

Source/bmalloc/bmalloc/BCompiler.h

@@ -93,6 +93,26 @@
 #define BFALLTHROUGH
 #endif
 
+/* BLIKELY */
+
+#if !defined(BLIKELY) && BCOMPILER(GCC_COMPATIBLE)
+#define BLIKELY(x) __builtin_expect(!!(x), 1)
+#endif
+
+#if !defined(BLIKELY)
+#define BLIKELY(x) (x)
+#endif
+
+/* BUNLIKELY */
+
+#if !defined(BUNLIKELY) && BCOMPILER(GCC_COMPATIBLE)
+#define BUNLIKELY(x) __builtin_expect(!!(x), 0)
+#endif
+
+#if !defined(BUNLIKELY)
+#define BUNLIKELY(x) (x)
+#endif
+
 /* BUNUSED_TYPE_ALIAS */
 
 #if !defined(BUNUSED_TYPE_ALIAS) && BCOMPILER(GCC_COMPATIBLE)

Source/bmalloc/bmalloc/IsoHeap.cpp

@@ -25,6 +25,8 @@
 
 #include "IsoHeap.h"
 
+#include "AllocationCounts.h"
+
 #if BUSE(LIBPAS)
 
 #include "bmalloc_heap_inlines.h"
@@ -38,28 +40,35 @@ void* isoAllocate(pas_heap_ref& heapRef)
     // FIXME: libpas should know how to do the fallback thing.
     // https://bugs.webkit.org/show_bug.cgi?id=227177
 
+    auto typeSize = pas_simple_type_size(reinterpret_cast<pas_simple_type>(heapRef.type));
     if (IsoMallocFallback::shouldTryToFallBack()) {
-        IsoMallocFallback::MallocResult result = IsoMallocFallback::tryMalloc(
-            pas_simple_type_size(reinterpret_cast<pas_simple_type>(heapRef.type)));
+        IsoMallocFallback::MallocResult result = IsoMallocFallback::tryMalloc(typeSize);
         if (result.didFallBack) {
             RELEASE_BASSERT(result.ptr);
+            BPROFILE_ALLOCATION_NON_JS_CELL(result.ptr, typeSize);
             return result.ptr;
         }
     }
 
-    return bmalloc_iso_allocate_inline(&heapRef);
+    void* result = bmalloc_iso_allocate_inline(&heapRef);
+    BPROFILE_ALLOCATION_NON_JS_CELL(result, typeSize);
+    return result;
 }
 
 void* isoTryAllocate(pas_heap_ref& heapRef)
 {
+    auto typeSize = pas_simple_type_size(reinterpret_cast<pas_simple_type>(heapRef.type));
     if (IsoMallocFallback::shouldTryToFallBack()) {
-        IsoMallocFallback::MallocResult result = IsoMallocFallback::tryMalloc(
-            pas_simple_type_size(reinterpret_cast<pas_simple_type>(heapRef.type)));
-        if (result.didFallBack)
+        IsoMallocFallback::MallocResult result = IsoMallocFallback::tryMalloc(typeSize);
+        if (result.didFallBack) {
+            BPROFILE_TRY_ALLOCATION(NON_JS_CELL, result.ptr, typeSize);
             return result.ptr;
+        }
     }
 
-    return bmalloc_try_iso_allocate_inline(&heapRef);
+    void* result = bmalloc_try_iso_allocate_inline(&heapRef);
+    BPROFILE_TRY_ALLOCATION(NON_JS_CELL, result, typeSize);
+    return result;
 }
 
 void isoDeallocate(void* ptr)

Source/bmalloc/bmalloc/bmalloc.h

@@ -25,6 +25,7 @@
 
 #pragma once
 
+#include "AllocationCounts.h"
 #include "AvailableMemory.h"
 #include "Cache.h"
 #include "DebugHeap.h"