When rolling back a bundle failure, reinstall all upgrade related bundles.

Fixes #3421

31 files changed, 515 insertions, 28 deletions

diff --git a/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h b/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
index 659901be..5c6258d0 100644
--- a/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
+++ b/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
@@ -209,6 +209,7 @@ enum BOOTSTRAPPER_APPLICATION_MESSAGE
     BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGEBEGIN,
     BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGECOMPLETE,
     BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANNEDCOMPATIBLEPACKAGE,
+    BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANRESTORERELATEDBUNDLE,
 };
 
 enum BOOTSTRAPPER_APPLYCOMPLETE_ACTION
@@ -1208,6 +1209,20 @@ struct BA_ONPLANRELATEDBUNDLE_RESULTS
     BOOTSTRAPPER_REQUEST_STATE requestedState;
 };
 
+struct BA_ONPLANRESTORERELATEDBUNDLE_ARGS
+{
+    DWORD cbSize;
+    LPCWSTR wzBundleId;
+    BOOTSTRAPPER_REQUEST_STATE recommendedState;
+};
+
+struct BA_ONPLANRESTORERELATEDBUNDLE_RESULTS
+{
+    DWORD cbSize;
+    BOOL fCancel;
+    BOOTSTRAPPER_REQUEST_STATE requestedState;
+};
+
 struct BA_ONPLANROLLBACKBOUNDARY_ARGS
 {
     DWORD cbSize;
diff --git a/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs b/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs
index f277425e..b08e66c0 100644
--- a/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs
+++ b/src/api/burn/WixToolset.Mba.Core/BootstrapperApplication.cs
@@ -271,6 +271,9 @@ namespace WixToolset.Mba.Core
         /// <inheritdoc/>
         public event EventHandler<SetUpdateCompleteEventArgs> SetUpdateComplete;
 
+        /// <inheritdoc/>
+        public event EventHandler<PlanRestoreRelatedBundleEventArgs> PlanRestoreRelatedBundle;
+
         /// <summary>
         /// Entry point that is called when the bootstrapper application is ready to run.
         /// </summary>
@@ -1321,6 +1324,19 @@ namespace WixToolset.Mba.Core
             }
         }
 
+        /// <summary>
+        /// Called by the engine, raises the <see cref="PlanRestoreRelatedBundle"/> event.
+        /// </summary>
+        /// <param name="args">Additional arguments for this event.</param>
+        protected virtual void OnPlanRestoreRelatedBundle(PlanRestoreRelatedBundleEventArgs args)
+        {
+            EventHandler<PlanRestoreRelatedBundleEventArgs> handler = this.PlanRestoreRelatedBundle;
+            if (null != handler)
+            {
+                handler(this, args);
+            }
+        }
+
         #region IBootstrapperApplication Members
 
         int IBootstrapperApplication.BAProc(int message, IntPtr pvArgs, IntPtr pvResults, IntPtr pvContext)
@@ -2042,6 +2058,16 @@ namespace WixToolset.Mba.Core
             return args.HResult;
         }
 
+        int IBootstrapperApplication.OnPlanRestoreRelatedBundle(string wzBundleId, RequestState recommendedState, ref RequestState pRequestedState, ref bool fCancel)
+        {
+            PlanRestoreRelatedBundleEventArgs args = new PlanRestoreRelatedBundleEventArgs(wzBundleId, recommendedState, pRequestedState, fCancel);
+            this.OnPlanRestoreRelatedBundle(args);
+
+            pRequestedState = args.State;
+            fCancel = args.Cancel;
+            return args.HResult;
+        }
+
         #endregion
     }
 }
diff --git a/src/api/burn/WixToolset.Mba.Core/EventArgs.cs b/src/api/burn/WixToolset.Mba.Core/EventArgs.cs
index d4d70651..2e1e1be3 100644
--- a/src/api/burn/WixToolset.Mba.Core/EventArgs.cs
+++ b/src/api/burn/WixToolset.Mba.Core/EventArgs.cs
@@ -2402,4 +2402,35 @@ namespace WixToolset.Mba.Core
         /// </summary>
         public string NewPackageId { get; private set; }
     }
+
+    /// <summary>
+    /// Event arguments for <see cref="IDefaultBootstrapperApplication.PlanRestoreRelatedBundle"/>
+    /// </summary>
+    [Serializable]
+    public class PlanRestoreRelatedBundleEventArgs : CancellableHResultEventArgs
+    {
+        /// <summary />
+        public PlanRestoreRelatedBundleEventArgs(string bundleId, RequestState recommendedState, RequestState state, bool cancelRecommendation)
+            : base(cancelRecommendation)
+        {
+            this.BundleId = bundleId;
+            this.RecommendedState = recommendedState;
+            this.State = state;
+        }
+
+        /// <summary>
+        /// Gets the identity of the bundle to plan for.
+        /// </summary>
+        public string BundleId { get; private set; }
+
+        /// <summary>
+        /// Gets the recommended requested state for the bundle.
+        /// </summary>
+        public RequestState RecommendedState { get; private set; }
+
+        /// <summary>
+        /// Gets or sets the requested state for the bundle.
+        /// </summary>
+        public RequestState State { get; set; }
+    }
 }
diff --git a/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs b/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs
index 05f96106..4fbe5e18 100644
--- a/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs
+++ b/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplication.cs
@@ -1136,6 +1136,18 @@ namespace WixToolset.Mba.Core
             [MarshalAs(UnmanagedType.LPWStr)] string wzPreviousPackageId,
             [MarshalAs(UnmanagedType.LPWStr)] string wzNewPackageId
             );
+
+        /// <summary>
+        /// See <see cref="IDefaultBootstrapperApplication.PlanRestoreRelatedBundle"/>.
+        /// </summary>
+        [PreserveSig]
+        [return: MarshalAs(UnmanagedType.I4)]
+        int OnPlanRestoreRelatedBundle(
+            [MarshalAs(UnmanagedType.LPWStr)] string wzBundleId,
+            [MarshalAs(UnmanagedType.U4)] RequestState recommendedState,
+            [MarshalAs(UnmanagedType.U4)] ref RequestState pRequestedState,
+            [MarshalAs(UnmanagedType.Bool)] ref bool fCancel
+            );
     }
 
     /// <summary>
diff --git a/src/api/burn/WixToolset.Mba.Core/IDefaultBootstrapperApplication.cs b/src/api/burn/WixToolset.Mba.Core/IDefaultBootstrapperApplication.cs
index ce06408e..c237cb9d 100644
--- a/src/api/burn/WixToolset.Mba.Core/IDefaultBootstrapperApplication.cs
+++ b/src/api/burn/WixToolset.Mba.Core/IDefaultBootstrapperApplication.cs
@@ -334,6 +334,11 @@ namespace WixToolset.Mba.Core
         event EventHandler<PlanRelatedBundleEventArgs> PlanRelatedBundle;
 
         /// <summary>
+        /// Fired when the engine has begun planning an upgrade related bundle for restoring in case of failure.
+        /// </summary>
+        event EventHandler<PlanRestoreRelatedBundleEventArgs> PlanRestoreRelatedBundle;
+
+        /// <summary>
         /// Fired when the engine is planning a rollback boundary.
         /// </summary>
         event EventHandler<PlanRollbackBoundaryEventArgs> PlanRollbackBoundary;
diff --git a/src/api/burn/balutil/inc/BAFunctions.h b/src/api/burn/balutil/inc/BAFunctions.h
index 84359d65..cbef88df 100644
--- a/src/api/burn/balutil/inc/BAFunctions.h
+++ b/src/api/burn/balutil/inc/BAFunctions.h
@@ -88,6 +88,7 @@ enum BA_FUNCTIONS_MESSAGE
     BA_FUNCTIONS_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGEBEGIN = BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGEBEGIN,
     BA_FUNCTIONS_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGECOMPLETE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGECOMPLETE,
     BA_FUNCTIONS_MESSAGE_ONPLANNEDCOMPATIBLEPACKAGE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANNEDCOMPATIBLEPACKAGE,
+    BA_FUNCTIONS_MESSAGE_ONPLANRESTORERELATEDBUNDLE = BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANRESTORERELATEDBUNDLE,
 
     BA_FUNCTIONS_MESSAGE_ONTHEMELOADED = 1024,
     BA_FUNCTIONS_MESSAGE_WNDPROC,
diff --git a/src/api/burn/balutil/inc/BalBaseBAFunctions.h b/src/api/burn/balutil/inc/BalBaseBAFunctions.h
index c6d0924f..e98ebc9f 100644
--- a/src/api/burn/balutil/inc/BalBaseBAFunctions.h
+++ b/src/api/burn/balutil/inc/BalBaseBAFunctions.h
@@ -849,6 +849,16 @@ public: // IBootstrapperApplication
         return S_OK;
     }
 
+    virtual STDMETHODIMP OnPlanRestoreRelatedBundle(
+        __in_z LPCWSTR /*wzBundleId*/,
+        __in BOOTSTRAPPER_REQUEST_STATE /*recommendedState*/,
+        __inout BOOTSTRAPPER_REQUEST_STATE* /*pRequestedState*/,
+        __inout BOOL* /*pfCancel*/
+        )
+    {
+        return S_OK;
+    }
+
 public: // IBAFunctions
     virtual STDMETHODIMP OnPlan(
         )
diff --git a/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h b/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h
index 5d5ff098..09cc189e 100644
--- a/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h
+++ b/src/api/burn/balutil/inc/BalBaseBAFunctionsProc.h
@@ -159,6 +159,7 @@ static HRESULT WINAPI BalBaseBAFunctionsProc(
         case BA_FUNCTIONS_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGEBEGIN:
         case BA_FUNCTIONS_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGECOMPLETE:
         case BA_FUNCTIONS_MESSAGE_ONPLANNEDCOMPATIBLEPACKAGE:
+        case BA_FUNCTIONS_MESSAGE_ONPLANRESTORERELATEDBUNDLE:
             hr = BalBaseBootstrapperApplicationProc((BOOTSTRAPPER_APPLICATION_MESSAGE)message, pvArgs, pvResults, pvContext);
             break;
         case BA_FUNCTIONS_MESSAGE_ONTHEMELOADED:
diff --git a/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h b/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
index e1a36fdf..6a24f24b 100644
--- a/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
+++ b/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
@@ -1047,6 +1047,17 @@ public: // IBootstrapperApplication
         return S_OK;
     }
 
+    virtual STDMETHODIMP OnPlanRestoreRelatedBundle(
+        __in_z LPCWSTR /*wzBundleId*/,
+        __in BOOTSTRAPPER_REQUEST_STATE /*recommendedState*/,
+        __inout BOOTSTRAPPER_REQUEST_STATE* /*pRequestedState*/,
+        __inout BOOL* pfCancel
+        )
+    {
+        *pfCancel |= CheckCanceled();
+        return S_OK;
+    }
+
 public: //CBalBaseBootstrapperApplication
     virtual STDMETHODIMP Initialize(
         __in const BOOTSTRAPPER_CREATE_ARGS* pCreateArgs
diff --git a/src/api/burn/balutil/inc/BalBaseBootstrapperApplicationProc.h b/src/api/burn/balutil/inc/BalBaseBootstrapperApplicationProc.h
index 1ee5258e..d40390e5 100644
--- a/src/api/burn/balutil/inc/BalBaseBootstrapperApplicationProc.h
+++ b/src/api/burn/balutil/inc/BalBaseBootstrapperApplicationProc.h
@@ -720,6 +720,15 @@ static HRESULT BalBaseBAProcOnSetUpdateComplete(
     return pBA->OnSetUpdateComplete(pArgs->hrStatus, pArgs->wzPreviousPackageId, pArgs->wzNewPackageId);
 }
 
+static HRESULT BalBaseBAProcOnPlanRestoreRelatedBundle(
+    __in IBootstrapperApplication* pBA,
+    __in BA_ONPLANRESTORERELATEDBUNDLE_ARGS* pArgs,
+    __inout BA_ONPLANRESTORERELATEDBUNDLE_RESULTS* pResults
+    )
+{
+    return pBA->OnPlanRestoreRelatedBundle(pArgs->wzBundleId, pArgs->recommendedState, &pResults->requestedState, &pResults->fCancel);
+}
+
 /*******************************************************************
 BalBaseBootstrapperApplicationProc - requires pvContext to be of type IBootstrapperApplication.
                                      Provides a default mapping between the new message based BA interface and
@@ -976,6 +985,9 @@ static HRESULT WINAPI BalBaseBootstrapperApplicationProc(
         case BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANNEDCOMPATIBLEPACKAGE:
             hr = BalBaseBAProcOnPlannedCompatiblePackage(pBA, reinterpret_cast<BA_ONPLANNEDCOMPATIBLEPACKAGE_ARGS*>(pvArgs), reinterpret_cast<BA_ONPLANNEDCOMPATIBLEPACKAGE_RESULTS*>(pvResults));
             break;
+        case BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANRESTORERELATEDBUNDLE:
+            hr = BalBaseBAProcOnPlanRestoreRelatedBundle(pBA, reinterpret_cast<BA_ONPLANRESTORERELATEDBUNDLE_ARGS*>(pvArgs), reinterpret_cast<BA_ONPLANRESTORERELATEDBUNDLE_RESULTS*>(pvResults));
+            break;
         }
     }
 
diff --git a/src/api/burn/balutil/inc/IBootstrapperApplication.h b/src/api/burn/balutil/inc/IBootstrapperApplication.h
index 640f609d..5932c06e 100644
--- a/src/api/burn/balutil/inc/IBootstrapperApplication.h
+++ b/src/api/burn/balutil/inc/IBootstrapperApplication.h
@@ -691,4 +691,12 @@ DECLARE_INTERFACE_IID_(IBootstrapperApplication, IUnknown, "53C31D56-49C0-426B-A
         __in_z_opt LPCWSTR wzPreviousPackageId,
         __in_z_opt LPCWSTR wzNewPackageId
         ) = 0;
+
+    // OnPlanRestoreRelatedBundle - called when the engine begins planning an upgrade related bundle for restoring in case of failure.
+    STDMETHOD(OnPlanRestoreRelatedBundle)(
+        __in_z LPCWSTR wzBundleId,
+        __in BOOTSTRAPPER_REQUEST_STATE recommendedState,
+        __inout BOOTSTRAPPER_REQUEST_STATE* pRequestedState,
+        __inout BOOL* pfCancel
+        ) = 0;
 };
diff --git a/src/burn/engine/apply.cpp b/src/burn/engine/apply.cpp
index 5cc63d02..73f8fc72 100644
--- a/src/burn/engine/apply.cpp
+++ b/src/burn/engine/apply.cpp
@@ -210,6 +210,11 @@ static HRESULT ExecuteRelatedBundle(
     __out BOOL* pfSuspend,
     __out BOOTSTRAPPER_APPLY_RESTART* pRestart
     );
+static HRESULT DoRestoreRelatedBundleActions(
+    __in BURN_ENGINE_STATE* pEngineState,
+    __in BURN_EXECUTE_CONTEXT* pContext,
+    __out BOOTSTRAPPER_APPLY_RESTART* pRestart
+    );
 static HRESULT ExecuteExePackage(
     __in BURN_ENGINE_STATE* pEngineState,
     __in BURN_EXECUTE_ACTION* pExecuteAction,
@@ -788,6 +793,9 @@ extern "C" HRESULT ApplyExecute(
                 {
                     if (pCheckpoint->pActiveRollbackBoundary->fVital)
                     {
+                        hrRollback = DoRestoreRelatedBundleActions(pEngineState, &context, pRestart);
+                        IgnoreRollbackError(hrRollback, "Failed rollback actions");
+
                         // If the rollback boundary is vital, end execution here.
                         break;
                     }
@@ -2590,6 +2598,48 @@ LExit:
     return hr;
 }
 
+static HRESULT DoRestoreRelatedBundleActions(
+    __in BURN_ENGINE_STATE* pEngineState,
+    __in BURN_EXECUTE_CONTEXT* pContext,
+    __out BOOTSTRAPPER_APPLY_RESTART* pRestart
+    )
+{
+    HRESULT hr = S_OK;
+    BOOL fRetryIgnored = FALSE;
+    BOOL fSuspendIgnored = FALSE;
+
+    // execute restore related bundle actions
+    for (DWORD i = 0; i < pEngineState->plan.cRestoreRelatedBundleActions; ++i)
+    {
+        BURN_EXECUTE_ACTION* pRestoreRelatedBundleAction = &pEngineState->plan.rgRestoreRelatedBundleActions[i];
+        if (pRestoreRelatedBundleAction->fDeleted)
+        {
+            continue;
+        }
+
+        BOOTSTRAPPER_APPLY_RESTART restart = BOOTSTRAPPER_APPLY_RESTART_NONE;
+        switch (pRestoreRelatedBundleAction->type)
+        {
+        case BURN_EXECUTE_ACTION_TYPE_RELATED_BUNDLE:
+            hr = ExecuteRelatedBundle(pEngineState, pRestoreRelatedBundleAction, pContext, TRUE, &fRetryIgnored, &fSuspendIgnored, &restart);
+            IgnoreRollbackError(hr, "Failed to restore related bundle package.");
+            break;
+
+        default:
+            hr = E_UNEXPECTED;
+            ExitOnFailure(hr, "Invalid restore related bundle action: %d.", pRestoreRelatedBundleAction->type);
+        }
+
+        if (*pRestart < restart)
+        {
+            *pRestart = restart;
+        }
+    }
+
+LExit:
+    return hr;
+}
+
 static HRESULT ExecuteExePackage(
     __in BURN_ENGINE_STATE* pEngineState,
     __in BURN_EXECUTE_ACTION* pExecuteAction,
diff --git a/src/burn/engine/bundlepackageengine.cpp b/src/burn/engine/bundlepackageengine.cpp
index cd84601f..b5fc51e5 100644
--- a/src/burn/engine/bundlepackageengine.cpp
+++ b/src/burn/engine/bundlepackageengine.cpp
@@ -51,7 +51,7 @@ extern "C" HRESULT BundlePackageEnginePlanCalculatePackage(
             execute = BOOTSTRAPPER_ACTION_STATE_NONE;
             break;
         case BOOTSTRAPPER_REQUEST_STATE_REPAIR:
-            execute = pPackage->Bundle.fRepairable ? BOOTSTRAPPER_ACTION_STATE_REPAIR : BOOTSTRAPPER_ACTION_STATE_NONE;
+            execute = BOOTSTRAPPER_ACTION_STATE_REPAIR;
             break;
         case BOOTSTRAPPER_REQUEST_STATE_ABSENT: __fallthrough;
         case BOOTSTRAPPER_REQUEST_STATE_CACHE:
diff --git a/src/burn/engine/core.cpp b/src/burn/engine/core.cpp
index 7f7a915e..8fac7bd0 100644
--- a/src/burn/engine/core.cpp
+++ b/src/burn/engine/core.cpp
@@ -558,7 +558,7 @@ extern "C" HRESULT CorePlan(
                 ExitOnFailure(hr, "Failed to plan packages.");
 
                 // Schedule the update of related bundles last.
-                hr = PlanRelatedBundlesComplete(&pEngineState->registration, &pEngineState->plan, &pEngineState->log, &pEngineState->variables, dwExecuteActionEarlyIndex);
+                hr = PlanRelatedBundlesComplete(&pEngineState->userExperience, &pEngineState->registration, &pEngineState->plan, &pEngineState->log, &pEngineState->variables, dwExecuteActionEarlyIndex);
                 ExitOnFailure(hr, "Failed to schedule related bundles.");
             }
         }
@@ -2309,7 +2309,7 @@ static void LogRelatedBundles(
 
             if (pRelatedBundle->fPlannable)
             {
-                LogId(REPORT_STANDARD, MSG_PLANNED_RELATED_BUNDLE, pPackage->sczId, LoggingRelationTypeToString(pRelatedBundle->relationType), LoggingRequestStateToString(pPackage->defaultRequested), LoggingRequestStateToString(pPackage->requested), LoggingActionStateToString(pPackage->execute), LoggingActionStateToString(pPackage->rollback), LoggingDependencyActionToString(pPackage->dependencyExecute));
+                LogId(REPORT_STANDARD, MSG_PLANNED_RELATED_BUNDLE, pPackage->sczId, LoggingRelationTypeToString(pRelatedBundle->relationType), LoggingRequestStateToString(pPackage->defaultRequested), LoggingRequestStateToString(pPackage->requested), LoggingActionStateToString(pPackage->execute), LoggingActionStateToString(pPackage->rollback), LoggingRequestStateToString(pRelatedBundle->defaultRequestedRestore), LoggingRequestStateToString(pRelatedBundle->requestedRestore), LoggingActionStateToString(pRelatedBundle->restore), LoggingDependencyActionToString(pPackage->dependencyExecute));
             }
         }
     }
diff --git a/src/burn/engine/engine.mc b/src/burn/engine/engine.mc
index b8c9a2d3..675a5644 100644
--- a/src/burn/engine/engine.mc
+++ b/src/burn/engine/engine.mc
@@ -352,13 +352,6 @@ Language=English
 Planned package: %1!ls!, state: %2!hs!, default requested: %3!hs!, ba requested: %4!hs!, execute: %5!hs!, rollback: %6!hs!, default cache strategy: %7!hs!, ba requested strategy: %8!hs!, cache: %9!hs!, uncache: %10!hs!, dependency: %11!hs!, expected install registration state: %12!hs!, expected cache registration state: %13!hs!
 .
 
-MessageId=202
-Severity=Success
-SymbolicName=MSG_PLANNED_BUNDLE_UX_CHANGED_REQUEST
-Language=English
-Planned bundle: %1!ls!, ba requested state: %2!hs! over default: %3!hs!
-.
-
 MessageId=203
 Severity=Success
 SymbolicName=MSG_PLANNED_MSI_FEATURE
@@ -391,7 +384,7 @@ MessageId=207
 Severity=Success
 SymbolicName=MSG_PLANNED_RELATED_BUNDLE
 Language=English
-Planned related bundle: %1!ls!, type: %2!hs!, default requested: %3!hs!, ba requested: %4!hs!, execute: %5!hs!, rollback: %6!hs!, dependency: %7!hs!
+Planned related bundle: %1!ls!, type: %2!hs!, default requested: %3!hs!, ba requested: %4!hs!, execute: %5!hs!, rollback: %6!hs!, default requested restore: %7!hs!, ba requested restore: %8!hs!, restore: %9!hs!, dependency: %10!hs!
 .
 
 MessageId=208
diff --git a/src/burn/engine/package.h b/src/burn/engine/package.h
index 6d1b5dd9..1e61d92b 100644
--- a/src/burn/engine/package.h
+++ b/src/burn/engine/package.h
@@ -309,7 +309,6 @@ typedef struct _BURN_PACKAGE
             LPCWSTR wzAncestors; // points directly into engine state.
             LPCWSTR wzEngineWorkingDirectory; // points directly into engine state.
 
-            BOOL fRepairable;
             BOOL fSupportsBurnProtocol;
 
             BURN_EXE_EXIT_CODE* rgExitCodes;
diff --git a/src/burn/engine/plan.cpp b/src/burn/engine/plan.cpp
index f850d49c..f1fb87b8 100644
--- a/src/burn/engine/plan.cpp
+++ b/src/burn/engine/plan.cpp
@@ -103,6 +103,10 @@ static HRESULT AppendCleanAction(
     __in BURN_PLAN* pPlan,
     __out BURN_CLEAN_ACTION** ppCleanAction
     );
+static HRESULT AppendRestoreRelatedBundleAction(
+    __in BURN_PLAN* pPlan,
+    __out BURN_EXECUTE_ACTION** ppExecuteAction
+    );
 static HRESULT ProcessPayloadGroup(
     __in BURN_PLAN* pPlan,
     __in BURN_PAYLOAD_GROUP* pPayloadGroup
@@ -196,6 +200,15 @@ extern "C" void PlanReset(
         MemFree(pPlan->rgRollbackActions);
     }
 
+    if (pPlan->rgRestoreRelatedBundleActions)
+    {
+        for (DWORD i = 0; i < pPlan->cRestoreRelatedBundleActions; ++i)
+        {
+            PlanUninitializeExecuteAction(&pPlan->rgRestoreRelatedBundleActions[i]);
+        }
+        MemFree(pPlan->rgRestoreRelatedBundleActions);
+    }
+
     if (pPlan->rgCleanActions)
     {
         // Nothing needs to be freed inside clean actions today.
@@ -1276,6 +1289,9 @@ extern "C" HRESULT PlanRelatedBundlesBegin(
             continue;
         }
 
+        pRelatedBundle->defaultRequestedRestore = BOOTSTRAPPER_REQUEST_STATE_NONE;
+        pRelatedBundle->requestedRestore = BOOTSTRAPPER_REQUEST_STATE_NONE;
+        pRelatedBundle->restore = BOOTSTRAPPER_ACTION_STATE_NONE;
         pRelatedBundle->package.defaultRequested = BOOTSTRAPPER_REQUEST_STATE_NONE;
         pRelatedBundle->package.requested = BOOTSTRAPPER_REQUEST_STATE_NONE;
 
@@ -1312,12 +1328,6 @@ extern "C" HRESULT PlanRelatedBundlesBegin(
         hr = UserExperienceOnPlanRelatedBundle(pUserExperience, pRelatedBundle->package.sczId, &pRelatedBundle->package.requested);
         ExitOnRootFailure(hr, "BA aborted plan related bundle.");
 
-        // Log when the BA changed the bundle state so the engine doesn't get blamed for planning the wrong thing.
-        if (pRelatedBundle->package.requested != pRelatedBundle->package.defaultRequested)
-        {
-            LogId(REPORT_STANDARD, MSG_PLANNED_BUNDLE_UX_CHANGED_REQUEST, pRelatedBundle->package.sczId, LoggingRequestStateToString(pRelatedBundle->package.requested), LoggingRequestStateToString(pRelatedBundle->package.defaultRequested));
-        }
-
         // If uninstalling and the dependent related bundle may be executed, ignore its provider key to allow for downgrades with ref-counting.
         if (BOOTSTRAPPER_ACTION_UNINSTALL == pPlan->action && BOOTSTRAPPER_RELATION_DEPENDENT == pRelatedBundle->relationType && BOOTSTRAPPER_REQUEST_STATE_NONE != pRelatedBundle->package.requested)
         {
@@ -1340,6 +1350,7 @@ LExit:
 }
 
 extern "C" HRESULT PlanRelatedBundlesComplete(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
     __in BURN_REGISTRATION* pRegistration,
     __in BURN_PLAN* pPlan,
     __in BURN_LOGGING* pLog,
@@ -1359,16 +1370,19 @@ extern "C" HRESULT PlanRelatedBundlesComplete(
     ExitOnFailure(hr, "Failed to create dictionary for planned packages.");
 
     BOOL fExecutingAnyPackage = FALSE;
+    BOOL fInstallingAnyPackage = FALSE;
 
     for (DWORD i = 0; i < pPlan->cExecuteActions; ++i)
     {
+        BOOTSTRAPPER_ACTION_STATE packageAction = BOOTSTRAPPER_ACTION_STATE_NONE;
+
         switch (pPlan->rgExecuteActions[i].type)
         {
         case BURN_EXECUTE_ACTION_TYPE_RELATED_BUNDLE:
-            if (BOOTSTRAPPER_ACTION_STATE_NONE != pPlan->rgExecuteActions[i].relatedBundle.action)
-            {
-                fExecutingAnyPackage = TRUE;
+            packageAction = pPlan->rgExecuteActions[i].relatedBundle.action;
 
+            if (BOOTSTRAPPER_ACTION_STATE_NONE != packageAction)
+            {
                 BURN_PACKAGE* pPackage = &pPlan->rgExecuteActions[i].relatedBundle.pRelatedBundle->package;
                 if (pPackage->cDependencyProviders)
                 {
@@ -1380,21 +1394,24 @@ extern "C" HRESULT PlanRelatedBundlesComplete(
             break;
 
         case BURN_EXECUTE_ACTION_TYPE_EXE_PACKAGE:
-            fExecutingAnyPackage |= (BOOTSTRAPPER_ACTION_STATE_NONE != pPlan->rgExecuteActions[i].exePackage.action);
+            packageAction = pPlan->rgExecuteActions[i].exePackage.action;
             break;
 
         case BURN_EXECUTE_ACTION_TYPE_MSI_PACKAGE:
-            fExecutingAnyPackage |= (BOOTSTRAPPER_ACTION_STATE_NONE != pPlan->rgExecuteActions[i].msiPackage.action);
+            packageAction = pPlan->rgExecuteActions[i].msiPackage.action;
             break;
 
         case BURN_EXECUTE_ACTION_TYPE_MSP_TARGET:
-            fExecutingAnyPackage |= (BOOTSTRAPPER_ACTION_STATE_NONE != pPlan->rgExecuteActions[i].mspTarget.action);
+            packageAction = pPlan->rgExecuteActions[i].mspTarget.action;
             break;
 
         case BURN_EXECUTE_ACTION_TYPE_MSU_PACKAGE:
-            fExecutingAnyPackage |= (BOOTSTRAPPER_ACTION_STATE_NONE != pPlan->rgExecuteActions[i].msuPackage.action);
+            packageAction = pPlan->rgExecuteActions[i].msuPackage.action;
             break;
         }
+
+        fExecutingAnyPackage |= BOOTSTRAPPER_ACTION_STATE_NONE != packageAction;
+        fInstallingAnyPackage |= BOOTSTRAPPER_ACTION_STATE_INSTALL == packageAction || BOOTSTRAPPER_ACTION_STATE_MINOR_UPGRADE == packageAction;
     }
 
     for (DWORD i = 0; i < pRegistration->relatedBundles.cRelatedBundles; ++i)
@@ -1492,6 +1509,62 @@ extern "C" HRESULT PlanRelatedBundlesComplete(
             hr = DependencyPlanPackageComplete(&pRelatedBundle->package, pPlan);
             ExitOnFailure(hr, "Failed to complete plan dependency actions for related bundle package: %ls", pRelatedBundle->package.sczId);
         }
+
+        if (fInstallingAnyPackage && BOOTSTRAPPER_RELATION_UPGRADE == pRelatedBundle->relationType)
+        {
+            BURN_EXECUTE_ACTION* pAction = NULL;
+
+            pRelatedBundle->defaultRequestedRestore = pRelatedBundle->requestedRestore = BOOTSTRAPPER_REQUEST_STATE_FORCE_PRESENT;
+
+            hr = UserExperienceOnPlanRestoreRelatedBundle(pUserExperience, pRelatedBundle->package.sczId, &pRelatedBundle->requestedRestore);
+            ExitOnRootFailure(hr, "BA aborted plan restore related bundle.");
+
+            switch (pRelatedBundle->requestedRestore)
+            {
+            case BOOTSTRAPPER_REQUEST_STATE_REPAIR:
+                pRelatedBundle->restore = BOOTSTRAPPER_ACTION_STATE_REPAIR;
+                break;
+            case BOOTSTRAPPER_REQUEST_STATE_ABSENT: __fallthrough;
+            case BOOTSTRAPPER_REQUEST_STATE_CACHE: __fallthrough;
+            case BOOTSTRAPPER_REQUEST_STATE_FORCE_ABSENT:
+                pRelatedBundle->restore = BOOTSTRAPPER_ACTION_STATE_UNINSTALL;
+                break;
+            case BOOTSTRAPPER_REQUEST_STATE_FORCE_PRESENT:
+                pRelatedBundle->restore = BOOTSTRAPPER_ACTION_STATE_INSTALL;
+                break;
+            default:
+                pRelatedBundle->restore = BOOTSTRAPPER_ACTION_STATE_NONE;
+                break;
+            }
+
+            if (BOOTSTRAPPER_ACTION_STATE_NONE != pRelatedBundle->restore)
+            {
+                hr = AppendRestoreRelatedBundleAction(pPlan, &pAction);
+                ExitOnFailure(hr, "Failed to append restore related bundle action to plan.");
+
+                pAction->type = BURN_EXECUTE_ACTION_TYPE_RELATED_BUNDLE;
+                pAction->relatedBundle.pRelatedBundle = pRelatedBundle;
+                pAction->relatedBundle.action = pRelatedBundle->restore;
+
+                if (pRelatedBundle->package.Bundle.sczIgnoreDependencies)
+                {
+                    hr = StrAllocString(&pAction->relatedBundle.sczIgnoreDependencies, pRelatedBundle->package.Bundle.sczIgnoreDependencies, 0);
+                    ExitOnFailure(hr, "Failed to allocate the list of dependencies to ignore.");
+                }
+
+                if (pRelatedBundle->package.Bundle.wzAncestors)
+                {
+                    hr = StrAllocString(&pAction->relatedBundle.sczAncestors, pRelatedBundle->package.Bundle.wzAncestors, 0);
+                    ExitOnFailure(hr, "Failed to allocate the list of ancestors.");
+                }
+
+                if (pRelatedBundle->package.Bundle.wzEngineWorkingDirectory)
+                {
+                    hr = StrAllocString(&pAction->relatedBundle.sczEngineWorkingDirectory, pRelatedBundle->package.Bundle.wzEngineWorkingDirectory, 0);
+                    ExitOnFailure(hr, "Failed to allocate the custom working directory.");
+                }
+            }
+        }
     }
 
 LExit:
@@ -2269,6 +2342,23 @@ LExit:
     return hr;
 }
 
+static HRESULT AppendRestoreRelatedBundleAction(
+    __in BURN_PLAN* pPlan,
+    __out BURN_EXECUTE_ACTION** ppExecuteAction
+    )
+{
+    HRESULT hr = S_OK;
+
+    hr = MemEnsureArraySizeForNewItems(reinterpret_cast<LPVOID*>(&pPlan->rgRestoreRelatedBundleActions), pPlan->cRestoreRelatedBundleActions, 1, sizeof(BURN_EXECUTE_ACTION), 5);
+    ExitOnFailure(hr, "Failed to grow plan's array of restore related bundle actions.");
+
+    *ppExecuteAction = pPlan->rgRestoreRelatedBundleActions + pPlan->cRestoreRelatedBundleActions;
+    ++pPlan->cRestoreRelatedBundleActions;
+
+LExit:
+    return hr;
+}
+
 static HRESULT ProcessPayloadGroup(
     __in BURN_PLAN* pPlan,
     __in BURN_PAYLOAD_GROUP* pPayloadGroup
@@ -2725,6 +2815,28 @@ static void ExecuteActionLog(
     }
 }
 
+static void RestoreRelatedBundleActionLog(
+    __in DWORD iAction,
+    __in BURN_EXECUTE_ACTION* pAction
+    )
+{
+    switch (pAction->type)
+    {
+    case BURN_EXECUTE_ACTION_TYPE_RELATED_BUNDLE:
+        LogStringLine(PlanDumpLevel, "Restore action[%u]: RELATED_BUNDLE package id: %ls, action: %hs, ignore dependencies: %ls", iAction, pAction->relatedBundle.pRelatedBundle->package.sczId, LoggingActionStateToString(pAction->relatedBundle.action), pAction->relatedBundle.sczIgnoreDependencies);
+        break;
+
+    default:
+        AssertSz(FALSE, "Unknown execute action type.");
+        break;
+    }
+
+    if (pAction->fDeleted)
+    {
+        LogStringLine(PlanDumpLevel, "      (deleted action)");
+    }
+}
+
 static void CleanActionLog(
     __in DWORD iAction,
     __in BURN_CLEAN_ACTION* pAction
@@ -2784,6 +2896,11 @@ extern "C" void PlanDump(
         ExecuteActionLog(i, pPlan->rgRollbackActions + i, TRUE);
     }
 
+    for (DWORD i = 0; i < pPlan->cRestoreRelatedBundleActions; ++i)
+    {
+        RestoreRelatedBundleActionLog(i, pPlan->rgRestoreRelatedBundleActions + i);
+    }
+
     for (DWORD i = 0; i < pPlan->cCleanActions; ++i)
     {
         CleanActionLog(i, pPlan->rgCleanActions + i);
diff --git a/src/burn/engine/plan.h b/src/burn/engine/plan.h
index 0734e39f..3dce8e5d 100644
--- a/src/burn/engine/plan.h
+++ b/src/burn/engine/plan.h
@@ -280,6 +280,9 @@ typedef struct _BURN_PLAN
     BURN_EXECUTE_ACTION* rgRollbackActions;
     DWORD cRollbackActions;
 
+    BURN_EXECUTE_ACTION* rgRestoreRelatedBundleActions;
+    DWORD cRestoreRelatedBundleActions;
+
     BURN_CLEAN_ACTION* rgCleanActions;
     DWORD cCleanActions;
 
@@ -394,6 +397,7 @@ HRESULT PlanRelatedBundlesBegin(
     __in BURN_PLAN* pPlan
     );
 HRESULT PlanRelatedBundlesComplete(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
     __in BURN_REGISTRATION* pRegistration,
     __in BURN_PLAN* pPlan,
     __in BURN_LOGGING* pLog,
diff --git a/src/burn/engine/pseudobundle.cpp b/src/burn/engine/pseudobundle.cpp
index 94b095c5..082c4487 100644
--- a/src/burn/engine/pseudobundle.cpp
+++ b/src/burn/engine/pseudobundle.cpp
@@ -51,7 +51,6 @@ extern "C" HRESULT PseudoBundleInitializeRelated(
     pPackage->fVital = FALSE;
 
     pPackage->fPermanent = FALSE;
-    pPackage->Bundle.fRepairable = TRUE;
     pPackage->Bundle.fSupportsBurnProtocol = fSupportsBurnProtocol;
 
     hr = StrAllocString(&pPackage->sczId, wzId, 0);
diff --git a/src/burn/engine/registration.h b/src/burn/engine/registration.h
index 0ae61974..64191828 100644
--- a/src/burn/engine/registration.h
+++ b/src/burn/engine/registration.h
@@ -61,6 +61,10 @@ typedef struct _BURN_RELATED_BUNDLE
     BOOL fPlannable;
 
     BURN_PACKAGE package;
+
+    BOOTSTRAPPER_REQUEST_STATE defaultRequestedRestore;
+    BOOTSTRAPPER_REQUEST_STATE requestedRestore;
+    BOOTSTRAPPER_ACTION_STATE restore;
 } BURN_RELATED_BUNDLE;
 
 typedef struct _BURN_RELATED_BUNDLES
diff --git a/src/burn/engine/userexperience.cpp b/src/burn/engine/userexperience.cpp
index 1439f5f2..59988bef 100644
--- a/src/burn/engine/userexperience.cpp
+++ b/src/burn/engine/userexperience.cpp
@@ -2176,6 +2176,36 @@ LExit:
     return hr;
 }
 
+EXTERN_C BAAPI UserExperienceOnPlanRestoreRelatedBundle(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __in_z LPCWSTR wzBundleId,
+    __inout BOOTSTRAPPER_REQUEST_STATE* pRequestedState
+    )
+{
+    HRESULT hr = S_OK;
+    BA_ONPLANRESTORERELATEDBUNDLE_ARGS args = { };
+    BA_ONPLANRESTORERELATEDBUNDLE_RESULTS results = { };
+
+    args.cbSize = sizeof(args);
+    args.wzBundleId = wzBundleId;
+    args.recommendedState = *pRequestedState;
+
+    results.cbSize = sizeof(results);
+    results.requestedState = *pRequestedState;
+
+    hr = SendBAMessage(pUserExperience, BOOTSTRAPPER_APPLICATION_MESSAGE_ONPLANRESTORERELATEDBUNDLE, &args, &results);
+    ExitOnFailure(hr, "BA OnPlanRestoreRelatedBundle failed.");
+
+    if (results.fCancel)
+    {
+        hr = HRESULT_FROM_WIN32(ERROR_INSTALL_USEREXIT);
+    }
+    *pRequestedState = results.requestedState;
+
+LExit:
+    return hr;
+}
+
 EXTERN_C BAAPI UserExperienceOnPlanRollbackBoundary(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __in_z LPCWSTR wzRollbackBoundaryId,
diff --git a/src/burn/engine/userexperience.h b/src/burn/engine/userexperience.h
index e8341120..8106d7f7 100644
--- a/src/burn/engine/userexperience.h
+++ b/src/burn/engine/userexperience.h
@@ -497,6 +497,11 @@ BAAPI UserExperienceOnPlanRelatedBundle(
     __in_z LPCWSTR wzBundleId,
     __inout BOOTSTRAPPER_REQUEST_STATE* pRequestedState
     );
+BAAPI UserExperienceOnPlanRestoreRelatedBundle(
+    __in BURN_USER_EXPERIENCE* pUserExperience,
+    __in_z LPCWSTR wzBundleId,
+    __inout BOOTSTRAPPER_REQUEST_STATE* pRequestedState
+    );
 BAAPI UserExperienceOnPlanRollbackBoundary(
     __in BURN_USER_EXPERIENCE* pUserExperience,
     __in_z LPCWSTR wzRollbackBoundaryId,
diff --git a/src/burn/test/BurnUnitTest/PlanTest.cpp b/src/burn/test/BurnUnitTest/PlanTest.cpp
index 99644115..81484234 100644
--- a/src/burn/test/BurnUnitTest/PlanTest.cpp
+++ b/src/burn/test/BurnUnitTest/PlanTest.cpp
@@ -170,6 +170,10 @@ namespace Bootstrapper
             Assert::Equal(7ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            ValidateRestoreRelatedBundle(pPlan, dwIndex++, L"{FD9920AD-DBCA-4C6C-8CD5-B47431CE8D21}", BOOTSTRAPPER_ACTION_STATE_INSTALL, NULL);
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
             UINT uIndex = 0;
@@ -277,6 +281,9 @@ namespace Bootstrapper
             Assert::Equal(3ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             ValidateCleanAction(pPlan, dwIndex++, L"PackageC");
             ValidateCleanAction(pPlan, dwIndex++, L"PackageB");
             ValidateCleanAction(pPlan, dwIndex++, L"PackageA");
@@ -350,6 +357,9 @@ namespace Bootstrapper
             Assert::Equal(1ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             ValidateCleanAction(pPlan, dwIndex++, L"PackageA");
             ValidateCleanCompatibleAction(pPlan, dwIndex++, L"PackageA");
             Assert::Equal(dwIndex, pPlan->cCleanActions);
@@ -438,6 +448,9 @@ namespace Bootstrapper
             Assert::Equal(2ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
             UINT uIndex = 0;
@@ -508,6 +521,9 @@ namespace Bootstrapper
             Assert::Equal(1ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
             UINT uIndex = 0;
@@ -579,6 +595,9 @@ namespace Bootstrapper
             Assert::Equal(2ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             ValidateCleanAction(pPlan, dwIndex++, L"PackageA");
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
@@ -658,6 +677,10 @@ namespace Bootstrapper
             Assert::Equal(3ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            ValidateRestoreRelatedBundle(pPlan, dwIndex++, L"{FD9920AD-DBCA-4C6C-8CD5-B47431CE8D21}", BOOTSTRAPPER_ACTION_STATE_INSTALL, NULL);
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
             UINT uIndex = 0;
@@ -744,6 +767,10 @@ namespace Bootstrapper
             Assert::Equal(3ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            ValidateRestoreRelatedBundle(pPlan, dwIndex++, L"{FD9920AD-DBCA-4C6C-8CD5-B47431CE8D21}", BOOTSTRAPPER_ACTION_STATE_INSTALL, NULL);
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
             UINT uIndex = 0;
@@ -805,6 +832,9 @@ namespace Bootstrapper
             Assert::Equal(0ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             ValidateCleanAction(pPlan, dwIndex++, L"PackageA");
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
@@ -879,6 +909,9 @@ namespace Bootstrapper
             Assert::Equal(1ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             ValidateCleanAction(pPlan, dwIndex++, L"PackageA");
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
@@ -946,6 +979,9 @@ namespace Bootstrapper
             Assert::Equal(0ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
             UINT uIndex = 0;
@@ -1057,6 +1093,9 @@ namespace Bootstrapper
             Assert::Equal(5ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             Assert::Equal(dwIndex, pPlan->cCleanActions);
 
             UINT uIndex = 0;
@@ -1147,6 +1186,9 @@ namespace Bootstrapper
             Assert::Equal(2ul, pPlan->cOverallProgressTicksTotal);
 
             dwIndex = 0;
+            Assert::Equal(dwIndex, pPlan->cRestoreRelatedBundleActions);
+
+            dwIndex = 0;
             ValidateCleanAction(pPlan, dwIndex++, L"PatchA");
             ValidateCleanAction(pPlan, dwIndex++, L"PackageA");
             Assert::Equal(dwIndex, pPlan->cCleanActions);
@@ -1829,6 +1871,28 @@ namespace Bootstrapper
             NativeAssert::StringEqual(wzName, pProvider->sczName);
         }
 
+        void ValidateRestoreRelatedBundle(
+            __in BURN_PLAN* pPlan,
+            __in DWORD dwIndex,
+            __in LPCWSTR wzPackageId,
+            __in BOOTSTRAPPER_ACTION_STATE action,
+            __in LPCWSTR wzIgnoreDependencies
+            )
+        {
+            BURN_EXECUTE_ACTION* pAction = ValidateRestoreRelatedBundleActionExists(pPlan, dwIndex);
+            Assert::Equal<DWORD>(BURN_EXECUTE_ACTION_TYPE_RELATED_BUNDLE, pAction->type);
+            NativeAssert::StringEqual(wzPackageId, pAction->relatedBundle.pRelatedBundle->package.sczId);
+            Assert::Equal<DWORD>(action, pAction->relatedBundle.action);
+            NativeAssert::StringEqual(wzIgnoreDependencies, pAction->relatedBundle.sczIgnoreDependencies);
+            Assert::Equal<BOOL>(FALSE, pAction->fDeleted);
+        }
+
+        BURN_EXECUTE_ACTION* ValidateRestoreRelatedBundleActionExists(BURN_PLAN* pPlan, DWORD dwIndex)
+        {
+            Assert::InRange(dwIndex + 1ul, 1ul, pPlan->cRestoreRelatedBundleActions);
+            return pPlan->rgRestoreRelatedBundleActions + dwIndex;
+        }
+
         void ValidateUninstallMsiCompatiblePackage(
             __in BURN_PLAN* pPlan,
             __in BOOL fRollback,
diff --git a/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp b/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
index 02c10472..8c7bce7e 100644
--- a/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
+++ b/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
@@ -2089,6 +2089,16 @@ private: // privates
         m_pfnBAFunctionsProc(BA_FUNCTIONS_MESSAGE_ONPLANCOMPATIBLEMSIPACKAGECOMPLETE, pArgs, pResults, m_pvBAFunctionsProcContext);
     }
 
+    void OnPlanRestoreRelatedBundleFallback(
+        __in BA_ONPLANRESTORERELATEDBUNDLE_ARGS* pArgs,
+        __inout BA_ONPLANRESTORERELATEDBUNDLE_RESULTS* pResults
+        )
+    {
+        BOOTSTRAPPER_REQUEST_STATE requestedState = pResults->requestedState;
+        m_pfnBAFunctionsProc(BA_FUNCTIONS_MESSAGE_ONPLANRESTORERELATEDBUNDLE, pArgs, pResults, m_pvBAFunctionsProcContext);
+        BalLogId(BOOTSTRAPPER_LOG_LEVEL_STANDARD, MSG_WIXSTDBA_PLANNED_RESTORE_RELATED_BUNDLE, m_hModule, pArgs->wzBundleId, LoggingRequestStateToString(requestedState), LoggingRequestStateToString(pResults->requestedState));
+    }
+
 
 public: //CBalBaseBootstrapperApplication
     virtual STDMETHODIMP Initialize(
diff --git a/src/ext/Bal/wixstdba/wixstdba.mc b/src/ext/Bal/wixstdba/wixstdba.mc
index 40acfe54..eeb69914 100644
--- a/src/ext/Bal/wixstdba/wixstdba.mc
+++ b/src/ext/Bal/wixstdba/wixstdba.mc
@@ -85,3 +85,10 @@ Language=English
 WIXSTDBA: Planned rollback boundary: %1!ls!, wixstdba requested transaction: %2!hs!, bafunctions requested transaction: %3!hs!
 .
 
+MessageId=9
+Severity=Success
+SymbolicName=MSG_WIXSTDBA_PLANNED_RESTORE_RELATED_BUNDLE
+Language=English
+WIXSTDBA: Planned restore related bundle: %1!ls!, wixstdba requested: %2!hs!, bafunctions requested: %3!hs!
+.
+
diff --git a/src/test/burn/TestData/UpgradeRelatedBundleTests/BundleAv3/BundleAv3.wixproj b/src/test/burn/TestData/UpgradeRelatedBundleTests/BundleAv3/BundleAv3.wixproj
new file mode 100644
index 00000000..958ceb47
--- /dev/null
+++ b/src/test/burn/TestData/UpgradeRelatedBundleTests/BundleAv3/BundleAv3.wixproj
@@ -0,0 +1,16 @@
+<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
+<Project Sdk="WixToolset.Sdk">
+  <Import Project="..\BundleAv1\BundleA.props" />
+  <PropertyGroup>
+    <Version>3.0.0.0</Version>
+  </PropertyGroup>
+  <ItemGroup>
+    <ProjectReference Include="..\PackageAv3\PackageAv3.wixproj" />
+    <ProjectReference Include="..\PackageF\PackageF.wixproj" />
+    <ProjectReference Include="..\..\TestBA\TestBAWixlib\testbawixlib.wixproj" />
+  </ItemGroup>
+  <ItemGroup>
+    <PackageReference Include="WixToolset.Bal.wixext" />
+    <PackageReference Include="WixToolset.NetFx.wixext" />
+  </ItemGroup>
+</Project>
\ No newline at end of file
diff --git a/src/test/burn/TestData/UpgradeRelatedBundleTests/BundleAv3/BundleAv3.wxs b/src/test/burn/TestData/UpgradeRelatedBundleTests/BundleAv3/BundleAv3.wxs
new file mode 100644
index 00000000..b969b504
--- /dev/null
+++ b/src/test/burn/TestData/UpgradeRelatedBundleTests/BundleAv3/BundleAv3.wxs
@@ -0,0 +1,11 @@
+<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
+
+
+<Wix xmlns="http://wixtoolset.org/schemas/v4/wxs">
+  <Fragment>
+    <PackageGroup Id="BundlePackages">
+      <MsiPackage Id="PackageA" SourceFile="$(var.PackageAv3.TargetPath)" />
+      <MsiPackage Id="PackageF" SourceFile="$(var.PackageF.TargetPath)" />
+    </PackageGroup>
+  </Fragment>
+</Wix>
diff --git a/src/test/burn/TestData/UpgradeRelatedBundleTests/PackageAv3/PackageAv3.wixproj b/src/test/burn/TestData/UpgradeRelatedBundleTests/PackageAv3/PackageAv3.wixproj
new file mode 100644
index 00000000..f3c121af
--- /dev/null
+++ b/src/test/burn/TestData/UpgradeRelatedBundleTests/PackageAv3/PackageAv3.wixproj
@@ -0,0 +1,7 @@
+<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
+<Project Sdk="WixToolset.Sdk">
+  <Import Project="..\PackageAv1\PackageA.props" />
+  <PropertyGroup>
+    <Version>3.0.0.0</Version>
+  </PropertyGroup>
+</Project>
\ No newline at end of file
diff --git a/src/test/burn/TestData/UpgradeRelatedBundleTests/PackageF/PackageF.wixproj b/src/test/burn/TestData/UpgradeRelatedBundleTests/PackageF/PackageF.wixproj
new file mode 100644
index 00000000..63d32e28
--- /dev/null
+++ b/src/test/burn/TestData/UpgradeRelatedBundleTests/PackageF/PackageF.wixproj
@@ -0,0 +1,12 @@
+<!-- Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information. -->
+<Project Sdk="WixToolset.Sdk">
+  <PropertyGroup>
+    <UpgradeCode>{8F6C8C4B-573C-416B-B1B0-467273256BD9}</UpgradeCode>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="..\..\Templates\PackageFail.wxs" Link="PackageFail.wxs" />
+  </ItemGroup>
+  <ItemGroup>
+    <PackageReference Include="WixToolset.Util.wixext" />
+  </ItemGroup>
+</Project>
\ No newline at end of file
diff --git a/src/test/burn/WixToolsetTest.BurnE2E/DependencyTests.cs b/src/test/burn/WixToolsetTest.BurnE2E/DependencyTests.cs
index 69a1fa68..ba02d8ee 100644
--- a/src/test/burn/WixToolsetTest.BurnE2E/DependencyTests.cs
+++ b/src/test/burn/WixToolsetTest.BurnE2E/DependencyTests.cs
@@ -539,7 +539,7 @@ namespace WixToolsetTest.BurnE2E
             }
         }
 
-        [Fact(Skip = "https://github.com/wixtoolset/issues/issues/3421")]
+        [Fact]
         public void DoesntLoseDependenciesOnFailedMajorUpgradeBundleFromMajorUpdateMsiFifo()
         {
             var packageAv1 = this.CreatePackageInstaller("PackageAv1");
@@ -611,7 +611,7 @@ namespace WixToolsetTest.BurnE2E
             packageGv2.VerifyInstalled(false);
         }
 
-        [Fact(Skip = "https://github.com/wixtoolset/issues/issues/3421")]
+        [Fact]
         public void DoesntLoseDependenciesOnFailedMajorUpgradeBundleFromMajorUpdateMsiLifo()
         {
             var packageAv1 = this.CreatePackageInstaller("PackageAv1");
diff --git a/src/test/burn/WixToolsetTest.BurnE2E/UpgradeRelatedBundleTests.cs b/src/test/burn/WixToolsetTest.BurnE2E/UpgradeRelatedBundleTests.cs
index 9eb5081e..32a04e5c 100644
--- a/src/test/burn/WixToolsetTest.BurnE2E/UpgradeRelatedBundleTests.cs
+++ b/src/test/burn/WixToolsetTest.BurnE2E/UpgradeRelatedBundleTests.cs
@@ -13,6 +13,33 @@ namespace WixToolsetTest.BurnE2E
         public UpgradeRelatedBundleTests(ITestOutputHelper testOutputHelper) : base(testOutputHelper) { }
 
         [Fact]
+        public void ReinstallsOlderBundleAfterFailure()
+        {
+            var packageAv2 = this.CreatePackageInstaller("PackageAv2");
+            var packageAv3 = this.CreatePackageInstaller("PackageAv3");
+            var bundleAv2 = this.CreateBundleInstaller("BundleAv2");
+            var bundleAv3 = this.CreateBundleInstaller("BundleAv3");
+
+            packageAv2.VerifyInstalled(false);
+            packageAv3.VerifyInstalled(false);
+
+            bundleAv2.Install();
+            bundleAv2.VerifyRegisteredAndInPackageCache();
+
+            packageAv2.VerifyInstalled(true);
+            packageAv3.VerifyInstalled(false);
+
+            // Verify https://github.com/wixtoolset/issues/issues/3421
+            var bundleAv3InstallLogFilePath = bundleAv3.Install((int)MSIExec.MSIExecReturnCode.ERROR_INSTALL_FAILURE);
+            bundleAv3.VerifyUnregisteredAndRemovedFromPackageCache();
+
+            Assert.True(LogVerifier.MessageInLogFileRegex(bundleAv3InstallLogFilePath, @"Applied execute package: PackageA, result: 0x0, restart: None"));
+
+            packageAv2.VerifyInstalled(true);
+            packageAv3.VerifyInstalled(false);
+        }
+
+        [Fact]
         public void ReportsRelatedBundleMissingFromCache()
         {
             var packageAv1 = this.CreatePackageInstaller("PackageAv1");