OOP JIT error handling

adding a new type of exception to indicate system call failed(mostly because content process is terminated/crash and VM is deleted, the JIT server is still trying to do the VM operations, which fails)
wrap all server call in exception handling
make relationship between scriptcontext and threadcontext in server side

lib/Backend/CodeGenNumberAllocator.cpp

@@ -345,9 +345,13 @@ Js::JavascriptNumber* XProcNumberPageSegmentImpl::AllocateNumber(Func* func, dou
         {
             Assert((unsigned int)((char*)tail->GetEndAddress() - (char*)tail->GetCommitEndAddress()) >= BlockSize);
             // TODO: implement guard pages (still necessary for OOP JIT?)
-            auto ret = ::VirtualAllocEx(hProcess, tail->GetCommitEndAddress(), BlockSize, MEM_COMMIT, PAGE_READWRITE);
-            if (!ret)
+            LPVOID addr = ::VirtualAllocEx(hProcess, tail->GetCommitEndAddress(), BlockSize, MEM_COMMIT, PAGE_READWRITE);
+            if (addr == nullptr)
             {
+                if (hProcess != GetCurrentProcess())
+                {
+                    Js::Throw::CheckAndThrowJITOperationFailed();
+                }
                 Js::Throw::OutOfMemory();
             }
             tail->committedEnd += BlockSize;
@@ -359,6 +363,10 @@ Js::JavascriptNumber* XProcNumberPageSegmentImpl::AllocateNumber(Func* func, dou
     void* pages = ::VirtualAllocEx(hProcess, nullptr, PageCount * AutoSystemInfo::PageSize, MEM_RESERVE, PAGE_READWRITE);
     if (pages == nullptr)
     {
+        if (hProcess != GetCurrentProcess())
+        {
+            Js::Throw::CheckAndThrowJITOperationFailed();
+        }
         Js::Throw::OutOfMemory();
     }
 

lib/Backend/EmitBuffer.cpp

@@ -290,6 +290,13 @@ EmitBufferAllocation* EmitBufferManager<SyncObject>::AllocateBuffer(__in size_t
 #if DBG
     MEMORY_BASIC_INFORMATION memBasicInfo;
     size_t resultBytes = VirtualQueryEx(this->processHandle, allocation->allocation->address, &memBasicInfo, sizeof(memBasicInfo));
+    if (resultBytes == 0) 
+    {
+        if (this->processHandle != GetCurrentProcess())
+        {
+            Js::Throw::CheckAndThrowJITOperationFailed();
+        }
+    }
     Assert(resultBytes != 0 && memBasicInfo.Protect == PAGE_EXECUTE);
 #endif
 

lib/Backend/NativeCodeGenerator.cpp

@@ -909,6 +909,9 @@ NativeCodeGenerator::CodeGen(PageAllocator * pageAllocator, CodeGenWorkItem* wor
             Js::Throw::OutOfMemory();
         case VBSERR_OutOfStack:
             throw Js::StackOverflowException();
+        case E_ACCESSDENIED:
+            // OOP JIT TODO: if server side can't handle request any more, use better error code and turn off JIT
+            throw Js::JITOperationFailedException(workItem->codeGenResult);
         default:
             Js::Throw::FatalInternalError();
         }
@@ -2026,18 +2029,6 @@ NativeCodeGenerator::UpdateJITState()
             scriptContext->InitializeRemoteScriptContext();
         }
 
-        bool allowPrereserveAlloc = true;
-#if !_M_X64_OR_ARM64
-        if (this->scriptContext->webWorkerId != Js::Constants::NonWebWorkerContextId)
-        {
-            allowPrereserveAlloc = false;
-        }
-#endif
-#ifndef _CONTROL_FLOW_GUARD
-        allowPrereserveAlloc = false;
-#endif
-        scriptContext->GetThreadContext()->EnsureJITThreadContext(allowPrereserveAlloc);
-
         // update all property records on server that have been changed since last jit
         ThreadContext::PropertyMap * pendingProps = scriptContext->GetThreadContext()->GetPendingJITProperties();
         PropertyRecordIDL ** newPropArray = nullptr;
@@ -2080,10 +2071,7 @@ NativeCodeGenerator::UpdateJITState()
             props.newRecordCount = newCount;
             props.newRecordArray = newPropArray;
             HRESULT hr = JITManager::GetJITManager()->UpdatePropertyRecordMap(scriptContext->GetThreadContext()->GetRemoteThreadContextAddr(), &props);
-            if (hr != S_OK)
-            {
-                Js::Throw::FatalInternalError();
-            }
+
             if (newPropArray)
             {
                 HeapDeleteArray(newCount, newPropArray);
@@ -2092,6 +2080,12 @@ NativeCodeGenerator::UpdateJITState()
             {
                 HeapDeleteArray(reclaimedCount, reclaimedPropArray);
             }
+
+            if (hr != S_OK)
+            {
+                // OOP JIT TODO: use better exception when failed to update JIT server state
+                Js::Throw::OutOfMemory();
+            }
         }
     }
 }

lib/Backend/ServerScriptContext.cpp

@@ -5,16 +5,16 @@
 
 #include "Backend.h"
 
-ServerScriptContext::ServerScriptContext(ScriptContextDataIDL * contextData) :
+ServerScriptContext::ServerScriptContext(ScriptContextDataIDL * contextData, ServerThreadContext* threadContextInfo) :
     m_contextData(*contextData),
+    threadContextInfo(threadContextInfo),
     m_isPRNGSeeded(false),
-    m_isClosed(false),
     m_domFastPathHelperMap(nullptr),
     m_moduleRecords(&HeapAllocator::Instance),
 #ifdef PROFILE_EXEC
     m_codeGenProfiler(nullptr),
 #endif
-    m_activeJITCount(0)
+    m_refCount(0)
 {
 #ifdef PROFILE_EXEC
     if (Js::Configuration::Global.flags.IsEnabled(Js::ProfileFlag))
@@ -279,21 +279,19 @@ ServerScriptContext::Close()
 }
 
 void
-ServerScriptContext::BeginJIT()
+ServerScriptContext::AddRef()
 {
-    InterlockedExchangeAdd(&m_activeJITCount, 1u);
+    InterlockedExchangeAdd(&m_refCount, 1u);
 }
 
 void
-ServerScriptContext::EndJIT()
+ServerScriptContext::Release()
 {
-    InterlockedExchangeSubtract(&m_activeJITCount, 1u);
-}
-
-bool
-ServerScriptContext::IsJITActive()
-{
-    return m_activeJITCount != 0;
+    InterlockedExchangeSubtract(&m_refCount, 1u);
+    if (m_isClosed && m_refCount == 0)
+    {
+        HeapDelete(this);
+    }
 }
 
 Js::Var*
@@ -328,3 +326,4 @@ ServerScriptContext::GetCodeGenProfiler() const
     return nullptr;
 #endif
 }
+

lib/Backend/ServerScriptContext.h

@@ -8,7 +8,7 @@
 class ServerScriptContext : public ScriptContextInfo
 {
 public:
-    ServerScriptContext(ScriptContextDataIDL * contextData);
+    ServerScriptContext(ScriptContextDataIDL * contextData, ServerThreadContext* threadContextInfo);
     ~ServerScriptContext();
     virtual intptr_t GetNullAddr() const override;
     virtual intptr_t GetUndefinedAddr() const override;
@@ -62,11 +62,12 @@ public:
     void AddModuleRecordInfo(unsigned int moduleId, __int64 localExportSlotsAddr);
 
     Js::ScriptContextProfiler *  GetCodeGenProfiler() const;
+    ServerThreadContext* GetThreadContext() { return threadContextInfo; }
 
     void Close();
-    void BeginJIT();
-    void EndJIT();
-    bool IsJITActive();
+    void AddRef();
+    void Release();
+
 private:
     JITDOMFastPathHelperMap * m_domFastPathHelperMap;
 #ifdef PROFILE_EXEC
@@ -74,7 +75,9 @@ private:
 #endif
 
     ScriptContextDataIDL m_contextData;
-    uint m_activeJITCount;
+
+    ServerThreadContext* threadContextInfo;
+    uint m_refCount;
 
     bool m_isPRNGSeeded;
     bool m_isClosed;

lib/Backend/ServerThreadContext.cpp

@@ -260,3 +260,20 @@ ServerThreadContext::AddToPropertyMap(const Js::PropertyRecord * origRecord)
 
     PropertyRecordTrace(_u("Added JIT property '%s' at 0x%08x, pid = %d\n"), record->GetBuffer(), record, record->pid);
 }
+
+void ServerThreadContext::AddRef()
+{
+    InterlockedExchangeAdd(&m_refCount, (uint)1);
+}
+void ServerThreadContext::Release()
+{
+    InterlockedExchangeSubtract(&m_refCount, (uint)1);
+    if (m_isClosed && m_refCount == 0)
+    {
+        HeapDelete(this);
+    }
+}
+void ServerThreadContext::Close()
+{
+    this->m_isClosed = true;
+}

lib/Backend/ServerThreadContext.h

@@ -45,6 +45,11 @@ public:
     void RemoveFromPropertyMap(Js::PropertyId reclaimedId);
     void AddToPropertyMap(const Js::PropertyRecord * propertyRecord);
     void SetWellKnownHostTypeId(Js::TypeId typeId) { this->wellKnownHostTypeHTMLAllCollectionTypeId = typeId; }
+
+    void AddRef();
+    void Release();
+    void Close();    
+
 private:
     intptr_t GetRuntimeChakraBaseAddress() const;
     intptr_t GetRuntimeCRTBaseAddress() const;
@@ -60,8 +65,9 @@ private:
     CodeGenAllocators m_codeGenAlloc;
 
     ThreadContextDataIDL m_threadContextData;
-
-    ThreadContext * m_threadContext;
+    
     intptr_t m_jitChakraBaseAddress;
     intptr_t m_jitCRTBaseAddress;
+    uint m_refCount;
+
 };

lib/Common/Common.h

@@ -83,6 +83,7 @@ template<> struct IntMath<int64> { using Type = Int64Math; };
 #include "Exceptions/StackOverflowException.h"
 #include "Exceptions/NotImplementedException.h"
 #include "Exceptions/AsmJsParseException.h"
+#include "Exceptions/JITOperationFailedException.h"
 
 #include "Memory/AutoPtr.h"
 #include "Memory/AutoAllocatorObjectPtr.h"

lib/Common/Common/Jobs.cpp

@@ -14,6 +14,7 @@
 #include "Exceptions/OperationAbortedException.h"
 #include "Exceptions/OutOfMemoryException.h"
 #include "Exceptions/StackOverflowException.h"
+#include "Exceptions/JITOperationFailedException.h"
 
 #include "TemplateParameter.h"
 #include "DataStructures/DoublyLinkedListElement.h"
@@ -1004,6 +1005,14 @@ namespace JsUtil
             job->failureReason = Job::FailureReason::Aborted;
 #endif
         }
+        catch (Js::JITOperationFailedException)
+        {
+            // This can happen for any reason a job needs to be aborted while executing, like for instance, if the script
+            // context is closed while the job is being processed in the background
+#if ENABLE_DEBUG_CONFIG_OPTIONS
+            job->failureReason = Job::FailureReason::Unknown;
+#endif
+        }
 
         // Since the background job processor processes jobs on a background thread, out-of-memory and stack overflow need to be
         // caught here. Script would not be active in the background thread, so (at the time of this writing) a

lib/Common/Exceptions/Chakra.Common.Exceptions.vcxproj

@@ -47,6 +47,7 @@
     <ClInclude Include="ExceptionCheck.h" />
     <ClInclude Include="InScriptExceptionBase.h" />
     <ClInclude Include="InternalErrorException.h" />
+    <ClInclude Include="JITOperationFailedException.h" />
     <ClInclude Include="NotImplementedException.h" />
     <ClInclude Include="OperationAbortedException.h" />
     <ClInclude Include="OutOfMemoryException.h" />
@@ -58,4 +59,4 @@
   </ItemGroup>
   <Import Project="$(BuildConfigPropsPath)Chakra.Build.targets" Condition="exists('$(BuildConfigPropsPath)Chakra.Build.targets')" />
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.targets" />
-</Project>
+</Project>

lib/Common/Exceptions/JITOperationFailedException.h

@@ -0,0 +1,18 @@
+//-------------------------------------------------------------------------------------------------------
+// Copyright (C) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE.txt file in the project root for full license information.
+//-------------------------------------------------------------------------------------------------------
+#pragma once
+
+namespace Js {
+
+    class JITOperationFailedException : public ExceptionBase
+    {
+    public:
+        JITOperationFailedException(DWORD lastError)
+            :LastError(lastError)
+        {}
+        DWORD LastError;
+    };
+
+} // namespace Js

lib/Common/Exceptions/Throw.cpp

@@ -18,6 +18,7 @@
 #include "InternalErrorException.h"
 #include "OutOfMemoryException.h"
 #include "NotImplementedException.h"
+#include "JITOperationFailedException.h"
 
 // Header files required before including ConfigFlagsTable.h
 
@@ -121,6 +122,30 @@ namespace Js {
         throw StackOverflowException();
     }
 
+    void Throw::JITOperationFailed(DWORD lastError)
+    {
+#ifdef ENABLE_DEBUG_CONFIG_OPTIONS
+        if (CONFIG_FLAG(PrintSystemException))
+        {
+            Output::Print(_u("SystemException: JITOperationFailed\n"));
+            Output::Flush();
+        }
+#endif
+        throw JITOperationFailedException(lastError);
+    }
+
+    void Throw::CheckAndThrowJITOperationFailed()
+    {
+        DWORD lastError = GetLastError();
+        // currently this is used for virtual memory(Virtual*Ex) operations only
+        // which 0 indicate succeed
+        if (lastError != 0) 
+        {
+            Throw::JITOperationFailed(lastError);
+        }
+        
+    }
+
     void Throw::NotImplemented()
     {
         AssertMsg(false, "This functionality is not yet implemented");

lib/Common/Exceptions/Throw.h

@@ -21,6 +21,9 @@ namespace Js {
         static void __declspec(noreturn) InternalError();
         static void __declspec(noreturn) FatalInternalError();
         static void __declspec(noreturn) FatalProjectionError();
+        static void __declspec(noreturn) JITOperationFailed(DWORD lastError);
+
+        static void CheckAndThrowJITOperationFailed();
 
         static void CheckAndThrowOutOfMemory(BOOLEAN status);
 

lib/Common/Memory/CommonMemoryPch.h

@@ -20,6 +20,7 @@ typedef _Return_type_success_(return >= 0) LONG NTSTATUS;
 // Exceptions
 #include "Exceptions/ExceptionBase.h"
 #include "Exceptions/OutOfMemoryException.h"
+#include "Exceptions/JITOperationFailedException.h"
 
 // Other Memory headers
 #include "Memory/LeakReport.h"

lib/Common/Memory/CustomHeap.cpp

@@ -223,7 +223,11 @@ Allocation* Heap::Alloc(size_t bytes, ushort pdataCount, ushort xdataSize, bool
         {
             MEMORY_BASIC_INFORMATION memBasicInfo;
             size_t resultBytes = VirtualQueryEx(this->processHandle, allocation->address, &memBasicInfo, sizeof(memBasicInfo));
-            Assert(resultBytes != 0 && memBasicInfo.Protect == PAGE_EXECUTE);
+            if (resultBytes == 0)
+            {
+                Js::Throw::CheckAndThrowJITOperationFailed();
+            }
+            Assert(memBasicInfo.Protect == PAGE_EXECUTE);
         }
 #endif
         return allocation;
@@ -254,7 +258,11 @@ Allocation* Heap::Alloc(size_t bytes, ushort pdataCount, ushort xdataSize, bool
 #if defined(DBG)
         MEMORY_BASIC_INFORMATION memBasicInfo;
         size_t resultBytes = VirtualQueryEx(this->processHandle, page->address, &memBasicInfo, sizeof(memBasicInfo));
-        Assert(resultBytes != 0 && memBasicInfo.Protect == PAGE_EXECUTE);
+        if (resultBytes == 0)
+        {
+            Js::Throw::CheckAndThrowJITOperationFailed();
+        }
+        Assert(memBasicInfo.Protect == PAGE_EXECUTE);
 #endif
 
         Allocation* allocation = nullptr;

lib/Common/Memory/PageAllocator.cpp

@@ -210,7 +210,11 @@ PageSegmentBase<T>::Initialize(DWORD allocFlags, bool excludeGuardPages)
         {
             DWORD oldProtect;
             BOOL vpresult = VirtualProtectEx(this->allocator->processHandle, this->address, this->GetAvailablePageCount() * AutoSystemInfo::PageSize, PAGE_NOACCESS, &oldProtect);
-            Assert(vpresult && oldProtect == PAGE_READWRITE);
+            if (vpresult == FALSE)
+            {
+                Js::Throw::CheckAndThrowJITOperationFailed();
+            }
+            Assert(oldProtect == PAGE_READWRITE);
         }
         return true;
     }
@@ -294,6 +298,10 @@ PageSegmentBase<T>::AllocPages(uint pageCount)
 #ifdef PAGEALLOCATOR_PROTECT_FREEPAGE
                 DWORD oldProtect;
                 BOOL vpresult = VirtualProtectEx(this->allocator->processHandle, allocAddress, pageCount * AutoSystemInfo::PageSize, PAGE_READWRITE, &oldProtect);
+                if (vpresult == FALSE)
+                {
+                    Js::Throw::CheckAndThrowJITOperationFailed();
+                }
                 Assert(vpresult && oldProtect == PAGE_NOACCESS);
 #endif
             return allocAddress;
@@ -394,6 +402,10 @@ PageSegmentBase<T>::ReleasePages(__in void * address, uint pageCount)
 #ifdef PAGEALLOCATOR_PROTECT_FREEPAGE
     DWORD oldProtect;
     BOOL vpresult = VirtualProtectEx(this->allocator->processHandle, address, pageCount * AutoSystemInfo::PageSize, PAGE_NOACCESS, &oldProtect);
+    if (vpresult == FALSE)
+    {
+        Js::Throw::CheckAndThrowJITOperationFailed();
+    }
     Assert(vpresult && oldProtect == PAGE_READWRITE);
 #endif
 
@@ -2370,6 +2382,7 @@ HeapPageAllocator<T>::ProtectPages(__in char* address, size_t pageCount, __in vo
         || address < segment->GetAddress()
         || ((uint)(((char *)address) - segment->GetAddress()) > (segment->GetPageCount() - pageCount) * AutoSystemInfo::PageSize))
     {
+        // OOPJIT TODO: don't bring down the whole JIT process
         CustomHeap_BadPageState_fatal_error((ULONG_PTR)this);
         return FALSE;
     }
@@ -2378,6 +2391,13 @@ HeapPageAllocator<T>::ProtectPages(__in char* address, size_t pageCount, __in vo
 
     // check old protection on all pages about to change, ensure the fidelity
     size_t bytes = VirtualQueryEx(this->processHandle, address, &memBasicInfo, sizeof(memBasicInfo));
+    if (bytes == 0)
+    {
+        if (this->processHandle != GetCurrentProcess())
+        {
+            Js::Throw::CheckAndThrowJITOperationFailed();
+        }
+    }
     if (bytes == 0
         || memBasicInfo.RegionSize < pageCount * AutoSystemInfo::PageSize
         || desiredOldProtectFlag != memBasicInfo.Protect)
@@ -2392,6 +2412,7 @@ HeapPageAllocator<T>::ProtectPages(__in char* address, size_t pageCount, __in vo
         (dwVirtualProtectFlags & (PAGE_EXECUTE | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE)) &&
         ((dwVirtualProtectFlags & PAGE_TARGETS_NO_UPDATE) == 0))
     {
+        // OOPJIT TODO: don't bring down the whole JIT process
         CustomHeap_BadPageState_fatal_error((ULONG_PTR)this);
         return FALSE;
     }

lib/Common/Memory/VirtualAllocWrapper.cpp

@@ -44,13 +44,27 @@ LPVOID VirtualAllocWrapper::Alloc(LPVOID lpAddress, size_t dwSize, DWORD allocat
         {
             allocProtectFlags = PAGE_EXECUTE_READWRITE;
         }
+
         address = VirtualAllocEx(process, lpAddress, dwSize, allocationType, allocProtectFlags);
-        VirtualProtectEx(process, address, dwSize, protectFlags, &oldProtectFlags);
+        if (address == nullptr && process != GetCurrentProcess())
+        {
+            Js::Throw::CheckAndThrowJITOperationFailed();
+        }
+
+        BOOL result = VirtualProtectEx(process, address, dwSize, protectFlags, &oldProtectFlags);
+        if (result == FALSE && process != GetCurrentProcess())
+        {
+            Js::Throw::CheckAndThrowJITOperationFailed();
+        }
     }
     else
 #endif
     {
         address = VirtualAllocEx(process, lpAddress, dwSize, allocationType, protectFlags);
+        if (address == nullptr && process != GetCurrentProcess())
+        {
+            Js::Throw::CheckAndThrowJITOperationFailed();
+        }
     }
 
     return address;
@@ -62,7 +76,12 @@ BOOL VirtualAllocWrapper::Free(LPVOID lpAddress, size_t dwSize, DWORD dwFreeType
     AnalysisAssert(dwFreeType == MEM_RELEASE || dwFreeType == MEM_DECOMMIT);
     size_t bytes = (dwFreeType == MEM_RELEASE)? 0 : dwSize;
 #pragma warning(suppress: 28160) // Calling VirtualFreeEx without the MEM_RELEASE flag frees memory but not address descriptors (VADs)
-    return VirtualFreeEx(process, lpAddress, bytes, dwFreeType);
+    BOOL ret = VirtualFreeEx(process, lpAddress, bytes, dwFreeType);
+    if (ret == FALSE && process != GetCurrentProcess())
+    {
+        // OOP JIT TODO: check if we need to cleanup the context related to this content process
+    }
+    return ret;
 }
 
 /*
@@ -88,9 +107,9 @@ PreReservedVirtualAllocWrapper::~PreReservedVirtualAllocWrapper()
         BOOL success = VirtualFreeEx(processHandle, preReservedStartAddress, 0, MEM_RELEASE);
         PreReservedHeapTrace(_u("MEM_RELEASE the PreReservedSegment. Start Address: 0x%p, Size: 0x%x * 0x%x bytes"), preReservedStartAddress, PreReservedAllocationSegmentCount,
             AutoSystemInfo::Data.GetAllocationGranularityPageSize());
-        if (!success)
+        if (!success && this->processHandle != GetCurrentProcess())
         {
-            Assert(false);
+            // OOP JIT TODO: check if we need to cleanup the context related to this content process
         }
 
 #if !_M_X64_OR_ARM64 && _CONTROL_FLOW_GUARD
@@ -119,10 +138,14 @@ PreReservedVirtualAllocWrapper::IsInRange(void * address)
     //Check if the region is in MEM_COMMIT state.
     MEMORY_BASIC_INFORMATION memBasicInfo;
     size_t bytes = VirtualQueryEx(processHandle, address, &memBasicInfo, sizeof(memBasicInfo));
-    if (bytes == 0 || memBasicInfo.State != MEM_COMMIT)
+    if (bytes == 0)
     {
-        AssertMsg(false, "Memory not committed? Checking for uncommitted address region?");
+        if (this->processHandle != GetCurrentProcess())
+        {
+            Js::Throw::CheckAndThrowJITOperationFailed();
+        }
     }
+    AssertMsg(memBasicInfo.State == MEM_COMMIT, "Memory not committed? Checking for uncommitted address region?");
 #endif
     return result;
 }
@@ -291,6 +314,13 @@ LPVOID PreReservedVirtualAllocWrapper::Alloc(LPVOID lpAddress, size_t dwSize, DW
             //Check if the region is not already in MEM_COMMIT state.
             MEMORY_BASIC_INFORMATION memBasicInfo;
             size_t bytes = VirtualQueryEx(processHandle, addressToReserve, &memBasicInfo, sizeof(memBasicInfo));
+            if (bytes == 0)
+            {
+                if (this->processHandle != GetCurrentProcess())
+                {
+                    Js::Throw::CheckAndThrowJITOperationFailed();
+                }
+            }
             if (bytes == 0
                 || memBasicInfo.RegionSize < requestedNumOfSegments * AutoSystemInfo::Data.GetAllocationGranularityPageSize()
                 || memBasicInfo.State == MEM_COMMIT
@@ -345,7 +375,11 @@ LPVOID PreReservedVirtualAllocWrapper::Alloc(LPVOID lpAddress, size_t dwSize, DW
 
                 if (allocatedAddress != nullptr)
                 {
-                    VirtualProtectEx(processHandle, allocatedAddress, dwSize, protectFlags, &oldProtect);
+                    BOOL result = VirtualProtectEx(processHandle, allocatedAddress, dwSize, protectFlags, &oldProtect);
+                    if (result == FALSE && this->processHandle != GetCurrentProcess())
+                    {
+                        Js::Throw::CheckAndThrowJITOperationFailed();
+                    }
                     AssertMsg(oldProtect == (PAGE_EXECUTE_READWRITE), "CFG Bitmap gets allocated and bits will be set to invalid only upon passing these flags.");
                 }
             }
@@ -353,6 +387,10 @@ LPVOID PreReservedVirtualAllocWrapper::Alloc(LPVOID lpAddress, size_t dwSize, DW
 #endif
             {
                 allocatedAddress = (char *)VirtualAllocEx(processHandle, addressToReserve, dwSize, MEM_COMMIT, protectFlags);
+                if (allocatedAddress == nullptr && this->processHandle != GetCurrentProcess())
+                {
+                    Js::Throw::CheckAndThrowJITOperationFailed();
+                }
             }
         }
         else
@@ -425,6 +463,12 @@ PreReservedVirtualAllocWrapper::Free(LPVOID lpAddress, size_t dwSize, DWORD dwFr
             freeSegments.SetRange(freeSegmentsBVIndex, static_cast<uint>(requestedNumOfSegments));
             PreReservedHeapTrace(_u("MEM_RELEASE: Address: 0x%p of size: 0x%x * 0x%x bytes\n"), lpAddress, requestedNumOfSegments, AutoSystemInfo::Data.GetAllocationGranularityPageSize());
         }
+
+        if (success == FALSE && process != GetCurrentProcess())
+        {
+            // OOP JIT TODO: check if we need to cleanup the context related to this content process
+        }
+
         return success;
     }
 }

lib/JITClient/JITManager.cpp

@@ -443,6 +443,7 @@ JITManager::UpdatePropertyRecordMap(
 HRESULT
 JITManager::InitializeScriptContext(
     __in ScriptContextDataIDL * data,
+    __in intptr_t threadContextInfoAddress,
     __out intptr_t * scriptContextInfoAddress)
 {
     Assert(IsOOPJITEnabled());
@@ -450,7 +451,7 @@ JITManager::InitializeScriptContext(
     HRESULT hr = E_FAIL;
     RpcTryExcept
     {
-        hr = ClientInitializeScriptContext(m_rpcBindingHandle, data, scriptContextInfoAddress);
+        hr = ClientInitializeScriptContext(m_rpcBindingHandle, data, threadContextInfoAddress, scriptContextInfoAddress);
     }
         RpcExcept(RpcExceptionFilter(RpcExceptionCode()))
     {

lib/JITClient/JITManager.h

@@ -47,6 +47,7 @@ public:
 
     HRESULT InitializeScriptContext(
         __in ScriptContextDataIDL * data,
+        __in  intptr_t threadContextInfoAddress,
         __out intptr_t *scriptContextInfoAddress);
 
     HRESULT CleanupProcess();
@@ -78,6 +79,7 @@ public:
 
     HRESULT Shutdown();
 
+
     static JITManager * GetJITManager();
 private:
     JITManager();
@@ -96,6 +98,7 @@ private:
     bool m_isJITServer;
 
     static JITManager s_jitManager;
+
 };
 
 #else  // !ENABLE_OOP_NATIVE_CODEGEN
@@ -148,6 +151,7 @@ public:
 
     HRESULT InitializeScriptContext(
         __in ScriptContextDataIDL * data,
+        __in intptr_t threadContextInfoAddress,
         __out intptr_t *scriptContextInfoAddress)
         { Assert(false); return E_FAIL; }
 

lib/JITIDL/ChakraJIT.idl

@@ -52,6 +52,7 @@ interface IChakraJIT
     HRESULT InitializeScriptContext(
         [in] handle_t binding,
         [in] ScriptContextDataIDL * scriptContextData,
+        [in] CHAKRA_PTR threadContextInfoAddress,
         [out] CHAKRA_PTR * scriptContextInfoAddress);
 
     HRESULT CloseScriptContext(

lib/JITServer/Chakra.JITServer.vcxproj

@@ -28,6 +28,7 @@
     <ClCompile Include="$(MSBuildThisFileDirectory)JITServer.cpp" />
   </ItemGroup>
   <ItemGroup>
+    <ClInclude Include="JITServer.h" />
     <ClInclude Include="JITServerPch.h" />
   </ItemGroup>
   <ItemGroup>

lib/JITServer/Chakra.JITServer.vcxproj.filters

@@ -1,11 +1,12 @@
 <?xml version="1.0" encoding="utf-8"?>
 <Project ToolsVersion="4.0" xmlns="http://schemas.microsoft.com/developer/msbuild/2003">
   <ItemGroup>
-    <ClCompile Include="JITServerPch.cpp" />
-    <ClCompile Include="JITServerStub.c" />
-    <ClCompile Include="JITServer.cpp" />
+    <ClCompile Include="$(MSBuildThisFileDirectory)JITServerStub.c" />
+    <ClCompile Include="$(MSBuildThisFileDirectory)JITServerPch.cpp" />
+    <ClCompile Include="$(MSBuildThisFileDirectory)JITServer.cpp" />
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="JITServerPch.h" />
+    <ClInclude Include="JITServer.h" />
   </ItemGroup>
 </Project>

lib/JITServer/JITServer.cpp

@@ -131,7 +131,13 @@ ServerInitializeThreadContext(
 {
     AUTO_NESTED_HANDLED_EXCEPTION_TYPE(static_cast<ExceptionType>(ExceptionType_OutOfMemory | ExceptionType_StackOverflow));
 
-    ServerThreadContext * contextInfo = HeapNew(ServerThreadContext, threadContextData);
+    ServerThreadContext * contextInfo = HeapNewNoThrow(ServerThreadContext, threadContextData);
+    if (contextInfo == nullptr) 
+    {
+        return E_OUTOFMEMORY;
+    }
+
+    ServerContextManager::RegisterThreadContext(contextInfo);
 
     *threadContextRoot = (intptr_t)EncodePointer(contextInfo);
     *prereservedRegionAddr = (intptr_t)contextInfo->GetPreReservedVirtualAllocator()->EnsurePreReservedRegion();
@@ -152,8 +158,15 @@ ServerCleanupThreadContext(
         return RPC_S_INVALID_ARG;
     }
 
-    while (threadContextInfo->IsJITActive()) { Sleep(1); }
-    HeapDelete(threadContextInfo);
+    if (!ServerContextManager::IsThreadContextAlive(threadContextInfo))
+    {
+        return E_ACCESSDENIED;
+    }
+
+    AutoReleaseContext<ServerThreadContext> autoThreadContext(threadContextInfo);
+
+    threadContextInfo->Close();
+    ServerContextManager::UnRegisterThreadContext(threadContextInfo);
 
     return S_OK;
 }
@@ -173,17 +186,26 @@ ServerUpdatePropertyRecordMap(
         return RPC_S_INVALID_ARG;
     }
 
-    for (uint i = 0; i < updatedProps->reclaimedPropertyCount; ++i)
+    if (!ServerContextManager::IsThreadContextAlive(threadContextInfo))
     {
-        threadContextInfo->RemoveFromPropertyMap((Js::PropertyId)updatedProps->reclaimedPropertyIdArray[i]);
+        return E_ACCESSDENIED;
     }
 
-    for (uint i = 0; i < updatedProps->newRecordCount; ++i)
+    AutoReleaseContext<ServerThreadContext> autoThreadContext(threadContextInfo);
+    return ServerCallWrapper(threadContextInfo, [&]()->HRESULT
     {
-        threadContextInfo->AddToPropertyMap((Js::PropertyRecord *)updatedProps->newRecordArray[i]);
-    }
+        for (uint i = 0; i < updatedProps->reclaimedPropertyCount; ++i)
+        {
+            threadContextInfo->RemoveFromPropertyMap((Js::PropertyId)updatedProps->reclaimedPropertyIdArray[i]);
+        }
 
-    return S_OK;
+        for (uint i = 0; i < updatedProps->newRecordCount; ++i)
+        {
+            threadContextInfo->AddToPropertyMap((Js::PropertyRecord *)updatedProps->newRecordArray[i]);
+        }
+
+        return S_OK;
+    });
 }
 
 HRESULT
@@ -202,9 +224,17 @@ ServerAddDOMFastPathHelper(
         return RPC_S_INVALID_ARG;
     }
 
-    scriptContextInfo->AddToDOMFastPathHelperMap(funcInfoAddr, (IR::JnHelperMethod)helper);
+    if (!ServerContextManager::IsScriptContextAlive(scriptContextInfo))
+    {
+        return RPC_S_INVALID_ARG;
+    }
 
-    return S_OK;
+    AutoReleaseContext<ServerScriptContext> autoScriptContext(scriptContextInfo);
+    return ServerCallWrapper(scriptContextInfo, [&]()->HRESULT
+    {
+        scriptContextInfo->AddToDOMFastPathHelperMap(funcInfoAddr, (IR::JnHelperMethod)helper);
+        return S_OK;
+    });
 }
 
 HRESULT
@@ -221,10 +251,19 @@ ServerAddModuleRecordInfo(
     {
         return RPC_S_INVALID_ARG;
     }
-    serverScriptContext->AddModuleRecordInfo(moduleId, localExportSlotsAddr);
-    HRESULT hr = E_FAIL;
 
-    return hr;
+    if (!ServerContextManager::IsScriptContextAlive(serverScriptContext))
+    {
+        return RPC_S_INVALID_ARG;
+    }
+
+    AutoReleaseContext<ServerScriptContext> autoScriptContext(serverScriptContext);
+    return ServerCallWrapper(serverScriptContext, [&]()->HRESULT
+    {
+        serverScriptContext->AddModuleRecordInfo(moduleId, localExportSlotsAddr);
+        return S_OK;
+    });
+
 }
 
 HRESULT 
@@ -239,27 +278,56 @@ ServerSetWellKnownHostTypeId(
     {
         return RPC_S_INVALID_ARG;
     }
-    threadContextInfo->SetWellKnownHostTypeId((Js::TypeId)typeId);
 
-    return S_OK;
+    if (!ServerContextManager::IsThreadContextAlive(threadContextInfo))
+    {
+        return E_ACCESSDENIED;
+    }
+
+    AutoReleaseContext<ServerThreadContext> autoThreadContext(threadContextInfo);
+    return ServerCallWrapper(threadContextInfo, [&]()->HRESULT
+    {
+        threadContextInfo->SetWellKnownHostTypeId((Js::TypeId)typeId);
+        return S_OK;
+    });
+
 }
 
 HRESULT
 ServerInitializeScriptContext(
     /* [in] */ handle_t binding,
     /* [in] */ __RPC__in ScriptContextDataIDL * scriptContextData,
+    /* [in] */ __RPC__in intptr_t threadContextInfoAddress,
     /* [out] */ __RPC__out intptr_t * scriptContextInfoAddress)
 {
     AUTO_NESTED_HANDLED_EXCEPTION_TYPE(static_cast<ExceptionType>(ExceptionType_OutOfMemory | ExceptionType_StackOverflow));
 
-    ServerScriptContext * contextInfo = HeapNew(ServerScriptContext, scriptContextData);
-    *scriptContextInfoAddress = (intptr_t)EncodePointer(contextInfo);
+    ServerThreadContext * threadContextInfo = (ServerThreadContext*)DecodePointer((void*)threadContextInfoAddress);
+
+    if (threadContextInfo == nullptr)
+    {
+        return E_ACCESSDENIED;
+    }
+
+    if (!ServerContextManager::IsThreadContextAlive(threadContextInfo))
+    {
+        return E_ACCESSDENIED;
+    }
+
+    return ServerCallWrapper(threadContextInfo, [&]()->HRESULT
+    {
+        ServerScriptContext * contextInfo = HeapNew(ServerScriptContext, scriptContextData, threadContextInfo);
+
+        ServerContextManager::RegisterScriptContext(contextInfo);
+
+        *scriptContextInfoAddress = (intptr_t)EncodePointer(contextInfo);
 
 #if !FLOATVAR
-    // TODO: should move this to ServerInitializeThreadContext, also for the fields in IDL
-    XProcNumberPageSegmentImpl::Initialize(contextInfo->IsRecyclerVerifyEnabled(), contextInfo->GetRecyclerVerifyPad());
+        // TODO: should move this to ServerInitializeThreadContext, also for the fields in IDL
+        XProcNumberPageSegmentImpl::Initialize(contextInfo->IsRecyclerVerifyEnabled(), contextInfo->GetRecyclerVerifyPad());
 #endif
-    return S_OK;
+        return S_OK;
+    });
 }
 
 HRESULT
@@ -274,6 +342,13 @@ ServerCloseScriptContext(
         return RPC_S_INVALID_ARG;
     }
 
+    if (!ServerContextManager::IsScriptContextAlive(scriptContextInfo))
+    {
+        return E_ACCESSDENIED;
+    }
+
+    AutoReleaseContext<ServerScriptContext> autoScriptContext(scriptContextInfo);
+
 #ifdef PROFILE_EXEC
     auto profiler = scriptContextInfo->GetCodeGenProfiler();
     if (profiler && profiler->IsInitialized())
@@ -283,6 +358,7 @@ ServerCloseScriptContext(
 #endif
 
     scriptContextInfo->Close();
+    ServerContextManager::UnRegisterScriptContext(scriptContextInfo);
     return S_OK;
 }
 
@@ -298,8 +374,15 @@ ServerCleanupScriptContext(
         return RPC_S_INVALID_ARG;
     }
 
-    while (scriptContextInfo->IsJITActive()) { Sleep(1); }
-    HeapDelete(scriptContextInfo);
+    if (!ServerContextManager::IsScriptContextAlive(scriptContextInfo))
+    {
+        return E_ACCESSDENIED;
+    }
+
+    AutoReleaseContext<ServerScriptContext> autoScriptContext(scriptContextInfo);
+
+    scriptContextInfo->Close();
+    ServerContextManager::UnRegisterScriptContext(scriptContextInfo);
     return S_OK;
 }
 
@@ -317,11 +400,17 @@ ServerFreeAllocation(
         return RPC_S_INVALID_ARG;
     }
 
-    //DeRegister Entry Point for CFG
-    context->SetValidCallTargetForCFG((PVOID)address, false);
-
-    bool succeeded = context->GetCodeGenAllocators()->emitBufferManager.FreeAllocation((void*)address);
-    return succeeded ? S_OK : E_FAIL;
+    if (!ServerContextManager::IsThreadContextAlive(context))
+    {
+        return E_ACCESSDENIED;
+    }
+    AutoReleaseContext<ServerThreadContext> autoThreadContext(context);
+    return ServerCallWrapper(context, [&]()->HRESULT
+    {
+        context->SetValidCallTargetForCFG((PVOID)address, false);
+        bool succeeded = context->GetCodeGenAllocators()->emitBufferManager.FreeAllocation((void*)address);
+        return succeeded ? S_OK : E_FAIL;
+    });
 }
 
 HRESULT
@@ -339,22 +428,32 @@ ServerIsNativeAddr(
         return RPC_S_INVALID_ARG;
     }
 
-    PreReservedVirtualAllocWrapper *preReservedVirtualAllocWrapper = context->GetPreReservedVirtualAllocator();
-    if (preReservedVirtualAllocWrapper->IsInRange((void*)address))
-    {
-        *result = true;
-    }
-    else if (!context->IsAllJITCodeInPreReservedRegion())
-    {
-        CustomHeap::CodePageAllocators::AutoLock autoLock(context->GetCodePageAllocators());
-        *result = context->GetCodePageAllocators()->IsInNonPreReservedPageAllocator((void*)address);
-    }
-    else
+    if (!ServerContextManager::IsThreadContextAlive(context))
     {
         *result = false;
+        return E_ACCESSDENIED;
     }
 
-    return S_OK;
+    AutoReleaseContext<ServerThreadContext> autoThreadContext(context);
+    return ServerCallWrapper(context, [&]()->HRESULT
+    {
+        PreReservedVirtualAllocWrapper *preReservedVirtualAllocWrapper = context->GetPreReservedVirtualAllocator();
+        if (preReservedVirtualAllocWrapper->IsInRange((void*)address))
+        {
+            *result = true;
+        }
+        else if (!context->IsAllJITCodeInPreReservedRegion())
+        {
+            CustomHeap::CodePageAllocators::AutoLock autoLock(context->GetCodePageAllocators());
+            *result = context->GetCodePageAllocators()->IsInNonPreReservedPageAllocator((void*)address);
+        }
+        else
+        {
+            *result = false;
+        }
+
+        return S_OK;
+    });
 }
 
 HRESULT
@@ -364,9 +463,24 @@ ServerSetIsPRNGSeeded(
     /* [in] */ boolean value)
 {
     ServerScriptContext * scriptContextInfo = (ServerScriptContext*)DecodePointer((void*)scriptContextInfoAddress);
-    scriptContextInfo->SetIsPRNGSeeded(value != FALSE);
 
-    return S_OK;
+    if (scriptContextInfo == nullptr)
+    {
+        return RPC_S_INVALID_ARG;
+    }
+
+    if (!ServerContextManager::IsScriptContextAlive(scriptContextInfo))
+    {
+        return RPC_S_INVALID_ARG;
+    }
+
+    AutoReleaseContext<ServerScriptContext> autoScriptContext(scriptContextInfo);    
+
+    return ServerCallWrapper(scriptContextInfo, [&]()->HRESULT
+    {
+        scriptContextInfo->SetIsPRNGSeeded(value != FALSE);
+        return S_OK;
+    });
 }
 
 HRESULT
@@ -380,13 +494,6 @@ ServerRemoteCodeGen(
     UNREFERENCED_PARAMETER(binding);
     AUTO_NESTED_HANDLED_EXCEPTION_TYPE(static_cast<ExceptionType>(ExceptionType_OutOfMemory | ExceptionType_StackOverflow));
 
-    LARGE_INTEGER start_time = { 0 };
-    if (PHASE_TRACE1(Js::BackEndPhase))
-    {
-        QueryPerformanceCounter(&start_time);
-    }
-    memset(jitData, 0, sizeof(JITOutputIDL));
-
     ServerThreadContext * threadContextInfo = (ServerThreadContext*)DecodePointer((void*)threadContextInfoAddress);
     ServerScriptContext * scriptContextInfo = (ServerScriptContext*)DecodePointer((void*)scriptContextInfoAddress);
 
@@ -395,51 +502,64 @@ ServerRemoteCodeGen(
         return RPC_S_INVALID_ARG;
     }
 
-    NoRecoverMemoryJitArenaAllocator jitArena(L"JITArena", threadContextInfo->GetPageAllocator(), Js::Throw::OutOfMemory);
+    if (!ServerContextManager::IsThreadContextAlive(threadContextInfo))
+    {
+        return E_ACCESSDENIED;
+    }
 
-    scriptContextInfo->BeginJIT(); // TODO: OOP JIT, improve how we do this
-    threadContextInfo->BeginJIT();
+    if (!ServerContextManager::IsScriptContextAlive(scriptContextInfo))
+    {
+        return E_ACCESSDENIED;
+    }
 
-    JITTimeWorkItem * jitWorkItem = Anew(&jitArena, JITTimeWorkItem, workItemData);
+    AutoReleaseContext<ServerThreadContext> autoThreadContext(threadContextInfo);
+    AutoReleaseContext<ServerScriptContext> autoScriptContext(scriptContextInfo);
 
-    if (PHASE_VERBOSE_TRACE_RAW(Js::BackEndPhase, jitWorkItem->GetJITTimeInfo()->GetSourceContextId(), jitWorkItem->GetJITTimeInfo()->GetLocalFunctionId()))
+    return ServerCallWrapper(threadContextInfo, [&]() ->HRESULT
     {
-        LARGE_INTEGER freq;
-        LARGE_INTEGER end_time;
-        QueryPerformanceCounter(&end_time);
-        QueryPerformanceFrequency(&freq);
+        LARGE_INTEGER start_time = { 0 };
+        if (PHASE_TRACE1(Js::BackEndPhase))
+        {
+            QueryPerformanceCounter(&start_time);
+        }
+        memset(jitData, 0, sizeof(JITOutputIDL));
 
-        Output::Print(
-            L"BackendMarshalIn - function: %s time:%8.6f mSec\r\n",
-            jitWorkItem->GetJITFunctionBody()->GetDisplayName(),
-            (((double)((end_time.QuadPart - workItemData->startTime)* (double)1000.0 / (double)freq.QuadPart))) / (1));
-        Output::Flush();
-    }
+        NoRecoverMemoryJitArenaAllocator jitArena(L"JITArena", threadContextInfo->GetPageAllocator(), Js::Throw::OutOfMemory);
+        JITTimeWorkItem * jitWorkItem = Anew(&jitArena, JITTimeWorkItem, workItemData);
 
-    auto profiler = scriptContextInfo->GetCodeGenProfiler();
+        if (PHASE_VERBOSE_TRACE_RAW(Js::BackEndPhase, jitWorkItem->GetJITTimeInfo()->GetSourceContextId(), jitWorkItem->GetJITTimeInfo()->GetLocalFunctionId()))
+        {
+            LARGE_INTEGER freq;
+            LARGE_INTEGER end_time;
+            QueryPerformanceCounter(&end_time);
+            QueryPerformanceFrequency(&freq);
+
+            Output::Print(
+                L"BackendMarshalIn - function: %s time:%8.6f mSec\r\n",
+                jitWorkItem->GetJITFunctionBody()->GetDisplayName(),
+                (((double)((end_time.QuadPart - workItemData->startTime)* (double)1000.0 / (double)freq.QuadPart))) / (1));
+            Output::Flush();
+        }
+
+        auto profiler = scriptContextInfo->GetCodeGenProfiler();
 #ifdef PROFILE_EXEC
-    if (profiler && !profiler->IsInitialized())
-    {
-        profiler->Initialize(threadContextInfo->GetPageAllocator(), nullptr);
-    }
+        if (profiler && !profiler->IsInitialized())
+        {
+            profiler->Initialize(threadContextInfo->GetPageAllocator(), nullptr);
+        }
 #endif
-    if (jitWorkItem->GetWorkItemData()->xProcNumberPageSegment)
-    {
-        jitData->numberPageSegments = (XProcNumberPageSegment*)midl_user_allocate(sizeof(XProcNumberPageSegment));
-        if (!jitData->numberPageSegments)
+        if (jitWorkItem->GetWorkItemData()->xProcNumberPageSegment)
         {
-            scriptContextInfo->EndJIT();
-            threadContextInfo->EndJIT();
-
-            return E_OUTOFMEMORY;
+            jitData->numberPageSegments = (XProcNumberPageSegment*)midl_user_allocate(sizeof(XProcNumberPageSegment));
+            if (!jitData->numberPageSegments)
+            {
+                return E_OUTOFMEMORY;
+            }
+            __analysis_assume(jitData->numberPageSegments);
+
+            memcpy_s(jitData->numberPageSegments, sizeof(XProcNumberPageSegment), jitWorkItem->GetWorkItemData()->xProcNumberPageSegment, sizeof(XProcNumberPageSegment));
         }
-        __analysis_assume(jitData->numberPageSegments);
 
-        memcpy_s(jitData->numberPageSegments, sizeof(XProcNumberPageSegment), jitWorkItem->GetWorkItemData()->xProcNumberPageSegment, sizeof(XProcNumberPageSegment));
-    }
-    HRESULT hr = S_OK;
-    try
-    {
         Func::Codegen(
             &jitArena,
             jitWorkItem,
@@ -455,6 +575,105 @@ ServerRemoteCodeGen(
 #endif
             profiler,
             true);
+
+
+#ifdef PROFILE_EXEC
+        if (profiler && profiler->IsInitialized())
+        {
+            profiler->ProfilePrint(Js::Configuration::Global.flags.Profile.GetFirstPhase());
+        }
+#endif
+
+        if (PHASE_VERBOSE_TRACE_RAW(Js::BackEndPhase, jitWorkItem->GetJITTimeInfo()->GetSourceContextId(), jitWorkItem->GetJITTimeInfo()->GetLocalFunctionId()))
+        {
+            LARGE_INTEGER freq;
+            LARGE_INTEGER end_time;
+            QueryPerformanceCounter(&end_time);
+            QueryPerformanceFrequency(&freq);
+
+            Output::Print(
+                L"EndBackEndInner - function: %s time:%8.6f mSec\r\n",
+                jitWorkItem->GetJITFunctionBody()->GetDisplayName(),
+                (((double)((end_time.QuadPart - start_time.QuadPart)* (double)1000.0 / (double)freq.QuadPart))) / (1));
+            Output::Flush();
+
+        }
+        LARGE_INTEGER out_time = { 0 };
+        if (PHASE_TRACE1(Js::BackEndPhase))
+        {
+            QueryPerformanceCounter(&out_time);
+            jitData->startTime = out_time.QuadPart;
+        }
+
+        return S_OK;
+    });
+}
+
+JsUtil::BaseHashSet<ServerThreadContext*, HeapAllocator> ServerContextManager::threadContexts(&HeapAllocator::Instance);
+JsUtil::BaseHashSet<ServerScriptContext*, HeapAllocator> ServerContextManager::scriptContexts(&HeapAllocator::Instance);
+CriticalSection ServerContextManager::cs;
+
+void ServerContextManager::RegisterThreadContext(ServerThreadContext* threadContext)
+{
+    AutoCriticalSection autoCS(&cs);
+    threadContexts.Add(threadContext);
+}
+
+void ServerContextManager::UnRegisterThreadContext(ServerThreadContext* threadContext)
+{
+    AutoCriticalSection autoCS(&cs);
+    threadContexts.Remove(threadContext);
+    auto iter = scriptContexts.GetIteratorWithRemovalSupport();
+    while (iter.IsValid())
+    {
+        if (iter.Current().Key()->GetThreadContext() == threadContext)
+        {
+            iter.Current().Key()->Close();
+            iter.RemoveCurrent();
+        }
+        iter.MoveNext();
+    }
+}
+
+bool ServerContextManager::IsThreadContextAlive(ServerThreadContext* threadContext)
+{
+    AutoCriticalSection autoCS(&cs);
+    if (threadContexts.LookupWithKey(threadContext)) 
+    {
+        return !threadContext->IsClosed();
+    }
+    return false;
+}
+
+void ServerContextManager::RegisterScriptContext(ServerScriptContext* scriptContext)
+{
+    AutoCriticalSection autoCS(&cs);
+    scriptContexts.Add(scriptContext);
+}
+
+void ServerContextManager::UnRegisterScriptContext(ServerScriptContext* scriptContext)
+{
+    AutoCriticalSection autoCS(&cs);
+    scriptContexts.Remove(scriptContext);
+}
+
+bool ServerContextManager::IsScriptContextAlive(ServerScriptContext* scriptContext)
+{
+    AutoCriticalSection autoCS(&cs);
+    if (scriptContexts.LookupWithKey(scriptContext))
+    {
+        return !scriptContext->IsClosed();
+    }
+    return false;
+}
+
+template<typename Fn>
+HRESULT ServerCallWrapper(ServerThreadContext* threadContextInfo, Fn fn)
+{
+    HRESULT hr = S_OK;
+    try
+    {
+        hr = fn();
     }
     catch (Js::OutOfMemoryException)
     {
@@ -468,35 +687,42 @@ ServerRemoteCodeGen(
     {
         hr = E_ABORT;
     }
-    scriptContextInfo->EndJIT();
-    threadContextInfo->EndJIT();
-
-#ifdef PROFILE_EXEC
-    if (profiler && profiler->IsInitialized())
+    catch (Js::JITOperationFailedException& ex)
     {
-        profiler->ProfilePrint(Js::Configuration::Global.flags.Profile.GetFirstPhase());
-    }
-#endif
+        hr = HRESULT_FROM_WIN32(ex.LastError);
 
-    if (PHASE_VERBOSE_TRACE_RAW(Js::BackEndPhase, jitWorkItem->GetJITTimeInfo()->GetSourceContextId(), jitWorkItem->GetJITTimeInfo()->GetLocalFunctionId()))
-    {
-        LARGE_INTEGER freq;
-        LARGE_INTEGER end_time;
-        QueryPerformanceCounter(&end_time);
-        QueryPerformanceFrequency(&freq);
+        if (hr == E_ACCESSDENIED)
+        {
+            // target process might be terminated
+            DWORD exitCode;
+            if (GetExitCodeProcess(threadContextInfo->GetProcessHandle(), &exitCode))
+            {
+                if (exitCode != STILL_ACTIVE)
+                {
+                    threadContextInfo->Close();
+                    ServerContextManager::UnRegisterThreadContext(threadContextInfo);
+
+                }
+            }
+        }
+        else
+        {
+            // OOPJIT TODO: other error code might need to be handled
+            Assert(false);
+        }
+    }
 
-        Output::Print(
-            L"EndBackEndInner - function: %s time:%8.6f mSec\r\n",
-            jitWorkItem->GetJITFunctionBody()->GetDisplayName(),
-            (((double)((end_time.QuadPart - start_time.QuadPart)* (double)1000.0 / (double)freq.QuadPart))) / (1));
-        Output::Flush();
+    return hr;
+}
 
-    }
-    LARGE_INTEGER out_time = { 0 };
-    if (PHASE_TRACE1(Js::BackEndPhase))
+template<typename Fn>
+HRESULT ServerCallWrapper(ServerScriptContext* scriptContextInfo, Fn fn)
+{
+    ServerThreadContext* threadContextInfo = scriptContextInfo->GetThreadContext();
+    if (!ServerContextManager::IsThreadContextAlive(threadContextInfo))
     {
-        QueryPerformanceCounter(&out_time);
-        jitData->startTime = out_time.QuadPart;
+        return E_ACCESSDENIED;
     }
-    return hr;
-}
+    AutoReleaseContext<ServerThreadContext> autoThreadContext(threadContextInfo);
+    return ServerCallWrapper(threadContextInfo, fn);
+}

lib/JITServer/JITServer.h

@@ -0,0 +1,43 @@
+
+//-------------------------------------------------------------------------------------------------------
+// Copyright (C) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE.txt file in the project root for full license information.
+//-------------------------------------------------------------------------------------------------------
+
+class ServerContextManager
+{
+public:
+    static void RegisterThreadContext(ServerThreadContext* threadContext);
+    static void UnRegisterThreadContext(ServerThreadContext* threadContext);
+    static bool IsThreadContextAlive(ServerThreadContext* threadContext);
+
+    static void RegisterScriptContext(ServerScriptContext* scriptContext);
+    static void UnRegisterScriptContext(ServerScriptContext* scriptContext);    
+    static bool IsScriptContextAlive(ServerScriptContext* scriptContext);
+private:
+    static JsUtil::BaseHashSet<ServerThreadContext*, HeapAllocator> threadContexts;
+    static JsUtil::BaseHashSet<ServerScriptContext*, HeapAllocator> scriptContexts;
+    static CriticalSection cs;
+};
+
+template<class T>
+struct AutoReleaseContext
+{
+    AutoReleaseContext(T* context)
+        :context(context)
+    {
+        context->AddRef();
+    }
+
+    ~AutoReleaseContext()
+    {
+        context->Release();
+    }
+
+    T* context;
+};
+
+template<typename Fn>
+HRESULT ServerCallWrapper(ServerThreadContext* threadContextInfo, Fn fn);
+template<typename Fn>
+HRESULT ServerCallWrapper(ServerScriptContext* scriptContextInfo, Fn fn);

lib/JITServer/JITServerPch.h

@@ -11,3 +11,4 @@
 
 #include "Runtime.h"
 #include "Backend.h"
+#include "JITServer.h"

lib/Runtime/Base/ScriptContext.cpp

@@ -4495,7 +4495,20 @@ void ScriptContext::RegisterPrototypeChainEnsuredToHaveOnlyWritableDataPropertie
         {
             contextData.vtableAddresses[i] = vtblAddresses[i];
         }
-        JITManager::GetJITManager()->InitializeScriptContext(&contextData, &m_remoteScriptContextAddr);
+
+        bool allowPrereserveAlloc = true;
+#if !_M_X64_OR_ARM64
+        if (this->webWorkerId != Js::Constants::NonWebWorkerContextId)
+        {
+            allowPrereserveAlloc = false;
+        }
+#endif
+#ifndef _CONTROL_FLOW_GUARD
+        allowPrereserveAlloc = false;
+#endif
+        this->GetThreadContext()->EnsureJITThreadContext(allowPrereserveAlloc);
+
+        JITManager::GetJITManager()->InitializeScriptContext(&contextData, this->GetThreadContext()->GetRemoteThreadContextAddr(), &m_remoteScriptContextAddr);
     }
 #endif
 

lib/Runtime/Base/ThreadContext.cpp

@@ -1979,6 +1979,7 @@ ThreadContext::EnsureJITThreadContext(bool allowPrereserveAlloc)
     {
         return;
     }
+
     ThreadContextDataIDL contextData;
     contextData.processHandle = (intptr_t)JITManager::GetJITManager()->GetJITTargetHandle();
 

lib/Runtime/Base/ThreadContextInfo.cpp

@@ -12,8 +12,7 @@
 
 ThreadContextInfo::ThreadContextInfo() :
     m_isAllJITCodeInPreReservedRegion(true),
-    wellKnownHostTypeHTMLAllCollectionTypeId(Js::TypeIds_Undefined),
-    m_activeJITCount(0)
+    wellKnownHostTypeHTMLAllCollectionTypeId(Js::TypeIds_Undefined)
 {
 }
 
@@ -411,25 +410,12 @@ ThreadContextInfo::SetValidCallTargetForCFG(PVOID callTargetAddress, bool isSetV
 #endif // _CONTROL_FLOW_GUARD
 }
 
-void
-ThreadContextInfo::BeginJIT()
-{
-    InterlockedExchangeAdd(&m_activeJITCount, (uint)1);
-}
-
-void
-ThreadContextInfo::EndJIT()
+bool 
+ThreadContextInfo::IsClosed()
 {
-    InterlockedExchangeSubtract(&m_activeJITCount, (uint)1);
+    return m_isClosed;
 }
 
-bool
-ThreadContextInfo::IsJITActive()
-{
-    return m_activeJITCount != 0;
-}
-
-
 intptr_t SHIFT_ADDR(const ThreadContextInfo*const context, intptr_t address)
 {
 #if ENABLE_OOP_NATIVE_CODEGEN

lib/Runtime/Base/ThreadContextInfo.h

@@ -100,9 +100,8 @@ public:
     bool CanBeFalsy(Js::TypeId typeId) { return typeId == this->wellKnownHostTypeHTMLAllCollectionTypeId; }
 
     bool IsCFGEnabled();
-    void BeginJIT();
-    void EndJIT();
-    bool IsJITActive();
+    bool IsClosed();
+    
 
 #if defined(ENABLE_GLOBALIZATION) && defined(_CONTROL_FLOW_GUARD)
     Js::DelayLoadWinCoreMemory * GetWinCoreMemoryLibrary();
@@ -113,10 +112,9 @@ public:
 #endif
 protected:
     Js::TypeId wellKnownHostTypeHTMLAllCollectionTypeId;
-private:
 
-    uint m_activeJITCount;
     bool m_isAllJITCodeInPreReservedRegion;
+    bool m_isClosed;
 
 };