3 files changed, 27 insertions, 15 deletions

diff --git a/src/burn/engine/apply.cpp b/src/burn/engine/apply.cpp
index d0db053a..41cb6ad4 100644
--- a/src/burn/engine/apply.cpp
+++ b/src/burn/engine/apply.cpp
@@ -329,8 +329,7 @@ static HRESULT ReportOverallProgressTicks(
     __in BURN_APPLY_CONTEXT* pApplyContext
     );
 static HRESULT ExecutePackageComplete(
-    __in BURN_USER_EXPERIENCE* pUX,
-    __in BURN_VARIABLES* pVariables,
+    __in BURN_ENGINE_STATE* pEngineState,
     __in LPCWSTR wzPackageId,
     __in BOOL fPackageVital,
     __in BOOL fAbandonedProcess,
@@ -2400,8 +2399,6 @@ static HRESULT DoExecuteAction(
             ExitOnFailure(hr, "Failed to execute dependency action.");
             break;
 
-            break;
-
         case BURN_EXECUTE_ACTION_TYPE_ROLLBACK_BOUNDARY_START: __fallthrough;
         case BURN_EXECUTE_ACTION_TYPE_ROLLBACK_BOUNDARY_END:
             *ppCheckpoint = NULL;
@@ -2653,7 +2650,7 @@ LExit:
     if (fBeginCalled)
     {
         pPackage->fAbandonedProcess = pContext->fAbandonedProcess;
-        hr = ExecutePackageComplete(&pEngineState->userExperience, &pEngineState->variables, pPackage->sczId, pPackage->fVital, pPackage->fAbandonedProcess, hr, hrExecute, fRollback, pRestart, pfRetry, pfSuspend);
+        hr = ExecutePackageComplete(pEngineState, pPackage->sczId, pPackage->fVital, pPackage->fAbandonedProcess, hr, hrExecute, fRollback, pRestart, pfRetry, pfSuspend);
     }
 
     return hr;
@@ -2778,7 +2775,7 @@ LExit:
     if (fBeginCalled)
     {
         pPackage->fAbandonedProcess = pContext->fAbandonedProcess;
-        hr = ExecutePackageComplete(&pEngineState->userExperience, &pEngineState->variables, pPackage->sczId, pPackage->fVital, pPackage->fAbandonedProcess, hr, hrExecute, fRollback, pRestart, pfRetry, pfSuspend);
+        hr = ExecutePackageComplete(pEngineState, pPackage->sczId, pPackage->fVital, pPackage->fAbandonedProcess, hr, hrExecute, fRollback, pRestart, pfRetry, pfSuspend);
     }
 
     return hr;
@@ -2858,7 +2855,7 @@ LExit:
     if (fBeginCalled)
     {
         pPackage->fAbandonedProcess = pContext->fAbandonedProcess;
-        hr = ExecutePackageComplete(&pEngineState->userExperience, &pEngineState->variables, pPackage->sczId, pPackage->fVital, pPackage->fAbandonedProcess, hr, hrExecute, fRollback, pRestart, pfRetry, pfSuspend);
+        hr = ExecutePackageComplete(pEngineState, pPackage->sczId, pPackage->fVital, pPackage->fAbandonedProcess, hr, hrExecute, fRollback, pRestart, pfRetry, pfSuspend);
     }
 
     return hr;
@@ -2923,7 +2920,7 @@ LExit:
     if (fBeginCalled)
     {
         Assert(!pContext->fAbandonedProcess);
-        hr = ExecutePackageComplete(&pEngineState->userExperience, &pEngineState->variables, pPackage->sczId, pPackage->fVital, FALSE, hr, hrExecute, fRollback, pRestart, pfRetry, pfSuspend);
+        hr = ExecutePackageComplete(pEngineState, pPackage->sczId, pPackage->fVital, FALSE, hr, hrExecute, fRollback, pRestart, pfRetry, pfSuspend);
     }
 
     return hr;
@@ -2997,7 +2994,7 @@ LExit:
     if (fBeginCalled)
     {
         Assert(!pContext->fAbandonedProcess);
-        hr = ExecutePackageComplete(&pEngineState->userExperience, &pEngineState->variables, pPackage->sczId, pPackage->fVital, FALSE, hr, hrExecute, fRollback, pRestart, pfRetry, pfSuspend);
+        hr = ExecutePackageComplete(pEngineState, pPackage->sczId, pPackage->fVital, FALSE, hr, hrExecute, fRollback, pRestart, pfRetry, pfSuspend);
     }
 
     return hr;
@@ -3078,7 +3075,7 @@ LExit:
     if (fBeginCalled)
     {
         pPackage->fAbandonedProcess = pContext->fAbandonedProcess;
-        hr = ExecutePackageComplete(&pEngineState->userExperience, &pEngineState->variables, pPackage->sczId, pPackage->fVital, pPackage->fAbandonedProcess, hr, hrExecute, fRollback, pRestart, pfRetry, pfSuspend);
+        hr = ExecutePackageComplete(pEngineState, pPackage->sczId, pPackage->fVital, pPackage->fAbandonedProcess, hr, hrExecute, fRollback, pRestart, pfRetry, pfSuspend);
     }
 
     return hr;
@@ -3398,7 +3395,7 @@ LExit:
     if (fBeginCalled)
     {
         Assert(!pContext->fAbandonedProcess);
-        hr = ExecutePackageComplete(&pEngineState->userExperience, &pEngineState->variables, pContext->wzExecutingPackageId, FALSE, FALSE, hr, hrExecute, fRollback, pRestart, pfRetry, pfSuspend);
+        hr = ExecutePackageComplete(pEngineState, pContext->wzExecutingPackageId, FALSE, FALSE, hr, hrExecute, fRollback, pRestart, pfRetry, pfSuspend);
     }
 
     return hr;
@@ -3574,8 +3571,7 @@ static HRESULT ReportOverallProgressTicks(
 }
 
 static HRESULT ExecutePackageComplete(
-    __in BURN_USER_EXPERIENCE* pUX,
-    __in BURN_VARIABLES* pVariables,
+    __in BURN_ENGINE_STATE* pEngineState,
     __in LPCWSTR wzPackageId,
     __in BOOL fPackageVital,
     __in BOOL fAbandonedProcess,
@@ -3589,6 +3585,13 @@ static HRESULT ExecutePackageComplete(
 {
     HRESULT hr = FAILED(hrOverall) ? hrOverall : hrExecute; // if the overall function failed use that otherwise use the execution result.
     BOOTSTRAPPER_EXECUTEPACKAGECOMPLETE_ACTION executePackageCompleteAction = FAILED(hrOverall) || SUCCEEDED(hrExecute) || fPackageVital ? BOOTSTRAPPER_EXECUTEPACKAGECOMPLETE_ACTION_NONE : BOOTSTRAPPER_EXECUTEPACKAGECOMPLETE_ACTION_IGNORE;
+    BURN_USER_EXPERIENCE* pUX = &pEngineState->userExperience;
+    BURN_VARIABLES* pVariables = &pEngineState->variables;
+
+    if (pEngineState->fCriticalShutdownInitiated)
+    {
+        *pRestart = BOOTSTRAPPER_APPLY_RESTART_INITIATED;
+    }
 
     // Send package execute complete to BA.
     UserExperienceOnExecutePackageComplete(pUX, wzPackageId, hr, *pRestart, &executePackageCompleteAction);
@@ -3596,12 +3599,15 @@ static HRESULT ExecutePackageComplete(
     {
         *pRestart = BOOTSTRAPPER_APPLY_RESTART_INITIATED;
     }
-    *pfRetry = (FAILED(hrExecute) && BOOTSTRAPPER_EXECUTEPACKAGECOMPLETE_ACTION_RETRY == executePackageCompleteAction && !fAbandonedProcess); // allow retry only on failures.
+    *pfRetry = (BOOTSTRAPPER_EXECUTEPACKAGECOMPLETE_ACTION_RETRY == executePackageCompleteAction &&
+                FAILED(hrExecute) && !fAbandonedProcess && BOOTSTRAPPER_APPLY_RESTART_INITIATED != *pRestart); // allow retry only on non-fatal failures.
     *pfSuspend = (BOOTSTRAPPER_EXECUTEPACKAGECOMPLETE_ACTION_SUSPEND == executePackageCompleteAction);
 
     // Remember this package as the package that initiated the forced restart.
-    if (BOOTSTRAPPER_APPLY_RESTART_INITIATED == *pRestart)
+    if (BOOTSTRAPPER_APPLY_RESTART_INITIATED == *pRestart && !pEngineState->wzRestartInitiatedPackageId)
     {
+        pEngineState->wzRestartInitiatedPackageId = wzPackageId;
+
         // Best effort to set the forced restart package variable.
         VariableSetString(pVariables, BURN_BUNDLE_FORCED_RESTART_PACKAGE, wzPackageId, TRUE, FALSE);
     }
diff --git a/src/burn/engine/core.h b/src/burn/engine/core.h
index c04de20a..28d7ea78 100644
--- a/src/burn/engine/core.h
+++ b/src/burn/engine/core.h
@@ -159,7 +159,9 @@ typedef struct _BURN_ENGINE_STATE
     BURN_PIPE_CONNECTION companionConnection;
     BURN_PIPE_CONNECTION embeddedConnection;
 
+    BOOL fCriticalShutdownInitiated;
     BURN_RESUME_MODE resumeMode;
+    LPCWSTR wzRestartInitiatedPackageId;
 
     BURN_ENGINE_COMMAND internalCommand;
 } BURN_ENGINE_STATE;
diff --git a/src/burn/engine/uithread.cpp b/src/burn/engine/uithread.cpp
index 187f3349..bb1e5972 100644
--- a/src/burn/engine/uithread.cpp
+++ b/src/burn/engine/uithread.cpp
@@ -18,6 +18,7 @@ struct UITHREAD_INFO
 {
     BOOL fElevatedEngine;
     BURN_USER_EXPERIENCE* pUserExperience;
+    BOOL* pfCriticalShutdownInitiated;
 };
 
 
@@ -132,6 +133,7 @@ static DWORD WINAPI ThreadProc(
 
     info.fElevatedEngine = fElevatedEngine;
     info.pUserExperience = &pEngineState->userExperience;
+    info.pfCriticalShutdownInitiated = &pEngineState->fCriticalShutdownInitiated;
 
     // Create the window to handle reboots without activating it.
     hWnd = ::CreateWindowExW(WS_EX_NOACTIVATE, wc.lpszClassName, NULL, WS_POPUP, 0, 0, 0, 0, HWND_DESKTOP, NULL, pContext->hInstance, &info);
@@ -208,6 +210,8 @@ static LRESULT CALLBACK WndProc(
             UserExperienceOnSystemShutdown(pInfo->pUserExperience, dwEndSession, &fCancel);
         }
 
+        *pInfo->pfCriticalShutdownInitiated |= fCritical;
+
         fRet = !fCancel;
         LogId(REPORT_STANDARD, MSG_SYSTEM_SHUTDOWN, LoggingBoolToString(fCritical), LoggingBoolToString(pInfo->fElevatedEngine), LoggingBoolToString(fRet));
         return fRet;