Ensure Environment.ExitCode works correctly after app domain unloaded. (#10842)

* Ensure Environment.ExitCode works correctly after app domain unloaded.

This PR addresses 2 problems of Environment.ExitCode:

1. Can't get correct exit code of main function.
2. Can't set %errorlevel%.

Details can be found on #6206

Fix #6206

9 files changed, 106 insertions, 21 deletions

diff --git a/src/coreclr/hosts/inc/coreclrhost.h b/src/coreclr/hosts/inc/coreclrhost.h
index f0d7952aa6..dd11cb6a51 100644
--- a/src/coreclr/hosts/inc/coreclrhost.h
+++ b/src/coreclr/hosts/inc/coreclrhost.h
@@ -29,6 +29,11 @@ CORECLR_HOSTING_API(coreclr_shutdown,
             void* hostHandle,
             unsigned int domainId);
 
+CORECLR_HOSTING_API(coreclr_shutdown_2,
+            void* hostHandle,
+            unsigned int domainId,
+            int* latchedExitCode);
+
 CORECLR_HOSTING_API(coreclr_create_delegate,
             void* hostHandle,
             unsigned int domainId,
diff --git a/src/coreclr/hosts/unixcoreruncommon/coreruncommon.cpp b/src/coreclr/hosts/unixcoreruncommon/coreruncommon.cpp
index d40fb424e6..52ffda8bb5 100644
--- a/src/coreclr/hosts/unixcoreruncommon/coreruncommon.cpp
+++ b/src/coreclr/hosts/unixcoreruncommon/coreruncommon.cpp
@@ -321,7 +321,7 @@ int ExecuteManagedAssembly(
     {
         coreclr_initialize_ptr initializeCoreCLR = (coreclr_initialize_ptr)dlsym(coreclrLib, "coreclr_initialize");
         coreclr_execute_assembly_ptr executeAssembly = (coreclr_execute_assembly_ptr)dlsym(coreclrLib, "coreclr_execute_assembly");
-        coreclr_shutdown_ptr shutdownCoreCLR = (coreclr_shutdown_ptr)dlsym(coreclrLib, "coreclr_shutdown");
+        coreclr_shutdown_2_ptr shutdownCoreCLR = (coreclr_shutdown_2_ptr)dlsym(coreclrLib, "coreclr_shutdown_2");
 
         if (initializeCoreCLR == nullptr)
         {
@@ -333,7 +333,7 @@ int ExecuteManagedAssembly(
         }
         else if (shutdownCoreCLR == nullptr)
         {
-            fprintf(stderr, "Function coreclr_shutdown not found in the libcoreclr.so\n");
+            fprintf(stderr, "Function coreclr_shutdown_2 not found in the libcoreclr.so\n");
         }
         else
         {
@@ -416,12 +416,18 @@ int ExecuteManagedAssembly(
                     exitCode = -1;
                 }
 
-                st = shutdownCoreCLR(hostHandle, domainId);
+                int latchedExitCode = 0;
+                st = shutdownCoreCLR(hostHandle, domainId, &latchedExitCode);
                 if (!SUCCEEDED(st))
                 {
                     fprintf(stderr, "coreclr_shutdown failed - status: 0x%08x\n", st);
                     exitCode = -1;
                 }
+
+                if (exitCode != -1)
+                {
+                    exitCode = latchedExitCode;
+                }
             }
         }
 
diff --git a/src/dlls/mscoree/mscorwks_ntdef.src b/src/dlls/mscoree/mscorwks_ntdef.src
index 8115475418..d7e6a2dcf4 100644
--- a/src/dlls/mscoree/mscorwks_ntdef.src
+++ b/src/dlls/mscoree/mscorwks_ntdef.src
@@ -21,6 +21,7 @@ EXPORTS
         coreclr_execute_assembly
         coreclr_initialize
         coreclr_shutdown
+        coreclr_shutdown_2
 
         ; il{d}asm
         MetaDataGetDispenser
diff --git a/src/dlls/mscoree/mscorwks_unixexports.src b/src/dlls/mscoree/mscorwks_unixexports.src
index f7862d3afe..28e9ac223f 100644
--- a/src/dlls/mscoree/mscorwks_unixexports.src
+++ b/src/dlls/mscoree/mscorwks_unixexports.src
@@ -3,6 +3,7 @@ coreclr_create_delegate
 coreclr_execute_assembly
 coreclr_initialize
 coreclr_shutdown
+coreclr_shutdown_2
 
 ; il{d}asm
 MetaDataGetDispenser
diff --git a/src/dlls/mscoree/unixinterface.cpp b/src/dlls/mscoree/unixinterface.cpp
index edd361c0c2..cf9bbc5c72 100644
--- a/src/dlls/mscoree/unixinterface.cpp
+++ b/src/dlls/mscoree/unixinterface.cpp
@@ -183,9 +183,9 @@ int coreclr_initialize(
     }
 #endif
 
-    ReleaseHolder<ICLRRuntimeHost2> host;
+    ReleaseHolder<ICLRRuntimeHost4> host;
 
-    hr = CorHost2::CreateObject(IID_ICLRRuntimeHost2, (void**)&host);
+    hr = CorHost2::CreateObject(IID_ICLRRuntimeHost4, (void**)&host);
     IfFailRet(hr);
 
     ConstWStringHolder appDomainFriendlyNameW = StringToUnicode(appDomainFriendlyName);
@@ -284,7 +284,7 @@ int coreclr_shutdown(
             void* hostHandle,
             unsigned int domainId)
 {
-    ReleaseHolder<ICLRRuntimeHost2> host(reinterpret_cast<ICLRRuntimeHost2*>(hostHandle));
+    ReleaseHolder<ICLRRuntimeHost4> host(reinterpret_cast<ICLRRuntimeHost4*>(hostHandle));
 
     HRESULT hr = host->UnloadAppDomain(domainId, true); // Wait until done
     IfFailRet(hr);
@@ -299,6 +299,37 @@ int coreclr_shutdown(
 }
 
 //
+// Shutdown CoreCLR. It unloads the app domain and stops the CoreCLR host.
+//
+// Parameters:
+//  hostHandle              - Handle of the host
+//  domainId                - Id of the domain
+//  latchedExitCode         - Latched exit code after domain unloaded
+//
+// Returns:
+//  HRESULT indicating status of the operation. S_OK if the assembly was successfully executed
+//
+extern "C"
+int coreclr_shutdown_2(
+            void* hostHandle,
+            unsigned int domainId,
+            int* latchedExitCode)
+{
+    ReleaseHolder<ICLRRuntimeHost4> host(reinterpret_cast<ICLRRuntimeHost4*>(hostHandle));
+
+    HRESULT hr = host->UnloadAppDomain2(domainId, true, latchedExitCode); // Wait until done
+    IfFailRet(hr);
+
+    hr = host->Stop();
+
+#ifdef FEATURE_PAL
+    PAL_Shutdown();
+#endif
+
+    return hr;
+}
+
+//
 // Create a native callable delegate for a managed method.
 //
 // Parameters:
@@ -321,7 +352,7 @@ int coreclr_create_delegate(
             const char* entryPointMethodName,
             void** delegate)
 {
-    ICLRRuntimeHost2* host = reinterpret_cast<ICLRRuntimeHost2*>(hostHandle);
+    ICLRRuntimeHost4* host = reinterpret_cast<ICLRRuntimeHost4*>(hostHandle);
 
     ConstWStringHolder entryPointAssemblyNameW = StringToUnicode(entryPointAssemblyName);
     ConstWStringHolder entryPointTypeNameW = StringToUnicode(entryPointTypeName);
@@ -366,7 +397,7 @@ int coreclr_execute_assembly(
     }
     *exitCode = -1;
 
-    ICLRRuntimeHost2* host = reinterpret_cast<ICLRRuntimeHost2*>(hostHandle);
+    ICLRRuntimeHost4* host = reinterpret_cast<ICLRRuntimeHost4*>(hostHandle);
 
     ConstWStringArrayHolder argvW;
     argvW.Set(StringArrayToUnicode(argc, argv), argc);
diff --git a/src/inc/corhost.h b/src/inc/corhost.h
index 3aabe9ed5d..59ab23cd27 100644
--- a/src/inc/corhost.h
+++ b/src/inc/corhost.h
@@ -137,6 +137,7 @@ protected:
 
     STDMETHODIMP UnloadAppDomain(DWORD dwDomainId, BOOL fWaitUntilDone);
 
+    STDMETHODIMP UnloadAppDomain2(DWORD dwDomainId, BOOL fWaitUntilDone, int *pLatchedExitCode);
 public:
     static ULONG GetHostVersion()
     {
@@ -275,7 +276,7 @@ class CorHost2 :
 #ifndef FEATURE_PAL    
     , public IPrivateManagedExceptionReporting /* This interface is for internal Watson testing only*/
 #endif // FEATURE_PAL    
-    , public ICLRRuntimeHost2
+    , public ICLRRuntimeHost4
     , public CorExecutionManager
 {
     friend struct _DacGlobals;
@@ -337,6 +338,8 @@ public:
 
     STDMETHODIMP UnloadAppDomain(DWORD dwDomainId, BOOL fWaitUntilDone);
 
+    STDMETHODIMP UnloadAppDomain2(DWORD dwDomainId, BOOL fWaitUntilDone, int *pLatchedExitCode);
+
     STDMETHODIMP GetCurrentAppDomainId(DWORD *pdwAppDomainId);
 
     STDMETHODIMP ExecuteApplication(LPCWSTR  pwzAppFullName,
diff --git a/src/pal/prebuilt/inc/mscoree.h b/src/pal/prebuilt/inc/mscoree.h
index 12d2172a85..29f7b261d3 100644
--- a/src/pal/prebuilt/inc/mscoree.h
+++ b/src/pal/prebuilt/inc/mscoree.h
@@ -112,8 +112,13 @@ typedef interface ICLRRuntimeHost ICLRRuntimeHost;
 #define __ICLRRuntimeHost2_FWD_DEFINED__
 typedef interface ICLRRuntimeHost2 ICLRRuntimeHost2;
 
-#endif 	/* __ICLRRuntimeHost2_FWD_DEFINED__ */
+#endif 	/* __ICLRRuntimeHost4_FWD_DEFINED__ */
 
+#ifndef __ICLRRuntimeHost4_FWD_DEFINED__
+#define __ICLRRuntimeHost4_FWD_DEFINED__
+typedef interface ICLRRuntimeHost4 ICLRRuntimeHost4;
+
+#endif  /* __ICLRRuntimeHost4_FWD_DEFINED__ */
 
 #ifndef __ICLRExecutionManager_FWD_DEFINED__
 #define __ICLRExecutionManager_FWD_DEFINED__
@@ -254,6 +259,7 @@ EXTERN_GUID(IID_ICLRErrorReportingManager, 0x980d2f1a, 0xbf79, 0x4c08, 0x81, 0x2
 EXTERN_GUID(IID_ICLRErrorReportingManager2, 0xc68f63b1, 0x4d8b, 0x4e0b, 0x95, 0x64, 0x9d, 0x2e, 0xfe, 0x2f, 0xa1, 0x8c);
 EXTERN_GUID(IID_ICLRRuntimeHost, 0x90F1A06C, 0x7712, 0x4762, 0x86, 0xB5, 0x7A, 0x5E, 0xBA, 0x6B, 0xDB, 0x02);
 EXTERN_GUID(IID_ICLRRuntimeHost2, 0x712AB73F, 0x2C22, 0x4807, 0xAD, 0x7E, 0xF5, 0x01, 0xD7, 0xb7, 0x2C, 0x2D);
+EXTERN_GUID(IID_ICLRRuntimeHost4, 0x64F6D366, 0xD7C2, 0x4F1F, 0xB4, 0xB2, 0xE8, 0x16, 0x0C, 0xAC, 0x43, 0xAF);
 EXTERN_GUID(IID_ICLRExecutionManager, 0x1000A3E7, 0xB420, 0x4620, 0xAE, 0x30, 0xFB, 0x19, 0xB5, 0x87, 0xAD, 0x1D);
 EXTERN_GUID(IID_ITypeName, 0xB81FF171, 0x20F3, 0x11d2, 0x8d, 0xcc, 0x00, 0xa0, 0xc9, 0xb0, 0x05, 0x22);
 EXTERN_GUID(IID_ITypeNameBuilder, 0xB81FF171, 0x20F3, 0x11d2, 0x8d, 0xcc, 0x00, 0xa0, 0xc9, 0xb0, 0x05, 0x23);
@@ -1819,6 +1825,14 @@ EXTERN_C const IID IID_ICLRRuntimeHost2;
         
     };
     
+    MIDL_INTERFACE("64F6D366-D7C2-4F1F-B4B2-E8160CAC43AF")
+    ICLRRuntimeHost4 : public ICLRRuntimeHost2
+    {
+        virtual HRESULT STDMETHODCALLTYPE UnloadAppDomain2(
+            /* [in] */ DWORD dwAppDomainId,
+            /* [in] */ BOOL fWaitUntilDone,
+            /* [out] */ int *pLatchedExitCode) = 0;
+    };
     
 #else 	/* C style interface */
 
diff --git a/src/vm/assembly.cpp b/src/vm/assembly.cpp
index 75430644c3..92c1ebd817 100644
--- a/src/vm/assembly.cpp
+++ b/src/vm/assembly.cpp
@@ -1849,10 +1849,7 @@ HRESULT RunMain(MethodDesc *pFD ,
         else 
         {
             *pParam->piRetVal = (INT32)threadStart.Call_RetArgSlot(&stackVar);
-            if (pParam->stringArgs == NULL) 
-            {
-                SetLatchedExitCode(*pParam->piRetVal);
-            }
+            SetLatchedExitCode(*pParam->piRetVal);
         }
 
         GCPROTECT_END();
diff --git a/src/vm/corhost.cpp b/src/vm/corhost.cpp
index 75adbada94..d935ddd8c8 100644
--- a/src/vm/corhost.cpp
+++ b/src/vm/corhost.cpp
@@ -1201,6 +1201,11 @@ HRESULT  GetCLRRuntimeHost(REFIID riid, IUnknown **ppUnk)
 
 STDMETHODIMP CorHost2::UnloadAppDomain(DWORD dwDomainId, BOOL fWaitUntilDone)
 {
+    return UnloadAppDomain2(dwDomainId, fWaitUntilDone, nullptr);
+}
+
+STDMETHODIMP CorHost2::UnloadAppDomain2(DWORD dwDomainId, BOOL fWaitUntilDone, int *pLatchedExitCode)
+{
     WRAPPER_NO_CONTRACT;
     STATIC_CONTRACT_SO_TOLERANT;
 
@@ -1249,14 +1254,23 @@ STDMETHODIMP CorHost2::UnloadAppDomain(DWORD dwDomainId, BOOL fWaitUntilDone)
         }
         END_ENTRYPOINT_NOTHROW;
 
+        if (pLatchedExitCode)
+        {
+            *pLatchedExitCode = GetLatchedExitCode();
+        }
+
         return hr;
     }
-    else
 
-    return CorRuntimeHostBase::UnloadAppDomain(dwDomainId, fWaitUntilDone);
+    return CorRuntimeHostBase::UnloadAppDomain2(dwDomainId, fWaitUntilDone, pLatchedExitCode);
 }
 
-HRESULT CorRuntimeHostBase::UnloadAppDomain(DWORD dwDomainId, BOOL fSync)
+HRESULT CorRuntimeHostBase::UnloadAppDomain(DWORD dwDomainId, BOOL fWaitUntilDone)
+{
+    return UnloadAppDomain2(dwDomainId, fWaitUntilDone, nullptr);
+}
+
+HRESULT CorRuntimeHostBase::UnloadAppDomain2(DWORD dwDomainId, BOOL fWaitUntilDone, int *pLatchedExitCode)
 {
     CONTRACTL
     {
@@ -1282,7 +1296,7 @@ HRESULT CorRuntimeHostBase::UnloadAppDomain(DWORD dwDomainId, BOOL fSync)
         //
         // However, for a thread that holds the loader lock, unloading the appDomain is
         // not a supported scenario. Thus, we should not be ending up in this code
-        // path for the FAULT violation. 
+        // path for the FAULT violation.
         //
         // Hence, the CONTRACT_VIOLATION below for overriding the FORBID_FAULT
         // for this scope only.
@@ -1292,18 +1306,23 @@ HRESULT CorRuntimeHostBase::UnloadAppDomain(DWORD dwDomainId, BOOL fSync)
         )
         {
             return HOST_E_CLRNOTAVAILABLE;
-        }   
+        }
     }
-    
+
     BEGIN_ENTRYPOINT_NOTHROW;
 
     // We do not use BEGIN_EXTERNAL_ENTRYPOINT here because
     // we do not want to setup Thread.  Process may be OOM, and we want Unload
     // to work.
-    hr =  AppDomain::UnloadById(ADID(dwDomainId), fSync);
+    hr =  AppDomain::UnloadById(ADID(dwDomainId), fWaitUntilDone);
 
     END_ENTRYPOINT_NOTHROW;
 
+    if (pLatchedExitCode)
+    {
+        *pLatchedExitCode = GetLatchedExitCode();
+    }
+
     return hr;
 }
 
@@ -1404,6 +1423,14 @@ HRESULT CorHost2::QueryInterface(REFIID riid, void **ppUnk)
 
         *ppUnk = static_cast<ICLRRuntimeHost2 *>(this);
     }
+    else if (riid == IID_ICLRRuntimeHost4)
+    {
+        ULONG version = 4;
+        if (m_Version == 0)
+            FastInterlockCompareExchange((LONG*)&m_Version, version, 0);
+
+        *ppUnk = static_cast<ICLRRuntimeHost4 *>(this);
+    }
     else if (riid == IID_ICLRExecutionManager)
     {
         ULONG version = 2;