fix codegen profiler under OOP JIT

lib/Backend/NativeCodeGenerator.cpp

@@ -3448,41 +3448,7 @@ NativeCodeGenerator::SetProfilerFromNativeCodeGen(NativeCodeGenerator * nativeCo
 void
 NativeCodeGenerator::ProfilePrint()
 {
-    Js::ScriptContextProfiler *codegenProfiler = this->backgroundCodeGenProfiler;
-    if (Js::Configuration::Global.flags.Verbose)
-    {
-        //Print individual CodegenProfiler information in verbose mode
-        while (codegenProfiler)
-        {
-            codegenProfiler->ProfilePrint(Js::Configuration::Global.flags.Profile.GetFirstPhase());
-            codegenProfiler = codegenProfiler->next;
-        }
-    }
-    else
-    {
-        //Merge all the codegenProfiler for single snapshot.
-        Js::ScriptContextProfiler* mergeToProfiler = codegenProfiler;
-
-        // find the first initialized profiler
-        while (mergeToProfiler != nullptr && !mergeToProfiler->IsInitialized())
-        {
-            mergeToProfiler = mergeToProfiler->next;
-        }
-        if (mergeToProfiler != nullptr)
-        {
-            // merge the rest profiler to the above initialized profiler
-            codegenProfiler = mergeToProfiler->next;
-            while (codegenProfiler)
-            {
-                if (codegenProfiler->IsInitialized())
-                {
-                    mergeToProfiler->ProfileMerge(codegenProfiler);
-                }
-                codegenProfiler = codegenProfiler->next;
-            }
-            mergeToProfiler->ProfilePrint(Js::Configuration::Global.flags.Profile.GetFirstPhase());
-        }
-    }
+    this->backgroundCodeGenProfiler->ProfilePrint();
 }
 
 void

lib/Backend/PageAllocatorPool.cpp

@@ -106,7 +106,12 @@ void PageAllocatorPool::ReturnPageAllocator(PageAllocator* pageAllocator)
 void PageAllocatorPool::IdleCleanup()
 {
     AutoCriticalSection autoCS(&cs);
-    if (Instance)
+    if (Instance
+#ifdef PROFILE_EXEC
+        // profiler keeps persistent reference to page allocators, so we can't cleanup on idle even if the allocators aren't currently being used for JIT
+        && !Js::Configuration::Global.flags.IsEnabled(Js::ProfileFlag)
+#endif
+        )
     {
         LARGE_INTEGER liDueTime;
         liDueTime.QuadPart = Js::Configuration::Global.flags.JITServerIdleTimeout * -10000LL; // wait for JITServerIdleTimeout milliseconds to do the cleanup

lib/Backend/ServerScriptContext.cpp

@@ -29,7 +29,7 @@ ServerScriptContext::ServerScriptContext(ScriptContextDataIDL * contextData, Ser
     m_codeGenAlloc(nullptr, nullptr, threadContextInfo, threadContextInfo->GetCodePageAllocators(), threadContextInfo->GetProcessHandle()),
     m_globalThisAddr(0),
 #ifdef PROFILE_EXEC
-    m_codeGenProfiler(nullptr),
+    codeGenProfiler(nullptr),
 #endif
     m_refCount(0),
     m_isClosed(false)
@@ -39,12 +39,6 @@ ServerScriptContext::ServerScriptContext(ScriptContextDataIDL * contextData, Ser
     m_codeGenAlloc.canCreatePreReservedSegment = threadContextInfo->CanCreatePreReservedSegment();
 #endif
 
-#ifdef PROFILE_EXEC
-    if (Js::Configuration::Global.flags.IsEnabled(Js::ProfileFlag))
-    {
-        m_codeGenProfiler = HeapNew(Js::ScriptContextProfiler);
-    }
-#endif
     m_domFastPathHelperMap = HeapNew(JITDOMFastPathHelperMap, &HeapAllocator::Instance, 17);
 }
 
@@ -57,9 +51,11 @@ ServerScriptContext::~ServerScriptContext()
     });
 
 #ifdef PROFILE_EXEC
-    if (m_codeGenProfiler)
+    while (this->codeGenProfiler)
     {
-        HeapDelete(m_codeGenProfiler);
+        Js::ScriptContextProfiler* profiler = this->codeGenProfiler;
+        this->codeGenProfiler = this->codeGenProfiler->next;
+        HeapDelete(profiler);
     }
 #endif
 }
@@ -416,14 +412,46 @@ ServerScriptContext::AddModuleRecordInfo(unsigned int moduleId, __int64 localExp
     m_moduleRecords.Add(moduleId, record);
 }
 
-Js::ScriptContextProfiler *
-ServerScriptContext::GetCodeGenProfiler() const
-{
 #ifdef PROFILE_EXEC
-    return m_codeGenProfiler;
-#else
+Js::ScriptContextProfiler*
+ServerScriptContext::GetCodeGenProfiler(_In_ PageAllocator* pageAllocator)
+{
+    if (Js::Configuration::Global.flags.IsEnabled(Js::ProfileFlag))
+    {
+        AutoCriticalSection cs(&this->profilerCS);
+
+        Js::ScriptContextProfiler* profiler = this->codeGenProfiler;
+        while (profiler)
+        {
+            if (profiler->pageAllocator == pageAllocator)
+            {
+                if (!profiler->IsInitialized())
+                {
+                    profiler->Initialize(pageAllocator, nullptr);
+                }
+                return profiler;
+            }
+            profiler = profiler->next;
+        }
+
+        // If we didn't find a profiler, allocate a new one
+
+        profiler = HeapNew(Js::ScriptContextProfiler);
+        profiler->Initialize(pageAllocator, nullptr);
+        profiler->next = this->codeGenProfiler;
+
+        this->codeGenProfiler = profiler;
+
+        return profiler;
+    }
     return nullptr;
-#endif
 }
 
-#endif
+Js::ScriptContextProfiler*
+ServerScriptContext::GetFirstCodeGenProfiler() const
+{
+    return this->codeGenProfiler;
+}
+#endif // PROFILE_EXEC
+
+#endif // ENABLE_OOP_NATIVE_CODEGEN

lib/Backend/ServerScriptContext.h

@@ -82,7 +82,10 @@ public:
     void UpdateGlobalObjectThisAddr(intptr_t globalThis);
     OOPEmitBufferManager * GetEmitBufferManager(bool asmJsManager);
     void DecommitEmitBufferManager(bool asmJsManager);
-    Js::ScriptContextProfiler *  GetCodeGenProfiler() const;
+#ifdef PROFILE_EXEC
+    Js::ScriptContextProfiler*  GetCodeGenProfiler(_In_ PageAllocator* pageAllocator);
+    Js::ScriptContextProfiler* GetFirstCodeGenProfiler() const;
+#endif
     ServerThreadContext* GetThreadContext() { return threadContextHolder.threadContextInfo; }
 
     OOPCodeGenAllocators * GetCodeGenAllocators();
@@ -94,7 +97,8 @@ public:
 private:
     JITDOMFastPathHelperMap * m_domFastPathHelperMap;
 #ifdef PROFILE_EXEC
-    Js::ScriptContextProfiler * m_codeGenProfiler;
+    Js::ScriptContextProfiler * codeGenProfiler;
+    CriticalSection profilerCS;
 #endif
     ArenaAllocator m_sourceCodeArena;
 

lib/JITServer/JITServer.cpp

@@ -430,11 +430,7 @@ ServerCloseScriptContext(
     return ServerCallWrapper(scriptContextInfo, [&]()->HRESULT
     {
 #ifdef PROFILE_EXEC
-        auto profiler = scriptContextInfo->GetCodeGenProfiler();
-        if (profiler && profiler->IsInitialized())
-        {
-            profiler->ProfilePrint(Js::Configuration::Global.flags.Profile.GetFirstPhase());
-        }
+        scriptContextInfo->GetFirstCodeGenProfiler()->ProfilePrint();
 #endif
         scriptContextInfo->Close();
         ServerContextManager::UnRegisterScriptContext(scriptContextInfo);
@@ -733,12 +729,10 @@ ServerRemoteCodeGen(
             Output::Flush();
         }
 
-        auto profiler = scriptContextInfo->GetCodeGenProfiler();
 #ifdef PROFILE_EXEC
-        if (profiler && !profiler->IsInitialized())
-        {
-            profiler->Initialize(pageAllocator, nullptr);
-        }
+        Js::ScriptContextProfiler* profiler = scriptContextInfo->GetCodeGenProfiler(pageAllocator);
+#else
+        Js::ScriptContextProfiler* profiler = nullptr;
 #endif
 
 #if !FLOATVAR

lib/Runtime/Base/ScriptContextProfiler.cpp

@@ -80,6 +80,46 @@ namespace Js
         profiler->Begin(Js::AllPhase);
     }
 
+    void
+    ScriptContextProfiler::ProfilePrint()
+    {
+        Js::ScriptContextProfiler* profiler = this;
+        if (Js::Configuration::Global.flags.Verbose)
+        {
+            //Print individual profiler information in verbose mode
+            while (profiler)
+            {
+                profiler->ProfilePrint(Js::Configuration::Global.flags.Profile.GetFirstPhase());
+                profiler = profiler->next;
+            }
+        }
+        else
+        {
+            //Merge all the profiler for single snapshot.
+            Js::ScriptContextProfiler* mergeToProfiler = profiler;
+
+            // find the first initialized profiler
+            while (mergeToProfiler != nullptr && !mergeToProfiler->IsInitialized())
+            {
+                mergeToProfiler = mergeToProfiler->next;
+            }
+            if (mergeToProfiler != nullptr)
+            {
+                // merge the rest profiler to the above initialized profiler
+                profiler = mergeToProfiler->next;
+                while (profiler)
+                {
+                    if (profiler->IsInitialized())
+                    {
+                        mergeToProfiler->ProfileMerge(profiler);
+                    }
+                    profiler = profiler->next;
+                }
+                mergeToProfiler->ProfilePrint(Js::Configuration::Global.flags.Profile.GetFirstPhase());
+            }
+        }
+    }
+
     ScriptContextProfiler::~ScriptContextProfiler()
     {
         if (profilerArena)

lib/Runtime/Base/ScriptContextProfiler.h

@@ -4,12 +4,14 @@
 //-------------------------------------------------------------------------------------------------------
 #pragma once
 
+class ServerScriptContext;
 namespace Js
 {
     class ScriptContextProfiler
     {
 #ifdef PROFILE_EXEC
         friend class NativeCodeGenerator;
+        friend class ServerScriptContext;
 
     public:
         ScriptContextProfiler();
@@ -29,7 +31,7 @@ namespace Js
         void ProfileResume(Js::Profiler::SuspendRecord * suspendRecord);
         void ProfilePrint(Js::Phase phase);
         void ProfileMerge(ScriptContextProfiler * profiler);
-
+        void ProfilePrint();
     private:
         ArenaAllocator * profilerArena;
         ArenaAllocator * backgroundRecyclerProfilerArena;