From ba52589f965230eb67b70f375a54d5acbed2d9a2 Mon Sep 17 00:00:00 2001
From: David Mason <davmason@microsoft.com>
Date: Fri, 6 Nov 2020 10:58:00 -0800
Subject: [PATCH] Fix ELT profiler tests (#44285)

---
 .../profiler/common/ProfilerTestRunner.cs     | 26 ++++---
 .../eltprofiler/slowpatheltprofiler.cpp       | 70 +++++++++++++++++++
 2 files changed, 87 insertions(+), 9 deletions(-)

diff --git a/src/tests/profiler/common/ProfilerTestRunner.cs b/src/tests/profiler/common/ProfilerTestRunner.cs
index bb5bae9d90924..b1bd4a92ddd78 100644
--- a/src/tests/profiler/common/ProfilerTestRunner.cs
+++ b/src/tests/profiler/common/ProfilerTestRunner.cs
@@ -53,7 +53,7 @@ public static int Run(string profileePath,
 
             if(!File.Exists(profilerPath))
             {
-                LogTestFailure("Profiler library not found at expected path: " + profilerPath);
+                FailFastWithMessage("Profiler library not found at expected path: " + profilerPath);
             }
 
             ProfileeOutputVerifier verifier = new ProfileeOutputVerifier();
@@ -83,17 +83,25 @@ public static int Run(string profileePath,
             process.BeginOutputReadLine();
 
             process.WaitForExit();
-            if (process.ExitCode == 100 && verifier.HasPassingOutput)
-            {
-                return 100;
-            }
-            else
+
+            // There are two conditions for profiler tests to pass, the output of the profiled program
+            // must contain the phrase "PROFILER TEST PASSES" and the return code must be 100. This is
+            // because lots of verification happen in the profiler code, where it is hard to change the
+            // program return value.
+
+            if (!verifier.HasPassingOutput)
             {
-                LogTestFailure("Profiler tests are expected to contain the text \'" + verifier.SuccessPhrase + "\' in the console output " +
+                FailFastWithMessage("Profiler tests are expected to contain the text \'" + verifier.SuccessPhrase + "\' in the console output " +
                     "of the profilee app to indicate a passing test. Usually it is printed from the Shutdown() method of the profiler implementation. This " +
                     "text was not found in the output above.");
-                return process.ExitCode == 100 ? process.ExitCode : -1;
             }
+
+            if (process.ExitCode != 100)
+            {
+                FailFastWithMessage($"Profilee returned exit code {process.ExitCode} instead of expected exit code 100.");
+            }
+
+            return 100;
         }
 
         private static string GetProfilerPath()
@@ -132,7 +140,7 @@ private static string GetCorerunPath()
             return Path.Combine(Environment.GetEnvironmentVariable("CORE_ROOT"), corerunName);
         }
 
-        private static void LogTestFailure(string error)
+        private static void FailFastWithMessage(string error)
         {
             Console.WriteLine("Test failed: " + error);
             throw new Exception(error);
diff --git a/src/tests/profiler/native/eltprofiler/slowpatheltprofiler.cpp b/src/tests/profiler/native/eltprofiler/slowpatheltprofiler.cpp
index 8f858cdb0899b..4551a1003cc6e 100644
--- a/src/tests/profiler/native/eltprofiler/slowpatheltprofiler.cpp
+++ b/src/tests/profiler/native/eltprofiler/slowpatheltprofiler.cpp
@@ -8,11 +8,14 @@
 #include <cctype>
 #include <iomanip>
 #include <algorithm>
+#include <thread>
+#include <chrono>
 
 using std::shared_ptr;
 using std::vector;
 using std::wcout;
 using std::endl;
+using std::atomic;
 
 shared_ptr<SlowPathELTProfiler> SlowPathELTProfiler::s_profiler;
 
@@ -24,18 +27,81 @@ shared_ptr<SlowPathELTProfiler> SlowPathELTProfiler::s_profiler;
 #define PROFILER_STUB EXTERN_C void STDMETHODCALLTYPE
 #endif // WIN32
 
+class ELTGuard
+{
+private:
+    static atomic<bool> s_preventHooks;
+    static atomic<int> s_hooksInProgress;
+
+public:
+    ELTGuard()
+    {
+        ++s_hooksInProgress;
+    }
+
+    ~ELTGuard()
+    {
+        --s_hooksInProgress;
+    }
+
+    static void Initialize()
+    {
+        s_preventHooks = false;
+        s_hooksInProgress = 0;
+    }
+
+    static bool HasShutdownStarted()
+    {
+        return s_preventHooks.load();
+    }
+
+    static void WaitForInProgressHooks()
+    {
+        s_preventHooks = true;
+
+        while (s_hooksInProgress.load() > 0)
+        {
+            std::this_thread::sleep_for(std::chrono::milliseconds(10));
+        }
+    }
+};
+
+atomic<bool> ELTGuard::s_preventHooks;
+atomic<int> ELTGuard::s_hooksInProgress;
+
 PROFILER_STUB EnterStub(FunctionIDOrClientID functionId, COR_PRF_ELT_INFO eltInfo)
 {
+    ELTGuard();
+
+    if (ELTGuard::HasShutdownStarted())
+    {
+        return;
+    }
+
     SlowPathELTProfiler::s_profiler->EnterCallback(functionId, eltInfo);
 }
 
 PROFILER_STUB LeaveStub(FunctionIDOrClientID functionId, COR_PRF_ELT_INFO eltInfo)
 {
+    ELTGuard();
+
+    if (ELTGuard::HasShutdownStarted())
+    {
+        return;
+    }
+
     SlowPathELTProfiler::s_profiler->LeaveCallback(functionId, eltInfo);
 }
 
 PROFILER_STUB TailcallStub(FunctionIDOrClientID functionId, COR_PRF_ELT_INFO eltInfo)
 {
+    ELTGuard();
+
+    if (ELTGuard::HasShutdownStarted())
+    {
+        return;
+    }
+
     SlowPathELTProfiler::s_profiler->TailcallCallback(functionId, eltInfo);
 }
 
@@ -50,6 +116,8 @@ HRESULT SlowPathELTProfiler::Initialize(IUnknown* pICorProfilerInfoUnk)
 {
     Profiler::Initialize(pICorProfilerInfoUnk);
 
+    ELTGuard::Initialize();
+
     HRESULT hr = S_OK;
     constexpr ULONG bufferSize = 1024;
     ULONG envVarLen = 0;
@@ -109,6 +177,8 @@ HRESULT SlowPathELTProfiler::Initialize(IUnknown* pICorProfilerInfoUnk)
 
 HRESULT SlowPathELTProfiler::Shutdown()
 {
+    ELTGuard::WaitForInProgressHooks();
+
     Profiler::Shutdown();
 
     if (_testType == TestType::EnterHooks)