[MERGE #6161 @MikeHolman] add JsEnableOOPJIT API

Merge pull request #6161 from MikeHolman:jsenableoopjit

Allows for chakra to run in the interpreter while the host asynchronously spins up a JIT process

Author: MikeHolman

lib/Jsrt/ChakraCoreWindows.h

@@ -13,10 +13,24 @@
 #include <rpc.h>
 
 /// <summary>
-///     Enables out-of-process JIT by connecting to a Chakra JIT process that was initialized by calling JsInitializeJITServer
+///     Globally enables out-of-process JIT.
 /// </summary>
 /// <remarks>
 ///     Should be called before JS code is executed.
+///     Code in all runtimes will run in interpreter until JsConnectJITProcess is called.
+/// </remarks>
+/// <returns>
+///     The code <c>JsNoError</c> if the operation succeeded, a failure code otherwise.
+/// </returns>
+CHAKRA_API
+JsEnableOOPJIT();
+
+/// <summary>
+///     Globally enables out-of-process JIT and connects to a Chakra JIT process that was initialized
+///     by calling JsInitializeJITServer
+/// </summary>
+/// <remarks>
+///     Out-of-process JIT should be enabled before JS code is executed.
 /// </remarks>
 /// <param name="processHandle">Handle to the JIT process</param>
 /// <param name="serverSecurityDescriptor">Optional pointer to an RPC SECURITY_DESCRIPTOR structure</param>

lib/Jsrt/Core/JsrtCore.cpp

@@ -1509,16 +1509,30 @@ JsDiscardBackgroundParse_Experimental(
 }
 
 #ifdef _WIN32
+CHAKRA_API
+JsEnableOOPJIT()
+{
+#ifdef ENABLE_OOP_NATIVE_CODEGEN
+    JITManager::GetJITManager()->EnableOOPJIT();
+    return JsNoError;
+#else
+    return JsErrorNotImplemented;
+#endif
+}
+
 CHAKRA_API
 JsConnectJITProcess(_In_ HANDLE processHandle, _In_opt_ void* serverSecurityDescriptor, _In_ UUID connectionId)
 {
 #ifdef ENABLE_OOP_NATIVE_CODEGEN
     JITManager::GetJITManager()->EnableOOPJIT();
     ThreadContext::SetJITConnectionInfo(processHandle, serverSecurityDescriptor, connectionId);
-#endif
     return JsNoError;
+#else
+    return JsErrorNotImplemented;
+#endif
 }
 #endif
+
 CHAKRA_API
 JsGetArrayBufferFreeFunction(
     _In_ JsValueRef arrayBuffer,

lib/Jsrt/JsrtCommonExports.inc

@@ -122,6 +122,7 @@
     JsCreateWeakReference
     JsDetachArrayBuffer
     JsGetArrayBufferExtraInfo
+    JsEnableOOPJIT
     JsExternalizeArrayBuffer
     JsGetAndClearExceptionWithMetadata
     JsGetArrayBufferFreeFunction

lib/Runtime/Base/ScriptContext.cpp

@@ -5412,11 +5412,36 @@ ScriptContext::GetJitFuncRangeCache()
         allowPrereserveAlloc = false;
 #endif
         // The EnsureJITThreadContext() call could fail if the JIT Server process has died. In such cases, we should not try to do anything further in the client process.
-        if (this->GetThreadContext()->EnsureJITThreadContext(allowPrereserveAlloc))
+        if (!this->GetThreadContext()->EnsureJITThreadContext(allowPrereserveAlloc))
         {
-            HRESULT hr = JITManager::GetJITManager()->InitializeScriptContext(&contextData, this->GetThreadContext()->GetRemoteThreadContextAddr(), &m_remoteScriptContextAddr);
+            return;
+        }
+
+        HRESULT hr = JITManager::GetJITManager()->InitializeScriptContext(&contextData, this->GetThreadContext()->GetRemoteThreadContextAddr(), &m_remoteScriptContextAddr);
+        JITManager::HandleServerCallResult(hr, RemoteCallType::StateUpdate);
+
+        if (!m_remoteScriptContextAddr)
+        {
+            return;
+        }
+
+        // Initialize mutable ScriptContext state if needed
+        if (this->IsPRNGSeeded())
+        {
+            hr = JITManager::GetJITManager()->SetIsPRNGSeeded(m_remoteScriptContextAddr, TRUE);
             JITManager::HandleServerCallResult(hr, RemoteCallType::StateUpdate);
         }
+
+        if (this->GetLibrary()->GetModuleRecordList())
+        {
+            this->GetLibrary()->GetModuleRecordList()->Map([this](int start, SourceTextModuleRecord* moduleRecord) {
+                HRESULT hr = JITManager::GetJITManager()->AddModuleRecordInfo(
+                    m_remoteScriptContextAddr,
+                    moduleRecord->GetModuleId(),
+                    (intptr_t)moduleRecord->GetLocalExportSlots());
+                JITManager::HandleServerCallResult(hr, RemoteCallType::StateUpdate);
+            });
+        }
     }
 #endif
 

lib/Runtime/Base/ThreadContext.cpp

@@ -2005,6 +2005,13 @@ ThreadContext::EnsureJITThreadContext(bool allowPrereserveAlloc)
         &m_jitThunkStartAddr);
     JITManager::HandleServerCallResult(hr, RemoteCallType::StateUpdate);
 
+    // Initialize mutable ThreadContext state if needed
+    Js::TypeId wellKnownType = this->wellKnownHostTypeIds[WellKnownHostType_HTMLAllCollection];
+    if (m_remoteThreadContextInfo && wellKnownType != Js::TypeIds_Undefined)
+    {
+        hr = JITManager::GetJITManager()->SetWellKnownHostTypeId(m_remoteThreadContextInfo, wellKnownType);
+        JITManager::HandleServerCallResult(hr, RemoteCallType::StateUpdate);
+    }
     return m_remoteThreadContextInfo != nullptr;
 #endif
 }

lib/Runtime/Language/SourceTextModuleRecord.cpp

@@ -1204,7 +1204,7 @@ namespace Js
 #if ENABLE_NATIVE_CODEGEN
             if (JITManager::GetJITManager()->IsOOPJITEnabled() && JITManager::GetJITManager()->IsConnected())
             {
-                PSCRIPTCONTEXT_HANDLE remoteScriptContext = this->scriptContext->GetRemoteScriptAddr();
+                PSCRIPTCONTEXT_HANDLE remoteScriptContext = this->scriptContext->GetRemoteScriptAddr(false);
                 if (remoteScriptContext)
                 {
                     HRESULT hr = JITManager::GetJITManager()->AddModuleRecordInfo(

lib/Runtime/Library/JavascriptLibrary.cpp

@@ -7045,12 +7045,12 @@ namespace Js
 #if ENABLE_NATIVE_CODEGEN
         if (JITManager::GetJITManager()->IsOOPJITEnabled() && JITManager::GetJITManager()->IsConnected())
         {
-            PSCRIPTCONTEXT_HANDLE remoteScriptContext = this->scriptContext->GetRemoteScriptAddr();
+            PSCRIPTCONTEXT_HANDLE remoteScriptContext = this->scriptContext->GetRemoteScriptAddr(false);
             if (remoteScriptContext)
             {
                 HRESULT hr = JITManager::GetJITManager()->SetIsPRNGSeeded(remoteScriptContext, val);
-            JITManager::HandleServerCallResult(hr, RemoteCallType::StateUpdate);
-        }
+                JITManager::HandleServerCallResult(hr, RemoteCallType::StateUpdate);
+            }
         }
 #endif
     }

lib/Runtime/Library/JavascriptLibrary.h

@@ -1198,6 +1198,7 @@ namespace Js
         void SetFakeGlobalFuncForUndefer(FunctionBody* functionBody) { this->fakeGlobalFuncForUndefer = functionBody; }
 
         ModuleRecordList* EnsureModuleRecordList();
+        ModuleRecordList* GetModuleRecordList() const { return this->moduleRecordList; }
         SourceTextModuleRecord* GetModuleRecord(uint moduleId);
 
     private: