Add ProcessContexts for OOP JIT

lib/Backend/ServerThreadContext.cpp

@@ -8,27 +8,26 @@
 #if ENABLE_OOP_NATIVE_CODEGEN
 #include "JITServer/JITServer.h"
 
-ServerThreadContext::ServerThreadContext(ThreadContextDataIDL * data, HANDLE processHandle) :
-    m_autoProcessHandle(processHandle),
-    m_processHandle(processHandle),
+ServerThreadContext::ServerThreadContext(ThreadContextDataIDL* data, ProcessContext* processContext) :
     m_threadContextData(*data),
     m_refCount(0),
     m_numericPropertyBV(nullptr),
-    m_preReservedSectionAllocator(processHandle),
-    m_sectionAllocator(processHandle),
-    m_thunkPageAllocators(nullptr, /* allocXData */ false, &m_sectionAllocator, nullptr, processHandle),
-    m_codePageAllocators(nullptr, ALLOC_XDATA, &m_sectionAllocator, &m_preReservedSectionAllocator, processHandle),
+    m_preReservedSectionAllocator(processContext->processHandle),
+    m_sectionAllocator(processContext->processHandle),
+    m_thunkPageAllocators(nullptr, /* allocXData */ false, &m_sectionAllocator, nullptr, processContext->processHandle),
+    m_codePageAllocators(nullptr, ALLOC_XDATA, &m_sectionAllocator, &m_preReservedSectionAllocator, processContext->processHandle),
 #if defined(_CONTROL_FLOW_GUARD) && !defined(_M_ARM)
-    m_jitThunkEmitter(this, &m_sectionAllocator, processHandle),
+    m_jitThunkEmitter(this, &m_sectionAllocator, processContext->processHandle),
 #endif
-    m_codeGenAlloc(nullptr, nullptr, this, &m_codePageAllocators, processHandle),
+    m_codeGenAlloc(nullptr, nullptr, this, &m_codePageAllocators, processContext->processHandle),
     m_pageAlloc(nullptr, Js::Configuration::Global.flags, PageAllocatorType_BGJIT,
         AutoSystemInfo::Data.IsLowMemoryProcess() ?
         PageAllocator::DefaultLowMaxFreePageCount :
         PageAllocator::DefaultMaxFreePageCount
-    )
+    ),
+    processContext(processContext)
 {
-    m_pid = GetProcessId(processHandle);
+    m_pid = GetProcessId(processContext->processHandle);
 
 #if !TARGET_64 && _CONTROL_FLOW_GUARD
     m_codeGenAlloc.canCreatePreReservedSegment = data->allowPrereserveAlloc != FALSE;
@@ -38,6 +37,7 @@ ServerThreadContext::ServerThreadContext(ThreadContextDataIDL * data, HANDLE pro
 
 ServerThreadContext::~ServerThreadContext()
 {
+    processContext->Release();
     if (this->m_numericPropertyBV != nullptr)
     {
         HeapDelete(m_numericPropertyBV);
@@ -112,7 +112,7 @@ ServerThreadContext::IsThreadBound() const
 HANDLE
 ServerThreadContext::GetProcessHandle() const
 {
-    return m_autoProcessHandle.GetHandle();
+    return this->processContext->processHandle;
 }
 
 CustomHeap::OOPCodePageAllocators *
@@ -150,13 +150,13 @@ ServerThreadContext::GetJITThunkEmitter()
 intptr_t
 ServerThreadContext::GetRuntimeChakraBaseAddress() const
 {
-    return static_cast<intptr_t>(m_threadContextData.chakraBaseAddress);
+    return this->processContext->chakraBaseAddress;
 }
 
 intptr_t
 ServerThreadContext::GetRuntimeCRTBaseAddress() const
 {
-    return static_cast<intptr_t>(m_threadContextData.crtBaseAddress);
+    return this->processContext->crtBaseAddress;
 }
 
 /* static */
@@ -215,4 +215,31 @@ void ServerThreadContext::Close()
     ServerContextManager::RecordCloseContext(this);
 #endif
 }
+
+ProcessContext::ProcessContext(HANDLE processHandle, intptr_t chakraBaseAddress, intptr_t crtBaseAddress) :
+    processHandle(processHandle),
+    chakraBaseAddress(chakraBaseAddress),
+    crtBaseAddress(crtBaseAddress),
+    refCount(0)
+{
+}
+ProcessContext::~ProcessContext()
+{
+    CloseHandle(processHandle);
+}
+
+void ProcessContext::AddRef()
+{
+    InterlockedExchangeAdd(&this->refCount, 1);
+}
+void ProcessContext::Release()
+{
+    InterlockedExchangeSubtract(&this->refCount, 1);
+}
+
+bool ProcessContext::HasRef()
+{
+    return this->refCount != 0;
+}
+
 #endif

lib/Backend/ServerThreadContext.h

@@ -5,13 +5,31 @@
 
 #pragma once
 
+#if ENABLE_OOP_NATIVE_CODEGEN
+class ProcessContext
+{
+private:
+    uint refCount;
+
+public:
+    HANDLE processHandle;
+    intptr_t chakraBaseAddress;
+    intptr_t crtBaseAddress;
+
+    ProcessContext(HANDLE processHandle, intptr_t chakraBaseAddress, intptr_t crtBaseAddress);
+    ~ProcessContext();
+    void AddRef();
+    void Release();
+    bool HasRef();
+
+};
+
 class ServerThreadContext : public ThreadContextInfo
 {
-#if ENABLE_OOP_NATIVE_CODEGEN
 public:
     typedef BVSparseNode<JitArenaAllocator> BVSparseNode;
 
-    ServerThreadContext(ThreadContextDataIDL * data, HANDLE processHandle);
+    ServerThreadContext(ThreadContextDataIDL * data, ProcessContext* processContext);
     ~ServerThreadContext();
 
     virtual HANDLE GetProcessHandle() const override;
@@ -58,7 +76,7 @@ public:
     static intptr_t GetJITCRTBaseAddress();
 
 private:
-    AutoCloseHandle m_autoProcessHandle;
+    ProcessContext* processContext;
 
     BVSparse<HeapAllocator> * m_numericPropertyBV;
 
@@ -72,12 +90,11 @@ private:
 #endif
     // only allocate with this from foreground calls (never from CodeGen calls)
     PageAllocator m_pageAlloc;
-    HANDLE m_processHandle;
     ThreadContextDataIDL m_threadContextData;
 
     DWORD m_pid; //save client process id for easier diagnose
 
     CriticalSection m_cs;
     uint m_refCount;
-#endif
 };
+#endif

lib/JITClient/JITManager.cpp

@@ -249,6 +249,9 @@ JITManager::ConnectRpcServer(__in HANDLE jitProcessHandle, __in_opt void* server
 
     m_jitConnectionId = connectionUuid;
 
+    hr = ConnectProcess();
+    HandleServerCallResult(hr, RemoteCallType::StateUpdate);
+
     return hr;
 
 FailureCleanup:
@@ -286,11 +289,45 @@ JITManager::Shutdown()
 }
 
 HRESULT
-JITManager::InitializeThreadContext(
-    __in ThreadContextDataIDL * data,
+JITManager::ConnectProcess()
+{
+    Assert(IsOOPJITEnabled());
+
+#ifdef USE_RPC_HANDLE_MARSHALLING
+    HANDLE processHandle;
+    if (!DuplicateHandle(GetCurrentProcess(), GetCurrentProcess(), GetCurrentProcess(), &processHandle, 0, false, DUPLICATE_SAME_ACCESS))
+    {
+        return false;
+    }
+#endif
+
+    HRESULT hr = E_FAIL;
+    RpcTryExcept
+    {
+        hr = ClientConnectProcess(
+            m_rpcBindingHandle,
 #ifdef USE_RPC_HANDLE_MARSHALLING
-    __in HANDLE processHandle,
+            processHandle,
 #endif
+            (intptr_t)AutoSystemInfo::Data.GetChakraBaseAddr(),
+            (intptr_t)AutoSystemInfo::Data.GetCRTHandle());
+    }
+        RpcExcept(RpcExceptionFilter(RpcExceptionCode()))
+    {
+        hr = HRESULT_FROM_WIN32(RpcExceptionCode());
+    }
+    RpcEndExcept;
+
+#ifdef USE_RPC_HANDLE_MARSHALLING
+    CloseHandle(processHandle);
+#endif
+
+    return hr;
+}
+
+HRESULT
+JITManager::InitializeThreadContext(
+    __in ThreadContextDataIDL * data,
     __out PPTHREADCONTEXT_HANDLE threadContextInfoAddress,
     __out intptr_t * prereservedRegionAddr,
     __out intptr_t * jitThunkAddr)
@@ -303,9 +340,6 @@ JITManager::InitializeThreadContext(
         hr = ClientInitializeThreadContext(
             m_rpcBindingHandle,
             data,
-#ifdef USE_RPC_HANDLE_MARSHALLING
-            processHandle,
-#endif
             threadContextInfoAddress,
             prereservedRegionAddr,
             jitThunkAddr);

lib/JITClient/JITManager.h

@@ -33,9 +33,6 @@ public:
 
     HRESULT InitializeThreadContext(
         __in ThreadContextDataIDL * data,
-#ifdef USE_RPC_HANDLE_MARSHALLING
-        __in HANDLE processHandle,
-#endif
         __out PPTHREADCONTEXT_HANDLE threadContextInfoAddress,
         __out intptr_t * prereservedRegionAddr,
         __out intptr_t * jitThunkAddr);
@@ -122,6 +119,8 @@ private:
         __in UUID* connectionUuid,
         __out RPC_BINDING_HANDLE* bindingHandle);
 
+    HRESULT ConnectProcess();
+
     RPC_BINDING_HANDLE m_rpcBindingHandle;
     UUID m_jitConnectionId;
     bool m_oopJitEnabled;
@@ -148,9 +147,6 @@ public:
 
     HRESULT InitializeThreadContext(
         __in ThreadContextDataIDL * data,
-#ifdef USE_RPC_HANDLE_MARSHALLING
-        __in HANDLE processHandle,
-#endif
         __out PPTHREADCONTEXT_HANDLE threadContextInfoAddress,
         __out intptr_t *prereservedRegionAddr,
         __out intptr_t * jitThunkAddr)

lib/JITIDL/ChakraJIT.idl

@@ -25,16 +25,23 @@ interface IChakraJIT
 {
     HRESULT Shutdown([in] handle_t binding);
 
-    HRESULT InitializeThreadContext(
+    HRESULT ConnectProcess(
         [in] handle_t binding,
-        [in] ThreadContextDataIDL * threadData,
 #ifdef USE_RPC_HANDLE_MARSHALLING
         [in, system_handle(sh_process, PROCESS_VM_OPERATION | PROCESS_VM_READ | PROCESS_VM_WRITE | PROCESS_QUERY_LIMITED_INFORMATION)] HANDLE processHandle,
 #endif
+        [in] CHAKRA_PTR chakraBaseAddress,
+        [in] CHAKRA_PTR crtBaseAddress
+    );
+
+    HRESULT InitializeThreadContext(
+        [in] handle_t binding,
+        [in] ThreadContextDataIDL * threadData,
         [out] PPTHREADCONTEXT_HANDLE threadContextInfoAddress,
         [out] CHAKRA_PTR * prereservedRegionAddr,
         [out] CHAKRA_PTR * jitThunkAddr);
 
+
     HRESULT CleanupThreadContext(
         [in] handle_t binding,
         [in, out] PPTHREADCONTEXT_HANDLE threadContextInfoAddress);

lib/JITIDL/JITTypes.h

@@ -314,8 +314,6 @@ typedef struct ThreadContextDataIDL
 
     IDL_PAD2(0)
     X64_PAD4(1)
-    CHAKRA_PTR chakraBaseAddress;
-    CHAKRA_PTR crtBaseAddress;
     CHAKRA_PTR threadStackLimitAddr;
     CHAKRA_PTR scriptStackLimit;
     CHAKRA_PTR bailOutRegisterSaveSpaceAddr;

lib/JITServer/JITServer.cpp

@@ -187,15 +187,54 @@ HRESULT CheckModuleAddress(HANDLE process, LPCVOID remoteImageBase, LPCVOID loca
     return S_OK;
 }
 
+HRESULT
+ServerConnectProcess(
+    handle_t binding,
+#ifdef USE_RPC_HANDLE_MARSHALLING
+    HANDLE processHandle,
+#endif
+    intptr_t chakraBaseAddress,
+    intptr_t crtBaseAddress
+)
+{
+    DWORD clientPid;
+    HRESULT hr = HRESULT_FROM_WIN32(I_RpcBindingInqLocalClientPID(binding, &clientPid));
+    if (FAILED(hr))
+    {
+        return hr;
+    }
+#ifdef USE_RPC_HANDLE_MARSHALLING
+    HANDLE targetHandle;
+    if (!DuplicateHandle(GetCurrentProcess(), processHandle, GetCurrentProcess(), &targetHandle, 0, false, DUPLICATE_SAME_ACCESS))
+    {
+        return E_ACCESSDENIED;
+    }
+#else
+    HANDLE targetHandle = OpenProcess(PROCESS_VM_OPERATION | PROCESS_VM_READ | PROCESS_VM_WRITE | PROCESS_QUERY_LIMITED_INFORMATION, false, clientPid);
+    if (!targetHandle)
+    {
+        return E_ACCESSDENIED;
+    }
+#endif
+    hr = CheckModuleAddress(targetHandle, (LPCVOID)chakraBaseAddress, (LPCVOID)AutoSystemInfo::Data.dllLoadAddress);
+    if (FAILED(hr))
+    {
+        return hr;
+    }
+    hr = CheckModuleAddress(targetHandle, (LPCVOID)crtBaseAddress, (LPCVOID)AutoSystemInfo::Data.GetCRTHandle());
+    if (FAILED(hr))
+    {
+        return hr;
+    }
+    return ProcessContextManager::RegisterNewProcess(clientPid, targetHandle, chakraBaseAddress, crtBaseAddress);
+}
+
 #pragma warning(push)
 #pragma warning(disable:6387 28196) // PREFast does not understand the out context can be null here
 HRESULT
 ServerInitializeThreadContext(
     /* [in] */ handle_t binding,
     /* [in] */ __RPC__in ThreadContextDataIDL * threadContextData,
-#ifdef USE_RPC_HANDLE_MARSHALLING
-    /* [in] */ __RPC__in HANDLE processHandle,
-#endif
     /* [out] */ __RPC__deref_out_opt PPTHREADCONTEXT_HANDLE threadContextInfoAddress,
     /* [out] */ __RPC__out intptr_t *prereservedRegionAddr,
     /* [out] */ __RPC__out intptr_t *jitThunkAddr)
@@ -218,28 +257,23 @@ ServerInitializeThreadContext(
     {
         return hr;
     }
-#ifdef USE_RPC_HANDLE_MARSHALLING
-    HANDLE targetHandle;
-    if (!DuplicateHandle(GetCurrentProcess(), processHandle, GetCurrentProcess(), &targetHandle, 0, false, DUPLICATE_SAME_ACCESS))
-    {
-        return E_ACCESSDENIED;
-    }
-#else
-    HANDLE targetHandle = OpenProcess(PROCESS_VM_OPERATION | PROCESS_VM_READ | PROCESS_VM_WRITE | PROCESS_QUERY_LIMITED_INFORMATION, false, clientPid);
-    if (!targetHandle)
+    ProcessContext* processContext = ProcessContextManager::GetProcessContext(clientPid);
+    if (processContext == nullptr)
     {
         return E_ACCESSDENIED;
     }
-#endif
     try
     {
         AUTO_NESTED_HANDLED_EXCEPTION_TYPE(static_cast<ExceptionType>(ExceptionType_OutOfMemory));
-        contextInfo = HeapNew(ServerThreadContext, threadContextData, targetHandle);
+        contextInfo = HeapNew(ServerThreadContext, threadContextData, processContext);
         ServerContextManager::RegisterThreadContext(contextInfo);
     }
     catch (Js::OutOfMemoryException)
     {
-        CloseHandle(targetHandle);
+        if (!contextInfo)
+        {
+            processContext->Release();
+        }
         return E_OUTOFMEMORY;
     }
 
@@ -249,16 +283,6 @@ ServerInitializeThreadContext(
         {
             return E_ACCESSDENIED;
         }
-        hr = CheckModuleAddress(targetHandle, (LPCVOID)contextInfo->GetRuntimeChakraBaseAddress(), (LPCVOID)AutoSystemInfo::Data.dllLoadAddress);
-        if (FAILED(hr))
-        {
-            return hr;
-        }
-        hr = CheckModuleAddress(targetHandle, (LPCVOID)contextInfo->GetRuntimeCRTBaseAddress(), (LPCVOID)contextInfo->GetJITCRTBaseAddress());
-        if (FAILED(hr))
-        {
-            return hr;
-        }
 
         *threadContextInfoAddress = (PTHREADCONTEXT_HANDLE)EncodePointer(contextInfo);
 
@@ -834,6 +858,60 @@ JsUtil::BaseHashSet<ServerThreadContext*, HeapAllocator> ServerContextManager::t
 JsUtil::BaseHashSet<ServerScriptContext*, HeapAllocator> ServerContextManager::scriptContexts(&HeapAllocator::Instance);
 CriticalSection ServerContextManager::cs;
 
+BaseDictionary<DWORD, ProcessContext*, HeapAllocator> ProcessContextManager::ProcessContexts(&HeapAllocator::Instance);
+CriticalSection ProcessContextManager::cs;
+
+HRESULT
+ProcessContextManager::RegisterNewProcess(DWORD pid, HANDLE processHandle, intptr_t chakraBaseAddress, intptr_t crtBaseAddress)
+{
+    AutoCriticalSection autoCS(&cs);
+    auto iter = ProcessContexts.GetIteratorWithRemovalSupport();
+    while (iter.IsValid())
+    {
+        ProcessContext* context = iter.CurrentValue();
+        // We can delete a ProcessContext if no ThreadContexts refer to it and the process is terminated
+        if (!context->HasRef() && WaitForSingleObject(context->processHandle, 0) != WAIT_TIMEOUT)
+        {
+            iter.RemoveCurrent();
+            HeapDelete(context);
+        }
+    }
+    // We cannot register multiple ProcessContexts for a single process
+    if (ProcessContexts.ContainsKey(pid))
+    {
+        return E_ACCESSDENIED;
+    }
+
+    try
+    {
+        AUTO_NESTED_HANDLED_EXCEPTION_TYPE(static_cast<ExceptionType>(ExceptionType_OutOfMemory));
+
+        ProcessContext* context = HeapNew(ProcessContext, processHandle, chakraBaseAddress, crtBaseAddress);
+        ProcessContexts.Add(pid, context);
+    }
+    catch (Js::OutOfMemoryException)
+    {
+        return E_OUTOFMEMORY;
+    }
+
+    return S_OK;
+}
+
+ProcessContext*
+ProcessContextManager::GetProcessContext(DWORD pid)
+{
+    AutoCriticalSection autoCS(&cs);
+    ProcessContext* context = nullptr;
+    // It is possible that we don't have a ProcessContext for a pid in case ProcessContext initialization failed,
+    // or if the calling process terminated and the ProcessContext was already cleaned up before we got here
+    if (ProcessContexts.ContainsKey(pid))
+    {
+        context = ProcessContexts.Item(pid);
+        context->AddRef();
+    }
+    return context;
+}
+
 #ifdef STACK_BACK_TRACE
 SList<ServerContextManager::ClosedContextEntry<ServerThreadContext>*, NoThrowHeapAllocator> ServerContextManager::ClosedThreadContextList(&NoThrowHeapAllocator::Instance);
 SList<ServerContextManager::ClosedContextEntry<ServerScriptContext>*, NoThrowHeapAllocator> ServerContextManager::ClosedScriptContextList(&NoThrowHeapAllocator::Instance);

lib/JITServer/JITServer.h

@@ -3,10 +3,23 @@
 // Licensed under the MIT license. See LICENSE.txt file in the project root for full license information.
 //-------------------------------------------------------------------------------------------------------
 
+class ProcessContextManager
+{
+private:
+
+    static BaseDictionary<DWORD, ProcessContext*, HeapAllocator> ProcessContexts;
+    static CriticalSection cs;
+
+public:
+    static HRESULT RegisterNewProcess(DWORD pid, HANDLE processHandle, intptr_t chakraBaseAddress, intptr_t crtBaseAddress);
+    static ProcessContext* GetProcessContext(DWORD pid);
+};
+
 class ServerContextManager
 {
 public:
     static void RegisterThreadContext(ServerThreadContext* threadContext);
+
     static void UnRegisterThreadContext(ServerThreadContext* threadContext);
 
     static void RegisterScriptContext(ServerScriptContext* scriptContext);

lib/Runtime/Base/ThreadContext.cpp

@@ -1986,18 +1986,7 @@ ThreadContext::EnsureJITThreadContext(bool allowPrereserveAlloc)
         return true;
     }
 
-#ifdef USE_RPC_HANDLE_MARSHALLING
-    HANDLE processHandle;
-    if (!DuplicateHandle(GetCurrentProcess(), GetCurrentProcess(), GetCurrentProcess(), &processHandle, 0, false, DUPLICATE_SAME_ACCESS))
-    {
-        return false;
-    }
-    AutoCloseHandle autoClose(processHandle);
-#endif
-
     ThreadContextDataIDL contextData;
-    contextData.chakraBaseAddress = (intptr_t)AutoSystemInfo::Data.GetChakraBaseAddr();
-    contextData.crtBaseAddress = (intptr_t)AutoSystemInfo::Data.GetCRTHandle();
     contextData.threadStackLimitAddr = reinterpret_cast<intptr_t>(GetAddressOfStackLimitForCurrentThread());
     contextData.bailOutRegisterSaveSpaceAddr = (intptr_t)bailOutRegisterSaveSpace;
     contextData.disableImplicitFlagsAddr = (intptr_t)GetAddressOfDisableImplicitFlags();
@@ -2022,9 +2011,6 @@ ThreadContext::EnsureJITThreadContext(bool allowPrereserveAlloc)
 
     HRESULT hr = JITManager::GetJITManager()->InitializeThreadContext(
         &contextData,
-#ifdef USE_RPC_HANDLE_MARSHALLING
-        processHandle,
-#endif
         &m_remoteThreadContextInfo,
         &m_prereservedRegionAddr,
         &m_jitThunkStartAddr);