From 3535b89149f98a387ab04b95ed81862a5b336c5c Mon Sep 17 00:00:00 2001
From: David Mason <davmason@microsoft.com>
Date: Mon, 7 Feb 2022 11:19:31 -0800
Subject: [PATCH] Fix profiling objectsallocatedbyclass (#63814) (#64284)

* Add test for ObjectsAllocatedByClass profiler callbacks

* Fix ObjectsAllocatedByClass profiler callbacks based on wrong flag

* Fix wrong args in test log

Co-authored-by: Olivier Giniaux <oginiaux@smartadserver.com>

Co-authored-by: Olivier Giniaux <oginiaux@gmail.com>
Co-authored-by: Olivier Giniaux <oginiaux@smartadserver.com>
---
 src/coreclr/vm/proftoeeinterfaceimpl.cpp       |  2 +-
 .../profiler/native/gcprofiler/gcprofiler.cpp  | 18 ++++++++++++++++++
 .../profiler/native/gcprofiler/gcprofiler.h    |  3 +++
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/src/coreclr/vm/proftoeeinterfaceimpl.cpp b/src/coreclr/vm/proftoeeinterfaceimpl.cpp
index 232aa7fefa663..35bc318f1b20e 100644
--- a/src/coreclr/vm/proftoeeinterfaceimpl.cpp
+++ b/src/coreclr/vm/proftoeeinterfaceimpl.cpp
@@ -1274,7 +1274,7 @@ bool AllocByClassHelper(Object * pBO, void * pv)
     _ASSERTE(pv != NULL);
 
     {
-        BEGIN_PROFILER_CALLBACK(CORProfilerTrackAllocations());
+        BEGIN_PROFILER_CALLBACK(CORProfilerTrackGC());
         // Pass along the call
         g_profControlBlock.AllocByClass(
             (ObjectID) pBO,
diff --git a/src/tests/profiler/native/gcprofiler/gcprofiler.cpp b/src/tests/profiler/native/gcprofiler/gcprofiler.cpp
index 42666a2d8eac8..ed6312cd89c46 100644
--- a/src/tests/profiler/native/gcprofiler/gcprofiler.cpp
+++ b/src/tests/profiler/native/gcprofiler/gcprofiler.cpp
@@ -37,6 +37,10 @@ HRESULT GCProfiler::Shutdown()
     {
         printf("GCProfiler::Shutdown: FAIL: Expected GarbageCollectionFinished to be called\n");
     }
+    else if (_allocatedByClassCalls == 0)
+    {
+        printf("GCProfiler::Shutdown: FAIL: Expected ObjectsAllocatedByClass to be called\n");
+    }
     else if (_pohObjectsSeenRootReferences == 0 || _pohObjectsSeenObjectReferences == 0)
     {
         printf("GCProfiler::Shutdown: FAIL: no POH objects seen. root references=%d object references=%d\n",
@@ -86,6 +90,20 @@ HRESULT GCProfiler::GarbageCollectionFinished()
     return S_OK;
 }
 
+HRESULT GCProfiler::ObjectsAllocatedByClass(ULONG cClassCount, ClassID classIds[], ULONG cObjects[])
+{
+    SHUTDOWNGUARD();
+
+    _allocatedByClassCalls++;
+    if (_gcStarts != _allocatedByClassCalls)
+    {
+        _failures++;
+        printf("GCProfiler::ObjectsAllocatedByClass: FAIL: Expected ObjectsAllocatedByClass Calls == GCStart. AllocatedByClassCalls=%d, GCStart=%d\n", (int)_allocatedByClassCalls, (int)_gcStarts);
+    }
+
+    return S_OK;
+}
+
 HRESULT GCProfiler::ObjectReferences(ObjectID objectId, ClassID classId, ULONG cObjectRefs, ObjectID objectRefIds[])
 {
     SHUTDOWNGUARD();
diff --git a/src/tests/profiler/native/gcprofiler/gcprofiler.h b/src/tests/profiler/native/gcprofiler/gcprofiler.h
index ab519f6f53399..eebdb06853f7a 100644
--- a/src/tests/profiler/native/gcprofiler/gcprofiler.h
+++ b/src/tests/profiler/native/gcprofiler/gcprofiler.h
@@ -16,6 +16,7 @@ class GCProfiler : public Profiler
     GCProfiler() : Profiler(),
         _gcStarts(0),
         _gcFinishes(0),
+        _allocatedByClassCalls(0),
         _failures(0),
         _pohObjectsSeenRootReferences(0),
         _pohObjectsSeenObjectReferences(0),
@@ -28,12 +29,14 @@ class GCProfiler : public Profiler
     virtual HRESULT STDMETHODCALLTYPE Shutdown();
     virtual HRESULT STDMETHODCALLTYPE GarbageCollectionStarted(int cGenerations, BOOL generationCollected[], COR_PRF_GC_REASON reason);
     virtual HRESULT STDMETHODCALLTYPE GarbageCollectionFinished();
+    virtual HRESULT STDMETHODCALLTYPE ObjectsAllocatedByClass(ULONG cClassCount, ClassID classIds[], ULONG cObjects[]);
     virtual HRESULT STDMETHODCALLTYPE ObjectReferences(ObjectID objectId, ClassID classId, ULONG cObjectRefs, ObjectID objectRefIds[]);
     virtual HRESULT STDMETHODCALLTYPE RootReferences(ULONG cRootRefs, ObjectID rootRefIds[]);
 
 private:
     std::atomic<int> _gcStarts;
     std::atomic<int> _gcFinishes;
+    std::atomic<int> _allocatedByClassCalls;
     std::atomic<int> _failures;
     std::atomic<int> _pohObjectsSeenRootReferences;
     std::atomic<int> _pohObjectsSeenObjectReferences;