Add AppWaitForSingleObject/MultipleObjects, ThreadWaitForCompletion.

8 files changed, 109 insertions, 104 deletions

diff --git a/src/burn/engine/apply.cpp b/src/burn/engine/apply.cpp
index e5b978f1..dd99a5cd 100644
--- a/src/burn/engine/apply.cpp
+++ b/src/burn/engine/apply.cpp
@@ -2333,6 +2333,8 @@ static HRESULT DoExecuteAction(
 
     HRESULT hr = S_OK;
     HANDLE rghWait[2] = { };
+    DWORD dwSignaledIndex = 0;
+    DWORD dwExitCode = 0;
     BOOTSTRAPPER_APPLY_RESTART restart = BOOTSTRAPPER_APPLY_RESTART_NONE;
     BOOL fRetry = FALSE;
     BOOL fStopWusaService = FALSE;
@@ -2356,27 +2358,23 @@ static HRESULT DoExecuteAction(
             // wait for cache sync-point
             rghWait[0] = pExecuteAction->waitCachePackage.pPackage->hCacheEvent;
             rghWait[1] = pContext->pApplyContext->hCacheThread;
-            switch (::WaitForMultipleObjects(rghWait[1] ? 2 : 1, rghWait, FALSE, INFINITE))
+
+            hr = AppWaitForMultipleObjects(rghWait[1] ? 2 : 1, rghWait, FALSE, INFINITE, &dwSignaledIndex);
+            ExitOnFailure(hr, "Failed to wait for cache check-point.");
+
+            switch (dwSignaledIndex)
             {
-            case WAIT_OBJECT_0:
+            case 0:
                 break;
-
-            case WAIT_OBJECT_0 + 1:
-                if (!::GetExitCodeThread(pContext->pApplyContext->hCacheThread, (DWORD*)&hr))
+            case 1:
+                if (!::GetExitCodeThread(pContext->pApplyContext->hCacheThread, &dwExitCode))
                 {
                     ExitWithLastError(hr, "Failed to get cache thread exit code.");
                 }
 
-                if (SUCCEEDED(hr))
-                {
-                    hr = E_UNEXPECTED;
-                }
-                ExitOnFailure(hr, "Cache thread exited unexpectedly.");
-
-            case WAIT_FAILED: __fallthrough;
-            default:
-                ExitWithLastError(hr, "Failed to wait for cache check-point.");
+                ExitWithRootFailure(hr, E_UNEXPECTED, "Cache thread exited unexpectedly with exit code: %u.", dwExitCode);
             }
+
             break;
 
         case BURN_EXECUTE_ACTION_TYPE_RELATED_BUNDLE:
diff --git a/src/burn/engine/cabextract.cpp b/src/burn/engine/cabextract.cpp
index 5a02ff8a..56146a39 100644
--- a/src/burn/engine/cabextract.cpp
+++ b/src/burn/engine/cabextract.cpp
@@ -262,10 +262,8 @@ extern "C" HRESULT CabExtractClose(
         }
 
         // wait for thread to terminate
-        if (WAIT_OBJECT_0 != ::WaitForSingleObject(pContext->Cabinet.hThread, INFINITE))
-        {
-            ExitWithLastError(hr, "Failed to wait for thread to terminate.");
-        }
+        hr = AppWaitForSingleObject(pContext->Cabinet.hThread, INFINITE);
+        ExitOnFailure(hr, "Failed to wait for thread to terminate.");
     }
 
 LExit:
@@ -306,29 +304,31 @@ static HRESULT WaitForOperation(
 {
     HRESULT hr = S_OK;
     HANDLE rghWait[2] = { };
+    DWORD dwSignaledIndex = 0;
 
     // wait for operation complete event
     rghWait[0] = pContext->Cabinet.hOperationCompleteEvent;
     rghWait[1] = pContext->Cabinet.hThread;
-    switch (::WaitForMultipleObjects(countof(rghWait), rghWait, FALSE, INFINITE))
+
+    hr = AppWaitForMultipleObjects(countof(rghWait), rghWait, FALSE, INFINITE, &dwSignaledIndex);
+    ExitOnFailure(hr, "Failed to wait for operation complete event.");
+
+    switch (dwSignaledIndex)
     {
-    case WAIT_OBJECT_0:
+    case 0:
         if (!::ResetEvent(pContext->Cabinet.hOperationCompleteEvent))
         {
             ExitWithLastError(hr, "Failed to reset operation complete event.");
         }
-        break;
 
-    case WAIT_OBJECT_0 + 1:
+        break;
+    case 1:
         if (!::GetExitCodeThread(pContext->Cabinet.hThread, (DWORD*)&hr))
         {
             ExitWithLastError(hr, "Failed to get extraction thread exit code.");
         }
-        ExitFunction();
 
-    case WAIT_FAILED: __fallthrough;
-    default:
-        ExitWithLastError(hr, "Failed to wait for operation complete event.");
+        ExitFunction();
     }
 
     // clear operation
@@ -430,10 +430,8 @@ static DWORD WINAPI ExtractThreadProc(
     }
 
     // wait for begin operation event
-    if (WAIT_FAILED == ::WaitForSingleObject(pContext->Cabinet.hBeginOperationEvent, INFINITE))
-    {
-        ExitWithLastError(hr, "Failed to wait for begin operation event.");
-    }
+    hr = AppWaitForSingleObject(pContext->Cabinet.hBeginOperationEvent, INFINITE);
+    ExitOnFailure(hr, "Failed to wait for begin operation event.");
 
     if (!::ResetEvent(pContext->Cabinet.hBeginOperationEvent))
     {
@@ -517,10 +515,8 @@ static INT_PTR CopyFileCallback(
     }
 
     // wait for begin operation event
-    if (WAIT_FAILED == ::WaitForSingleObject(pContext->Cabinet.hBeginOperationEvent, INFINITE))
-    {
-        ExitWithLastError(hr, "Failed to wait for begin operation event.");
-    }
+    hr = AppWaitForSingleObject(pContext->Cabinet.hBeginOperationEvent, INFINITE);
+    ExitOnFailure(hr, "Failed to wait for begin operation event.");
 
     if (!::ResetEvent(pContext->Cabinet.hBeginOperationEvent))
     {
@@ -552,10 +548,8 @@ static INT_PTR CopyFileCallback(
     }
 
     // wait for begin operation event
-    if (WAIT_FAILED == ::WaitForSingleObject(pContext->Cabinet.hBeginOperationEvent, INFINITE))
-    {
-        ExitWithLastError(hr, "Failed to wait for begin operation event.");
-    }
+    hr = AppWaitForSingleObject(pContext->Cabinet.hBeginOperationEvent, INFINITE);
+    ExitOnFailure(hr, "Failed to wait for begin operation event.");
 
     if (!::ResetEvent(pContext->Cabinet.hBeginOperationEvent))
     {
diff --git a/src/burn/engine/core.cpp b/src/burn/engine/core.cpp
index c1e3a12c..716e5af1 100644
--- a/src/burn/engine/core.cpp
+++ b/src/burn/engine/core.cpp
@@ -12,6 +12,9 @@ struct BURN_CACHE_THREAD_CONTEXT
 };
 
 
+static PFN_PROCWAITFORCOMPLETION vpfnProcWaitForCompletion = ProcWaitForCompletion;
+
+
 // internal function declarations
 
 static HRESULT CoreRecreateCommandLine(
@@ -65,9 +68,6 @@ static HRESULT DetectPackagePayloadsCached(
 static DWORD WINAPI CacheThreadProc(
     __in LPVOID lpThreadParameter
     );
-static HRESULT WaitForCacheThread(
-    __in HANDLE hCacheThread
-    );
 static void LogPackages(
     __in_opt const BURN_PACKAGE* pUpgradeBundlePackage,
     __in_opt const BURN_PACKAGE* pForwardCompatibleBundlePackage,
@@ -636,6 +636,7 @@ extern "C" HRESULT CoreApply(
     BURN_APPLY_CONTEXT applyContext = { };
     BOOL fDeleteApplyCs = FALSE;
     BURN_CACHE_THREAD_CONTEXT cacheThreadContext = { };
+    DWORD dwCacheExitCode = 0;
     BOOL fRollbackCache = FALSE;
     DWORD dwPhaseCount = 0;
     BOOTSTRAPPER_APPLYCOMPLETE_ACTION applyCompleteAction = BOOTSTRAPPER_APPLYCOMPLETE_ACTION_NONE;
@@ -744,7 +745,10 @@ extern "C" HRESULT CoreApply(
         // If we're not caching in parallel, wait for the cache thread to terminate.
         if (!pEngineState->fParallelCacheAndExecute)
         {
-            hr = WaitForCacheThread(applyContext.hCacheThread);
+            hr = ThrdWaitForCompletion(applyContext.hCacheThread, INFINITE, &dwCacheExitCode);
+            ExitOnFailure(hr, "Failed to wait for cache thread before execute.");
+
+            hr = (HRESULT)dwCacheExitCode;
             ExitOnFailure(hr, "Failed while caching, aborting execution.");
 
             ReleaseHandle(applyContext.hCacheThread);
@@ -761,10 +765,12 @@ extern "C" HRESULT CoreApply(
     // Wait for cache thread to terminate, this should return immediately unless we're waiting for layout to complete.
     if (applyContext.hCacheThread)
     {
-        HRESULT hrCached = WaitForCacheThread(applyContext.hCacheThread);
+        HRESULT hrCached = ThrdWaitForCompletion(applyContext.hCacheThread, INFINITE, &dwCacheExitCode);
+        ExitOnFailure(hrCached, "Failed to wait for cache thread after execute.");
+
         if (SUCCEEDED(hr))
         {
-            hr = hrCached;
+            hr = (HRESULT)dwCacheExitCode;
         }
     }
 
@@ -1940,6 +1946,22 @@ LExit:
     return hr;
 }
 
+extern "C" void CoreFunctionOverride(
+    __in_opt PFN_PROCWAITFORCOMPLETION pfnProcWaitForCompletion
+    )
+{
+    vpfnProcWaitForCompletion = pfnProcWaitForCompletion;
+}
+
+extern "C" HRESULT DAPI CoreWaitForProcCompletion(
+    __in HANDLE hProcess,
+    __in DWORD dwTimeout,
+    __out DWORD* pdwReturnCode
+    )
+{
+    return vpfnProcWaitForCompletion(hProcess, dwTimeout, pdwReturnCode);
+}
+
 // internal helper functions
 
 static HRESULT AppendEscapedArgumentToCommandLine(
@@ -2268,26 +2290,6 @@ LExit:
     return (DWORD)hr;
 }
 
-static HRESULT WaitForCacheThread(
-    __in HANDLE hCacheThread
-    )
-{
-    HRESULT hr = S_OK;
-
-    if (WAIT_OBJECT_0 != ::WaitForSingleObject(hCacheThread, INFINITE))
-    {
-        ExitWithLastError(hr, "Failed to wait for cache thread to terminate.");
-    }
-
-    if (!::GetExitCodeThread(hCacheThread, (DWORD*)&hr))
-    {
-        ExitWithLastError(hr, "Failed to get cache thread exit code.");
-    }
-
-LExit:
-    return hr;
-}
-
 static void LogPackages(
     __in_opt const BURN_PACKAGE* pUpgradeBundlePackage,
     __in_opt const BURN_PACKAGE* pForwardCompatibleBundlePackage,
diff --git a/src/burn/engine/core.h b/src/burn/engine/core.h
index 28d7ea78..14dbabcc 100644
--- a/src/burn/engine/core.h
+++ b/src/burn/engine/core.h
@@ -174,6 +174,12 @@ typedef struct _BURN_APPLY_CONTEXT
     DWORD dwCacheCheckpoint;
 } BURN_APPLY_CONTEXT;
 
+typedef HRESULT (DAPI *PFN_PROCWAITFORCOMPLETION)(
+    __in HANDLE hProcess,
+    __in DWORD dwTimeout,
+    __out DWORD* pReturnCode
+    );
+
 
 // function declarations
 
@@ -280,6 +286,14 @@ HRESULT CoreParseCommandLine(
     __inout HANDLE* phSectionFile,
     __inout HANDLE* phSourceEngineFile
     );
+void CoreFunctionOverride(
+    __in_opt PFN_PROCWAITFORCOMPLETION pfnProcWaitForCompletion
+    );
+HRESULT DAPI CoreWaitForProcCompletion(
+    __in HANDLE hProcess,
+    __in DWORD dwTimeout,
+    __out_opt DWORD* pdwReturnCode
+    );
 
 #if defined(__cplusplus)
 }
diff --git a/src/burn/engine/elevation.cpp b/src/burn/engine/elevation.cpp
index 4a5be8ec..e165a257 100644
--- a/src/burn/engine/elevation.cpp
+++ b/src/burn/engine/elevation.cpp
@@ -1687,15 +1687,8 @@ static HRESULT WaitForElevatedChildCacheThread(
     HRESULT hr = S_OK;
     DWORD dwExitCode = ERROR_SUCCESS;
 
-    if (WAIT_OBJECT_0 != ::WaitForSingleObject(hCacheThread, BURN_TIMEOUT))
-    {
-        ExitWithLastError(hr, "Failed to wait for cache thread to terminate.");
-    }
-
-    if (!::GetExitCodeThread(hCacheThread, &dwExitCode))
-    {
-        ExitWithLastError(hr, "Failed to get cache thread exit code.");
-    }
+    hr = ThrdWaitForCompletion(hCacheThread, BURN_TIMEOUT, &dwExitCode);
+    ExitOnFailure(hr, "Failed to wait for cache thread to complete.");
 
     AssertSz(dwExitCode == dwExpectedExitCode, "Cache thread should have exited with the expected exit code.");
 
diff --git a/src/burn/engine/netfxchainer.cpp b/src/burn/engine/netfxchainer.cpp
index 6f223eed..68ddb093 100644
--- a/src/burn/engine/netfxchainer.cpp
+++ b/src/burn/engine/netfxchainer.cpp
@@ -338,7 +338,8 @@ extern "C" HRESULT NetFxRunChainer(
     )
 {
     HRESULT hr = S_OK;
-    DWORD er = 0;
+    DWORD dwSignaledIndex = 0;
+    BOOL fTimedOut = 0;
     WCHAR wzGuid[GUID_STRING_LENGTH];
     LPWSTR sczEventName = NULL;
     LPWSTR sczSectionName = NULL;
@@ -381,9 +382,17 @@ extern "C" HRESULT NetFxRunChainer(
 
     for (;;)
     {
-        er = ::WaitForMultipleObjects(2, handles, FALSE, 100);
-        if (WAIT_OBJECT_0 == er)
+        hr = AppWaitForMultipleObjects(2, handles, FALSE, 100, &dwSignaledIndex);
+        ExitOnWaitObjectFailure(hr, fTimedOut, "Failed to wait for netfx chainer process to complete");
+
+        if (fTimedOut)
+        {
+            continue;
+        }
+
+        switch (dwSignaledIndex)
         {
+        case 0:
             // Process has exited
             *pdwExitCode = NetFxGetResult(pNetfxChainer, &hrInternalError);
             if (E_PENDING == *pdwExitCode)
@@ -396,21 +405,17 @@ extern "C" HRESULT NetFxRunChainer(
             else if (FAILED(hrInternalError))
             {
                 // push internal error message
-                OnNetFxError(pNetfxChainer, hrInternalError, pfnGenericMessageHandler, pvContext);
+                hr = OnNetFxError(pNetfxChainer, hrInternalError, pfnGenericMessageHandler, pvContext);
                 ExitOnFailure(hr, "Failed to send internal error message from netfx chainer.");
-            }           
+            }
 
-            break;
-        }
-        else if (WAIT_OBJECT_0 + 1 == er)
-        {
+            ExitFunction();
+        case 1:
             // Chainee has notified us of a change.
             hr = ProcessNetFxMessage(pNetfxChainer, pfnGenericMessageHandler, pvContext);
             ExitOnFailure(hr, "Failed to process netfx chainer message.");
-        }
-        else if (WAIT_FAILED == er)
-        {
-            ExitWithLastError(hr, "Failed to wait for netfx chainer process to complete");
+
+            break;
         }
     }
 
diff --git a/src/burn/engine/pipe.cpp b/src/burn/engine/pipe.cpp
index e6c9905b..9529ef40 100644
--- a/src/burn/engine/pipe.cpp
+++ b/src/burn/engine/pipe.cpp
@@ -422,6 +422,7 @@ extern "C" HRESULT PipeTerminateChildProcess(
     HRESULT hr = S_OK;
     BYTE* pbData = NULL;
     SIZE_T cbData = 0;
+    BOOL fTimedOut = FALSE;
 
     // Prepare the exit message.
     hr = BuffWriteNumber(&pbData, &cbData, dwParentExitCode);
@@ -443,31 +444,28 @@ extern "C" HRESULT PipeTerminateChildProcess(
     // If we were able to get a handle to the other process, wait for it to exit.
     if (pConnection->hProcess)
     {
-        if (WAIT_FAILED == ::WaitForSingleObject(pConnection->hProcess, PIPE_WAIT_FOR_CONNECTION * PIPE_RETRY_FOR_CONNECTION))
-        {
-            ExitWithLastError(hr, "Failed to wait for child process exit.");
-        }
+        hr = AppWaitForSingleObject(pConnection->hProcess, PIPE_WAIT_FOR_CONNECTION * PIPE_RETRY_FOR_CONNECTION);
+        ExitOnWaitObjectFailure(hr, fTimedOut, "Failed to wait for child process exit.");
+
+        AssertSz(!fTimedOut, "Timed out while waiting for child process to exit.");
+    }
 
 #ifdef DEBUG
+    if (pConnection->hProcess && !fTimedOut)
+    {
         DWORD dwChildExitCode = 0;
-        DWORD dwErrorCode = ERROR_SUCCESS;
-        BOOL fReturnedExitCode = ::GetExitCodeProcess(pConnection->hProcess, &dwChildExitCode);
-        if (!fReturnedExitCode)
-        {
-            dwErrorCode = ::GetLastError(); // if the other process is elevated and we are not, then we'll get ERROR_ACCESS_DENIED.
+        HRESULT hrDebug = S_OK;
 
-            // The unit test use a thread instead of a process so try to get the exit code from
-            // the thread because we failed to get it from the process.
-            if (ERROR_INVALID_HANDLE == dwErrorCode)
-            {
-                fReturnedExitCode = ::GetExitCodeThread(pConnection->hProcess, &dwChildExitCode);
-            }
+        hrDebug = CoreWaitForProcCompletion(pConnection->hProcess, 0, &dwChildExitCode);
+        if (E_ACCESSDENIED != hrDebug) // if the other process is elevated and we are not, then we'll get ERROR_ACCESS_DENIED.
+        {
+            TraceError(hrDebug, "Failed to wait for child process completion.");
         }
-        AssertSz((fReturnedExitCode && dwChildExitCode == dwParentExitCode) ||
-                 (!fReturnedExitCode && ERROR_ACCESS_DENIED == dwErrorCode),
+
+        AssertSz(E_ACCESSDENIED == hrDebug || dwChildExitCode == dwParentExitCode,
                  "Child elevated process did not return matching exit code to parent process.");
-#endif
     }
+#endif
 
 LExit:
     return hr;
diff --git a/src/burn/engine/precomp.h b/src/burn/engine/precomp.h
index bc7046f6..e2d1b5cd 100644
--- a/src/burn/engine/precomp.h
+++ b/src/burn/engine/precomp.h
@@ -46,6 +46,7 @@
 #include <srputil.h>
 #include <strutil.h>
 #include <svcutil.h>
+#include <thrdutil.h>
 #include <userutil.h>
 #include <wiutil.h>
 #include <wuautil.h>