3 files changed, 15 insertions, 8 deletions

diff --git a/src/burn/engine/engine.cpp b/src/burn/engine/engine.cpp
index ee848acf..a408ed4a 100644
--- a/src/burn/engine/engine.cpp
+++ b/src/burn/engine/engine.cpp
@@ -796,7 +796,7 @@ LExit:
         else if (BOOTSTRAPPER_SHUTDOWN_ACTION_RELOAD_BOOTSTRAPPER == shutdownAction)
         {
             LogId(REPORT_STANDARD, MSG_BA_REQUESTED_RELOAD);
-            *pfReloadApp = TRUE;
+            *pfReloadApp = SUCCEEDED(hr);
         }
         else if (BOOTSTRAPPER_SHUTDOWN_ACTION_SKIP_CLEANUP == shutdownAction)
         {
@@ -806,7 +806,7 @@ LExit:
     }
 
     // Unload BA.
-    UserExperienceUnload(&pEngineState->userExperience);
+    UserExperienceUnload(&pEngineState->userExperience, *pfReloadApp);
 
     return hr;
 }
diff --git a/src/burn/engine/userexperience.cpp b/src/burn/engine/userexperience.cpp
index 4325a6ee..87ef4de1 100644
--- a/src/burn/engine/userexperience.cpp
+++ b/src/burn/engine/userexperience.cpp
@@ -122,7 +122,6 @@ extern "C" HRESULT UserExperienceLoad(
 
     pUserExperience->pfnBAProc = results.pfnBootstrapperApplicationProc;
     pUserExperience->pvBAProcContext = results.pvBootstrapperApplicationProcContext;
-    pUserExperience->fDisableUnloading = results.fDisableUnloading;
 
 LExit:
     return hr;
@@ -133,10 +132,18 @@ LExit:
 
 *******************************************************************/
 extern "C" HRESULT UserExperienceUnload(
-    __in BURN_USER_EXPERIENCE* pUserExperience
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __in BOOL fReload
     )
 {
     HRESULT hr = S_OK;
+    BOOTSTRAPPER_DESTROY_ARGS args = { };
+    BOOTSTRAPPER_DESTROY_RESULTS results = { };
+
+    args.cbSize = sizeof(BOOTSTRAPPER_DESTROY_ARGS);
+    args.fReload = fReload;
+
+    results.cbSize = sizeof(BOOTSTRAPPER_DESTROY_RESULTS);
 
     if (pUserExperience->hUXModule)
     {
@@ -144,11 +151,11 @@ extern "C" HRESULT UserExperienceUnload(
         PFN_BOOTSTRAPPER_APPLICATION_DESTROY pfnDestroy = (PFN_BOOTSTRAPPER_APPLICATION_DESTROY)::GetProcAddress(pUserExperience->hUXModule, "BootstrapperApplicationDestroy");
         if (pfnDestroy)
         {
-            pfnDestroy();
+            pfnDestroy(&args, &results);
         }
 
         // Free BA DLL if it supports it.
-        if (!pUserExperience->fDisableUnloading && !::FreeLibrary(pUserExperience->hUXModule))
+        if (!results.fDisableUnloading && !::FreeLibrary(pUserExperience->hUXModule))
         {
             hr = HRESULT_FROM_WIN32(::GetLastError());
             TraceError(hr, "Failed to unload BA DLL.");
diff --git a/src/burn/engine/userexperience.h b/src/burn/engine/userexperience.h
index e7489710..94b73f7d 100644
--- a/src/burn/engine/userexperience.h
+++ b/src/burn/engine/userexperience.h
@@ -22,7 +22,6 @@ typedef struct _BURN_USER_EXPERIENCE
     HMODULE hUXModule;
     PFN_BOOTSTRAPPER_APPLICATION_PROC pfnBAProc;
     LPVOID pvBAProcContext;
-    BOOL fDisableUnloading;
     LPWSTR sczTempDirectory;
 
     CRITICAL_SECTION csEngineActive;    // Changing the engine active state in the user experience must be
@@ -61,7 +60,8 @@ HRESULT UserExperienceLoad(
     __in BOOTSTRAPPER_COMMAND* pCommand
     );
 HRESULT UserExperienceUnload(
-    __in BURN_USER_EXPERIENCE* pUserExperience
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __in BOOL fReload
     );
 HRESULT UserExperienceEnsureWorkingFolder(
     __in BURN_CACHE* pCache,